SQLite extension loading via sqlx.toml for CLI and query macros

Cargo.toml

@@ -21,6 +21,7 @@ members = [
     "examples/postgres/mockable-todos",
     "examples/postgres/transaction",
     "examples/sqlite/todos",
+    "examples/sqlite/extension",
 ]
 
 [workspace.package]

examples/sqlite/extension/Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "sqlx-example-sqlite-extension"
+version = "0.1.0"
+license.workspace = true
+edition.workspace = true
+repository.workspace = true
+keywords.workspace = true
+categories.workspace = true
+authors.workspace = true
+
+[dependencies]
+sqlx = { path = "../../../", features = [ "sqlite", "runtime-tokio", "tls-native-tls" ] }
+tokio = { version = "1.20.0", features = ["rt", "macros"]}
+anyhow = "1.0"
+
+[lints]
+workspace = true

examples/sqlite/extension/download-extension.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+# This grabs a pre-compiled version of the extension used in this
+# example, and stores it in a temporary directory. That's a bit
+# unusual. Normally, any extensions you need will be installed into a
+# directory on the library search path, either by using the system
+# package manager or by compiling and installing it yourself.
+
+mkdir /tmp/sqlite3-lib && wget -O /tmp/sqlite3-lib/ipaddr.so https://github.com/nalgeon/sqlean/releases/download/0.15.2/ipaddr.so

examples/sqlite/extension/migrations/20250203094951_addresses.sql

@@ -0,0 +1,25 @@
+create table addresses (address text, family integer);
+
+-- The `ipfamily` function is provided by the
+-- [ipaddr](https://github.com/nalgeon/sqlean/blob/main/docs/ipaddr.md)
+-- sqlite extension, and so this migration can not run if that
+-- extension is not loaded.
+insert into addresses (address, family) values
+  ('fd04:3d29:9f41::1', ipfamily('fd04:3d29:9f41::1')),
+  ('10.0.0.1', ipfamily('10.0.0.1')),
+  ('10.0.0.2', ipfamily('10.0.0.2')),
+  ('fd04:3d29:9f41::2', ipfamily('fd04:3d29:9f41::2')),
+  ('fd04:3d29:9f41::3', ipfamily('fd04:3d29:9f41::3')),
+  ('10.0.0.3', ipfamily('10.0.0.3')),
+  ('fd04:3d29:9f41::4', ipfamily('fd04:3d29:9f41::4')),
+  ('fd04:3d29:9f41::5', ipfamily('fd04:3d29:9f41::5')),
+  ('fd04:3d29:9f41::6', ipfamily('fd04:3d29:9f41::6')),
+  ('10.0.0.4', ipfamily('10.0.0.4')),
+  ('10.0.0.5', ipfamily('10.0.0.5')),
+  ('10.0.0.6', ipfamily('10.0.0.6')),
+  ('10.0.0.7', ipfamily('10.0.0.7')),
+  ('fd04:3d29:9f41::7', ipfamily('fd04:3d29:9f41::7')),
+  ('fd04:3d29:9f41::8', ipfamily('fd04:3d29:9f41::8')),
+  ('10.0.0.8', ipfamily('10.0.0.8')),
+  ('fd04:3d29:9f41::9', ipfamily('fd04:3d29:9f41::9')),
+  ('10.0.0.9', ipfamily('10.0.0.9'));

examples/sqlite/extension/sqlx.toml

@@ -0,0 +1,12 @@
+[common.drivers.sqlite]
+# Including the full path to the extension is somewhat unusual,
+# because normally an extension will be installed in a standard
+# directory which is part of the library search path. If that were the
+# case here, the load-extensions value could just be `["ipaddr"]`
+#
+# When the extension file is installed in a non-standard location, as
+# in this example, there are two options:
+# * Provide the full path the the extension, as seen below.
+# * Add the non-standard location to the library search path, which on
+#   Linux means adding it to the LD_LIBRARY_PATH environment variable.
+load-extensions = ["/tmp/sqlite3-lib/ipaddr"]

examples/sqlite/extension/src/main.rs

@@ -0,0 +1,33 @@
+use std::str::FromStr;
+
+use sqlx::{query, sqlite::{SqlitePool, SqliteConnectOptions}};
+
+#[tokio::main(flavor = "current_thread")]
+async fn main() -> anyhow::Result<()> {
+    let opts = SqliteConnectOptions::from_str(&std::env::var("DATABASE_URL")?)?
+        // The sqlx.toml file controls loading extensions for the CLI
+        // and for the query checking macros, *not* for the
+        // application while it's running. Thus, if we want the
+        // extension to be available during program execution, we need
+        // to load it.
+        //
+        // Note that while in this case the extension path is the same
+        // when checking the program (sqlx.toml) and when running it
+        // (here), this is not required. The runtime environment can
+        // be entirely different from the development one.
+        //
+        // The extension can be described with a full path, as seen
+        // here, but in many cases that will not be necessary. As long
+        // as the extension is installed in a directory on the library
+        // search path, it is sufficient to just provide the extension
+        // name, like "ipaddr"
+        .extension("/tmp/sqlite3-lib/ipaddr");
+
+    let db = SqlitePool::connect_with(opts).await?;
+
+    query!("insert into addresses (address, family) values (?1, ipfamily(?1))", "10.0.0.10").execute(&db).await?;
+
+    println!("Query which requires the extension was successfully executed.");
+
+    Ok(())
+}

examples/x.py

@@ -85,3 +85,4 @@ def project(name, database=None, driver=None):
 project("mysql/todos", driver="mysql_8", database="todos")
 project("postgres/todos", driver="postgres_12", database="todos")
 project("sqlite/todos", driver="sqlite", database="todos.db")
+project("sqlite/extension", driver="sqlite", database="extension.db")

sqlx-cli/src/lib.rs

@@ -131,7 +131,7 @@ pub async fn run(opt: Opt) -> Result<()> {
 
 /// Attempt to connect to the database server, retrying up to `ops.connect_timeout`.
 async fn connect(opts: &ConnectOpts) -> anyhow::Result<AnyConnection> {
-    retry_connect_errors(opts, AnyConnection::connect).await
+    retry_connect_errors(opts, AnyConnection::connect_with_config).await
 }
 
 /// Attempt an operation that may return errors like `ConnectionRefused`,

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

@@ -39,6 +39,20 @@ impl AnyConnection {
         })
     }
 
