From 920b248101bfcda819abdda191670eee3ef42380 Mon Sep 17 00:00:00 2001
From: Miles Liu <miles@bung.cc>
Date: Sat, 24 Sep 2022 15:47:55 +0800
Subject: [PATCH] [SQLite] Add option to execute `PRAGMA optimize;` on close of
 a connection

---
 sqlx-core/src/sqlite/connection/mod.rs | 15 ++++++++++-
 sqlx-core/src/sqlite/options/mod.rs    | 35 ++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/sqlx-core/src/sqlite/connection/mod.rs b/sqlx-core/src/sqlite/connection/mod.rs
index 35e5f20012..e3a4e437cc 100644
--- a/sqlx-core/src/sqlite/connection/mod.rs
+++ b/sqlx-core/src/sqlite/connection/mod.rs
@@ -1,5 +1,5 @@
 use std::cmp::Ordering;
-use std::fmt::{self, Debug, Formatter};
+use std::fmt::{self, Debug, Formatter, Write};
 use std::ptr::NonNull;
 
 use futures_core::future::BoxFuture;
@@ -12,6 +12,7 @@ pub(crate) use handle::{ConnectionHandle, ConnectionHandleRaw};
 use crate::common::StatementCache;
 use crate::connection::{Connection, LogSettings};
 use crate::error::Error;
+use crate::executor::Executor;
 use crate::sqlite::connection::establish::EstablishParams;
 use crate::sqlite::connection::worker::ConnectionWorker;
 use crate::sqlite::statement::VirtualStatement;
@@ -39,6 +40,8 @@ mod worker;
 /// You can explicitly call [`.close()`][Self::close] to ensure the database is closed successfully
 /// or get an error otherwise.
 pub struct SqliteConnection {
+    optimize_on_close: bool,
+    analysis_limit: Option<u32>,
     pub(crate) worker: ConnectionWorker,
     pub(crate) row_channel_size: usize,
 }
@@ -70,6 +73,8 @@ impl SqliteConnection {
         let params = EstablishParams::from_options(options)?;
         let worker = ConnectionWorker::establish(params).await?;
         Ok(Self {
+            optimize_on_close: options.optimize_on_close,
+            analysis_limit: options.analysis_limit,
             worker,
             row_channel_size: options.row_channel_size,
         })
@@ -146,6 +151,14 @@ impl Connection for SqliteConnection {
 
     fn close(mut self) -> BoxFuture<'static, Result<(), Error>> {
         Box::pin(async move {
+            if self.optimize_on_close {
+                let mut pragma_string = String::new();
+                if let Some(limit) = self.analysis_limit {
+                    write!(pragma_string, "PRAGMA analysis_limit = {}; ", limit).ok();
+                }
+                write!(pragma_string, "PRAGMA optimize;").ok();
+                self.execute(&*pragma_string).await?;
+            }
             let shutdown = self.worker.shutdown();
             // Drop the statement worker, which should
             // cover all references to the connection handle outside of the worker thread
diff --git a/sqlx-core/src/sqlite/options/mod.rs b/sqlx-core/src/sqlite/options/mod.rs
index 7070ec4dda..3cb7fb1060 100644
--- a/sqlx-core/src/sqlite/options/mod.rs
+++ b/sqlx-core/src/sqlite/options/mod.rs
@@ -79,6 +79,9 @@ pub struct SqliteConnectOptions {
 
     pub(crate) serialized: bool,
     pub(crate) thread_name: Arc<DebugFn<dyn Fn(u64) -> String + Send + Sync + 'static>>,
+
+    pub(crate) optimize_on_close: bool,
+    pub(crate) analysis_limit: Option<u32>,
 }
 
 impl Default for SqliteConnectOptions {
@@ -185,6 +188,8 @@ impl SqliteConnectOptions {
             thread_name: Arc::new(DebugFn(|id| format!("sqlx-sqlite-worker-{}", id))),
             command_channel_size: 50,
             row_channel_size: 50,
+            optimize_on_close: false,
+            analysis_limit: None,
         }
     }
 
@@ -458,4 +463,34 @@ impl SqliteConnectOptions {
             .insert(extension_name.into(), Some(entry_point.into()));
         self
     }
+    /// Execute `PRAGMA optimize;` on the SQLite connection before closing.
+    ///
+    /// The SQLite manual recommends using this for long-lived databases.
+    ///
+    /// This will collect and store statistics about the layout of data in your tables to help the query planner make better decisions.
+    /// Over the connection's lifetime, the query planner will make notes about which tables could use up-to-date statistics so this
+    /// command doesn't have to scan the whole database every time. Thus, the best time to execute this is on connection close.
+    ///
+    /// You may also wish to set the value for [`analyze_limit`][Self::analysis_limit] so database connections close quickly,
+    /// even if your database is very large.
+    ///
+    /// Not enabled by default.
+    ///
+    /// See [the SQLite manual](https://www.sqlite.org/lang_analyze.html#automatically_running_analyze) for details.
+    pub fn optimize_on_close(mut self, enabled: bool) -> Self {
+        self.optimize_on_close = enabled;
+        self
+    }
+
+    /// Set a soft limit on the number of rows that `ANALYZE` touches per index.
+    ///
+    /// This also affects `PRAGMA optimize` which is set by [Self::optimize_on_close].
+    ///
+    /// The value recommended by SQLite is `400`. There is no default.
+    ///
+    /// See [the SQLite manual](https://www.sqlite.org/lang_analyze.html#approx) for details.
+    pub fn analysis_limit(mut self, limit: impl Into<Option<u32>>) -> Self {
+        self.analysis_limit = limit.into();
+        self
+    }
 }