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

Cargo.toml

@@ -11,6 +11,7 @@ members = [
     "sqlx-postgres",
     "sqlx-sqlite",
     "examples/mysql/todos",
+    "examples/postgres/axum-multi-tenant",
     "examples/postgres/axum-social-with-tests",
     "examples/postgres/chat",
     "examples/postgres/files",
@@ -23,6 +24,7 @@ members = [
     "examples/postgres/todos",
     "examples/postgres/transaction",
     "examples/sqlite/todos",
+    "examples/sqlite/extension",
 ]
 
 [workspace.package]

examples/postgres/axum-multi-tenant/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "multi-tenant"
+version.workspace = true
+license.workspace = true
+edition.workspace = true
+repository.workspace = true
+keywords.workspace = true
+categories.workspace = true
+authors.workspace = true
+
+[dependencies]
+accounts = { path = "accounts" }
+payments = { path = "payments" }
+
+sqlx = { path = "../../..", features = ["runtime-tokio", "postgres"] }
+
+[lints]
+workspace = true

examples/postgres/axum-multi-tenant/README.md

@@ -0,0 +1,11 @@
+# Axum App with Multi-tenant Database
+
+This example project involves three crates, each owning a different schema in one database,
+with their own set of migrations.
+
+* The main crate, an Axum app. 
+  * Owns the `public` schema (tables are referenced unqualified).
+* `accounts`: a subcrate simulating a reusable account-management crate.
+  * Owns schema `accounts`.
+* `payments`: a subcrate simulating a wrapper for a payments API.
+  * Owns schema `payments`.

examples/postgres/axum-multi-tenant/accounts/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "accounts"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+sqlx = { workspace = true, features = ["postgres", "time"] }
+argon2 = { version = "0.5.3", features = ["password-hash"] }
+tokio = { version = "1", features = ["rt", "sync"] }
+
+uuid = "1"
+thiserror = "1"
+rand = "0.8"

examples/postgres/axum-multi-tenant/accounts/migrations/02_account.sql

@@ -0,0 +1,8 @@
+create table account
+(
+    account_id uuid primary key default gen_random_uuid(),
+    email text unique not null,
+    password_hash text not null,
+    created_at timestamptz not null default now(),
+    updated_at timestamptz
+);

examples/postgres/axum-multi-tenant/accounts/sqlx.toml

@@ -0,0 +1,6 @@
+[migrate]
+create-schemas = ["accounts"]
+migrations-table = "accounts._sqlx_migrations"
+
+[macros.table-overrides.'accounts.account']
+'account_id' = "crate::AccountId"

examples/postgres/axum-multi-tenant/accounts/src/lib.rs

@@ -0,0 +1,133 @@
+use std::error::Error;
+use argon2::{password_hash, Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
+
+use password_hash::PasswordHashString;
+
+use sqlx::{PgConnection, PgTransaction};
+use sqlx::types::Text;
+
+use uuid::Uuid;
+
+use tokio::sync::Semaphore;
+
+#[derive(sqlx::Type)]
+#[sqlx(transparent)]
+pub struct AccountId(pub Uuid);
+
+
+pub struct AccountsManager {
+    hashing_semaphore: Semaphore,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum CreateError {
+    #[error("email in-use")]
+    EmailInUse,
+    General(#[source]
+            #[from] GeneralError),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum AuthenticateError {
+    #[error("unknown email")]
+    UnknownEmail,
+    #[error("invalid password")]
+    InvalidPassword,
+    General(#[source]
+            #[from] GeneralError),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum GeneralError {
+    Sqlx(#[source]
+         #[from] sqlx::Error),
+    PasswordHash(#[source] #[from] argon2::password_hash::Error),
+    Task(#[source]
+         #[from] tokio::task::JoinError),
+}
+
+impl AccountsManager {
+    pub async fn new(conn: &mut PgConnection, max_hashing_threads: usize) -> Result<Self, GeneralError> {
+        sqlx::migrate!().run(conn).await?;
+
+        AccountsManager {
+            hashing_semaphore: Semaphore::new(max_hashing_threads)
+        }
+    }
+
+    async fn hash_password(&self, password: String) -> Result<PasswordHash, GeneralError> {
+        let guard = self.hashing_semaphore.acquire().await
+            .expect("BUG: this semaphore should not be closed");
+
+        // We transfer ownership to the blocking task and back to ensure Tokio doesn't spawn
+        // excess threads.
+        let (_guard, res) = tokio::task::spawn_blocking(move || {
+            let salt = argon2::password_hash::SaltString::generate(rand::thread_rng());
+            (guard, Argon2::default().hash_password(password.as_bytes(), &salt))
+        })
+            .await?;
+
+        Ok(res?)
+    }
+
+    async fn verify_password(&self, password: String, hash: PasswordHashString) -> Result<(), AuthenticateError> {
+        let guard = self.hashing_semaphore.acquire().await
+            .expect("BUG: this semaphore should not be closed");
+
+        let (_guard, res) = tokio::task::spawn_blocking(move || {
+            (guard, Argon2::default().verify_password(password.as_bytes(), &hash.password_hash()))
+        }).await.map_err(GeneralError::from)?;
+
+        if let Err(password_hash::Error::Password) = res {
+            return Err(AuthenticateError::InvalidPassword);
+        }
+
+        res.map_err(GeneralError::from)?;
+
+        Ok(())
+    }
+
+    pub async fn create(&self, txn: &mut PgTransaction, email: &str, password: String) -> Result<AccountId, CreateError> {
+        // Hash password whether the account exists or not to make it harder
+        // to tell the difference in the timing.
+        let hash = self.hash_password(password).await?;
+
+        // language=PostgreSQL
+        sqlx::query!(
+            "insert into accounts.account(email, password_hash) \
+             values ($1, $2) \
+             returning account_id",
+            email,
+            Text(hash) as Text<PasswordHash<'static>>,
+        )
+            .fetch_one(&mut *txn)
+            .await
+            .map_err(|e| if e.constraint() == Some("account_account_id_key") {
+                CreateError::EmailInUse
+            } else {
+                GeneralError::from(e).into()
+            })
+    }
+
+    pub async fn authenticate(&self, conn: &mut PgConnection, email: &str, password: String) -> Result<AccountId, AuthenticateError> {
+        let maybe_account = sqlx::query!(
+            "select account_id, password_hash as \"password_hash: Text<PasswordHashString>\" \
+             from accounts.account \
+             where email_id = $1",
+            email
+        )
+            .fetch_optional(&mut *conn)
+            .await
+            .map_err(GeneralError::from)?;
+
+        let Some(account) = maybe_account else {
+            // Hash the password whether the account exists or not to hide the difference in timing.
+            self.hash_password(password).await.map_err(GeneralError::from)?;
+            return Err(AuthenticateError::UnknownEmail);
+        };
+
+        self.verify_password(password, account.password_hash.into())?;
+
+        Ok(account.account_id)
+    }
+}

examples/postgres/axum-multi-tenant/payments/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "payments"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+sqlx = { workspace = true, features = ["postgres", "time"] }

examples/postgres/axum-multi-tenant/payments/src/lib.rs

@@ -0,0 +1,14 @@
+pub fn add(left: usize, right: usize) -> usize {
+    left + right
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn it_works() {
+        let result = add(2, 2);
+        assert_eq!(result, 4);
+    }
+}

examples/postgres/axum-multi-tenant/src/main.rs

@@ -0,0 +1,3 @@
+fn main() {
+    println!("Hello, world!");
+}

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

@@ -1,10 +1,12 @@
 use std::io;
+use std::path::PathBuf;
 use std::time::Duration;
 
 use futures::{Future, TryFutureExt};
 
 use sqlx::{AnyConnection, Connection};
 use tokio::{select, signal};
+use anyhow::Context;
 
 use crate::opt::{Command, ConnectOpts, DatabaseCommand, MigrateCommand};
 
@@ -188,7 +190,7 @@ async fn do_run(opt: Opt) -> anyhow::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`,
@@ -230,3 +232,18 @@ where
     )
     .await
 }
+
+async fn config_from_current_dir() -> anyhow::Result<&'static Config> {
+    // Tokio does file I/O on a background task anyway
+    tokio::task::spawn_blocking(|| {
+        let path = PathBuf::from("sqlx.toml");
+
+        if path.exists() {
+            eprintln!("Found `sqlx.toml` in current directory; reading...");
+        }
+
+        Config::read_with_or_default(move || Ok(path))
+    })
+    .await
+    .context("unexpected error loading config")
+}

sqlx-cli/src/migrate.rs

@@ -2,7 +2,9 @@ use crate::config::Config;
 use crate::opt::{AddMigrationOpts, ConnectOpts, MigrationSourceOpt};
 use anyhow::{bail, Context};
 use console::style;
-use sqlx::migrate::{AppliedMigration, Migrate, MigrateError, MigrationType, Migrator};
+use sqlx::migrate::{
+    AppliedMigration, Migrate, MigrateError, MigrationType, Migrator, ResolveWith,
+};
 use sqlx::Connection;
 use std::borrow::Cow;
 use std::collections::{HashMap, HashSet};