+    /// UNSTABLE: for use with `sqlx-cli`
+    ///
+    /// Connect to the database, and instruct the nested driver to
+    /// read options from the sqlx.toml file as appropriate.
+    #[doc(hidden)]
+    pub fn connect_with_config(url: &str) -> BoxFuture<'static, Result<Self, Error>>
+    where
+        Self: Sized,
+    {
+        let options: Result<AnyConnectOptions, Error> = url.parse();
+
+        Box::pin(async move { Self::connect_with(&options?.allow_config_file()).await })
+    }
+
     pub(crate) fn connect_with_db<DB: Database>(
         options: &AnyConnectOptions,
     ) -> BoxFuture<'_, crate::Result<Self>>

sqlx-core/src/any/options.rs

@@ -19,6 +19,7 @@ use url::Url;
 pub struct AnyConnectOptions {
     pub database_url: Url,
     pub log_settings: LogSettings,
+    pub enable_config: bool,
 }
 impl FromStr for AnyConnectOptions {
     type Err = Error;
@@ -29,6 +30,7 @@ impl FromStr for AnyConnectOptions {
                 .parse::<Url>()
                 .map_err(|e| Error::Configuration(e.into()))?,
             log_settings: LogSettings::default(),
+            enable_config: false,
         })
     }
 }
@@ -40,6 +42,7 @@ impl ConnectOptions for AnyConnectOptions {
         Ok(AnyConnectOptions {
             database_url: url.clone(),
             log_settings: LogSettings::default(),
+            enable_config: false,
         })
     }
 
@@ -63,3 +66,15 @@ impl ConnectOptions for AnyConnectOptions {
         self
     }
 }
+
+impl AnyConnectOptions {
+    /// UNSTABLE: for use with `sqlx-cli`
+    ///
+    /// Allow nested drivers to extract configuration information from
+    /// the sqlx.toml file.
+    #[doc(hidden)]
+    pub fn allow_config_file(mut self) -> Self {
+        self.enable_config = true;
+        self
+    }
+}

sqlx-core/src/config/common.rs

