sqlite improvements (#1965)

* use direct blocking calls for SQLite in `sqlx_macros`
    * this also ensures the database is closed properly, cleaning up tempfiles
* don't send `PRAGMA journal_mode` unless set
    * this previously defaulted to WAL mode which is a permanent setting
      on databases which doesn't necessarily apply to all use-cases
    * changing into or out of WAL mode acquires an exclusive lock on the database
      that can't be waited on by `sqlite3_busy_timeout()`
    * for consistency, `sqlx-cli` commands that create databases will still
      create SQLite databases in WAL mode; added a flag to disable this.
* in general, don't send `PRAGMA`s unless different than default
    * we were sending a bunch of `PRAGMA`s with their default values just to enforce
      an execution order on them, but we can also do this by inserting empty slots
      for their keys into the `IndexMap`
* add error code to `SqliteError` printout
* document why `u64` is not supported

Author: abonander

sqlx-cli/src/database.rs

@@ -11,6 +11,12 @@ pub async fn create(connect_opts: &ConnectOpts) -> anyhow::Result<()> {
     let exists = crate::retry_connect_errors(connect_opts, Any::database_exists).await?;
 
     if !exists {
+        #[cfg(feature = "sqlite")]
+        sqlx::sqlite::CREATE_DB_WAL.store(
+            connect_opts.sqlite_create_db_wal,
+            std::sync::atomic::Ordering::Release,
+        );
+
         Any::create_database(&connect_opts.database_url).await?;
     }
 

sqlx-cli/src/opt.rs

@@ -221,6 +221,18 @@ pub struct ConnectOpts {
     /// returning an error.
     #[clap(long, default_value = "10")]
     pub connect_timeout: u64,
+
+    /// Set whether or not to create SQLite databases in Write-Ahead Log (WAL) mode:
+    /// https://www.sqlite.org/wal.html
+    ///
+    /// WAL mode is enabled by default for SQLite databases created by `sqlx-cli`.
+    ///
+    /// However, if your application sets a `journal_mode` on `SqliteConnectOptions` to something
+    /// other than `Wal`, then it will have to take the database file out of WAL mode on connecting,
+    /// which requires an exclusive lock and may return a `database is locked` (`SQLITE_BUSY`) error.
+    #[cfg(feature = "sqlite")]
+    #[clap(long, action = clap::ArgAction::Set, default_value = "true")]
+    pub sqlite_create_db_wal: bool,
 }
 
 /// Argument for automatic confirmation.

sqlx-core/src/sqlite/connection/describe.rs

@@ -8,7 +8,7 @@ use crate::sqlite::{Sqlite, SqliteColumn};
 use either::Either;
 use std::convert::identity;
 
-pub(super) fn describe(conn: &mut ConnectionState, query: &str) -> Result<Describe<Sqlite>, Error> {
+pub(crate) fn describe(conn: &mut ConnectionState, query: &str) -> Result<Describe<Sqlite>, Error> {
     // describing a statement from SQLite can be involved
     // each SQLx statement is comprised of multiple SQL statements
 

sqlx-core/src/sqlite/connection/execute.rs

@@ -53,6 +53,16 @@ fn bind(
     Ok(n)
 }
 
+impl ExecuteIter<'_> {
+    pub fn finish(&mut self) -> Result<(), Error> {
+        for res in self {
+            let _ = res?;
+        }
+
+        Ok(())
+    }
+}
+
 impl Iterator for ExecuteIter<'_> {
     type Item = Result<Either<SqliteQueryResult, SqliteRow>, Error>;
 

sqlx-core/src/sqlite/connection/mod.rs

@@ -19,9 +19,9 @@ use crate::sqlite::{Sqlite, SqliteConnectOptions};
 use crate::transaction::Transaction;
 
 pub(crate) mod collation;
-mod describe;
-mod establish;
-mod execute;
+pub(crate) mod describe;
+pub(crate) mod establish;
+pub(crate) mod execute;
 mod executor;
 mod explain;
 mod handle;

sqlx-core/src/sqlite/error.rs

@@ -39,7 +39,11 @@ impl SqliteError {
 
 impl Display for SqliteError {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        f.pad(&self.message)
+        // We include the code as some produce ambiguous messages:
+        // SQLITE_BUSY: "database is locked"
+        // SQLITE_LOCKED: "database table is locked"
+        // Sadly there's no function to get the string label back from an error code.
+        write!(f, "(code: {}) {}", self.code, self.message)
     }
 }
 

sqlx-core/src/sqlite/migrate.rs

@@ -7,20 +7,32 @@ use crate::migrate::{Migrate, MigrateDatabase};
 use crate::query::query;
 use crate::query_as::query_as;
 use crate::query_scalar::query_scalar;
-use crate::sqlite::{Sqlite, SqliteConnectOptions, SqliteConnection};
+use crate::sqlite::{Sqlite, SqliteConnectOptions, SqliteConnection, SqliteJournalMode};
 use futures_core::future::BoxFuture;
 use sqlx_rt::fs;
 use std::str::FromStr;
+use std::sync::atomic::Ordering;
 use std::time::Duration;
 use std::time::Instant;
 
 impl MigrateDatabase for Sqlite {
     fn create_database(url: &str) -> BoxFuture<'_, Result<(), Error>> {
         Box::pin(async move {
+            let mut opts = SqliteConnectOptions::from_str(url)?.create_if_missing(true);
+
+            // Since it doesn't make sense to include this flag in the connection URL,
+            // we just use an `AtomicBool` to pass it.
+            if super::CREATE_DB_WAL.load(Ordering::Acquire) {
+                opts = opts.journal_mode(SqliteJournalMode::Wal);
+            }
+
             // Opening a connection to sqlite creates the database
             let _ = SqliteConnectOptions::from_str(url)?
                 .create_if_missing(true)
                 .connect()
+                .await?
+                // Ensure WAL mode tempfiles are cleaned up
+                .close()
                 .await?;
 
             Ok(())

sqlx-core/src/sqlite/mod.rs

@@ -16,11 +16,15 @@ pub use options::{
 pub use query_result::SqliteQueryResult;
 pub use row::SqliteRow;
 pub use statement::SqliteStatement;
+use std::sync::atomic::AtomicBool;
 pub use transaction::SqliteTransactionManager;
 pub use type_info::SqliteTypeInfo;
 pub use value::{SqliteValue, SqliteValueRef};
 
+use crate::describe::Describe;
+use crate::error::Error;
 use crate::executor::Executor;
+use crate::sqlite::connection::establish::EstablishParams;
 
 mod arguments;
 mod column;
@@ -60,3 +64,24 @@ impl_into_maybe_pool!(Sqlite, SqliteConnection);
 
 // required because some databases have a different handling of NULL
 impl_encode_for_option!(Sqlite);
+
+/// UNSTABLE: for use by `sqlx-cli` only.
+#[doc(hidden)]
+pub static CREATE_DB_WAL: AtomicBool = AtomicBool::new(true);
+
+/// UNSTABLE: for use by `sqlite_macros` only.
+#[doc(hidden)]
+pub fn describe_blocking(
+    opts: &SqliteConnectOptions,
+    query: &str,
+) -> Result<Describe<Sqlite>, Error> {
+    let params = EstablishParams::from_options(opts)?;
+    let mut conn = params.establish()?;
+
+    // Execute any ancillary `PRAGMA`s
+    connection::execute::iter(&mut conn, &opts.pragma_string(), None, false)?.finish()?;
+
+    connection::describe::describe(&mut conn, query)
+
+    // SQLite database is closed immediately when `conn` is dropped
+}

sqlx-core/src/sqlite/options/connect.rs

@@ -17,25 +17,8 @@ impl ConnectOptions for SqliteConnectOptions {
         Box::pin(async move {
             let mut conn = SqliteConnection::establish(self).await?;
 
-            // send an initial sql statement comprised of options
-            let mut init = String::new();
-
-            // This is a special case for sqlcipher. When the `key` pragma
-            // is set, we have to make sure it's executed first in order.
-            if let Some(pragma_key_password) = self.pragmas.get("key") {
-                write!(init, "PRAGMA key = {}; ", pragma_key_password).ok();
-            }
-
-            for (key, value) in &self.pragmas {
-                // Since we've already written the possible `key` pragma
-                // above, we shall skip it now.
-                if key == "key" {
-                    continue;
-                }
-                write!(init, "PRAGMA {} = {}; ", key, value).ok();
-            }
-
-            conn.execute(&*init).await?;
+            // Execute PRAGMAs
+            conn.execute(&*self.pragma_string()).await?;
 
             if !self.collations.is_empty() {
                 let mut locked = conn.lock_handle().await?;
@@ -59,3 +42,18 @@ impl ConnectOptions for SqliteConnectOptions {
         self
     }
 }
+
+impl SqliteConnectOptions {
+    /// Collect all `PRAMGA` commands into a single string
+    pub(crate) fn pragma_string(&self) -> String {
+        let mut string = String::new();
+
+        for (key, opt_value) in &self.pragmas {
+            if let Some(value) = opt_value {
+                write!(string, "PRAGMA {} = {}; ", key, value).ok();
+            }
+        }
+
+        string
+    }
+}