@@ -40,10 +40,49 @@ pub struct Config {
     /// The query macros used in `foo` will use `FOO_DATABASE_URL`,
     /// and the ones used in `bar` will use `BAR_DATABASE_URL`.
     pub database_url_var: Option<String>,
+
+    /// Settings for specific database drivers.
+    ///
+    /// These settings apply when checking queries, or when applying
+    /// migrations via `sqlx-cli`. These settings *do not* apply when
+    /// applying migrations via the macro, as that uses the run-time
+    /// database connection configured by the application.
+    pub drivers: Drivers,
 }
 
 impl Config {
     pub fn database_url_var(&self) -> &str {
         self.database_url_var.as_deref().unwrap_or("DATABASE_URL")
     }
 }
+
+/// Configuration for specific database drivers.
+#[derive(Debug, Default)]
+#[cfg_attr(
+    feature = "sqlx-toml",
+    derive(serde::Deserialize),
+    serde(default, rename_all = "kebab-case")
+)]
+pub struct Drivers {
+    /// Specify options for the SQLite driver.
+    pub sqlite: SQLite,
+}
+
+/// Configuration for the SQLite database driver.
+#[derive(Debug, Default)]
+#[cfg_attr(
+    feature = "sqlx-toml",
+    derive(serde::Deserialize),
+    serde(default, rename_all = "kebab-case")
+)]
+pub struct SQLite {
+    /// Specify extensions to load.
+    ///
+    /// # Example: Load the "uuid" and "vsv" extensions
+    /// `sqlx.toml`:
+    /// ```toml
+    /// [common.drivers.sqlite]
+    /// load-extensions = ["uuid", "vsv"]
+    /// ```
+    pub load_extensions: Vec<String>,
+}

sqlx-core/src/config/reference.toml

@@ -15,6 +15,12 @@
 # If not specified, defaults to `DATABASE_URL`
 database-url-var = "FOO_DATABASE_URL"
 
+[common.drivers.sqlite]
+# Load extensions into SQLite when running macros or migrations
+#
+# Defaults to an empty list, which has no effect.
+load-extensions = ["uuid", "vsv"]
+
 ###############################################################################################
 
 # Configuration for the `query!()` family of macros.

sqlx-sqlite/src/any.rs

@@ -191,6 +191,14 @@ impl<'a> TryFrom<&'a AnyConnectOptions> for SqliteConnectOptions {
     fn try_from(opts: &'a AnyConnectOptions) -> Result<Self, Self::Error> {
         let mut opts_out = SqliteConnectOptions::from_url(&opts.database_url)?;
         opts_out.log_settings = opts.log_settings.clone();
+
+        if opts.enable_config {
+            let config = sqlx_core::config::Config::from_crate();
+            for extension in config.common.drivers.sqlite.load_extensions.iter() {
+                opts_out = opts_out.extension(extension);
+            }
+        }
+
         Ok(opts_out)
     }
 }

sqlx-sqlite/src/lib.rs

@@ -124,8 +124,15 @@ pub static CREATE_DB_WAL: AtomicBool = AtomicBool::new(true);
 /// UNSTABLE: for use by `sqlite-macros-core` only.
 #[doc(hidden)]
 pub fn describe_blocking(query: &str, database_url: &str) -> Result<Describe<Sqlite>, Error> {
-    let opts: SqliteConnectOptions = database_url.parse()?;
+    let mut opts: SqliteConnectOptions = database_url.parse()?;
+
+    let config = sqlx_core::config::Config::from_crate();
+    for extension in config.common.drivers.sqlite.load_extensions.iter() {
+        opts = opts.extension(extension);
+    }
+
     let params = EstablishParams::from_options(&opts)?;
+
     let mut conn = params.establish()?;
 
     // Execute any ancillary `PRAGMA`s

tests/docker.py

@@ -17,9 +17,10 @@ def start_database(driver, database, cwd):
         database = path.join(cwd, database)
         (base_path, ext) = path.splitext(database)
         new_database = f"{base_path}.test{ext}"
-        shutil.copy(database, new_database)
+        if path.exists(database):
+            shutil.copy(database, new_database)
         # short-circuit for sqlite
-        return f"sqlite://{path.join(cwd, new_database)}"
+        return f"sqlite://{path.join(cwd, new_database)}?mode=rwc"
 
     res = subprocess.run(
         ["docker-compose", "up", "-d", driver],