From 53c3bc59e2b9061b830f7c1ae85e04a8e892e24b Mon Sep 17 00:00:00 2001
From: Sergio Benitez <sb@sergio.bz>
Date: Wed, 1 May 2024 20:04:00 -0700
Subject: [PATCH] Support sqlx_sqlite extensions in db_pools.

Resolves #2762.
---
 contrib/db_pools/lib/src/config.rs   | 16 ++++++++++++++--
 contrib/db_pools/lib/src/database.rs |  4 ++--
 contrib/db_pools/lib/src/lib.rs      |  5 ++++-
 contrib/db_pools/lib/src/pool.rs     |  6 ++++++
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/contrib/db_pools/lib/src/config.rs b/contrib/db_pools/lib/src/config.rs
index fcbbc345..8d45ea63 100644
--- a/contrib/db_pools/lib/src/config.rs
+++ b/contrib/db_pools/lib/src/config.rs
@@ -15,11 +15,14 @@ use rocket::serde::{Deserialize, Serialize};
 /// [default.databases.db_name]
 /// url = "/path/to/db.sqlite"
 ///
-/// # only `url` is required. `Initializer` provides defaults for the rest.
+/// # Only `url` is required. These have sane defaults and are optional.
 /// min_connections = 64
 /// max_connections = 1024
 /// connect_timeout = 5
 /// idle_timeout = 120
+///
+/// # This option is only supported by the `sqlx_sqlite` driver.
+/// extensions = ["memvfs", "rot13"]
 /// ```
 ///
 /// Alternatively, a custom provider can be used. For example, a custom `Figment`
@@ -36,6 +39,7 @@ use rocket::serde::{Deserialize, Serialize};
 ///             max_connections: 1024,
 ///             connect_timeout: 3,
 ///             idle_timeout: None,
+///             extensions: None,
 ///         }));
 ///
 ///     rocket::custom(figment)
@@ -45,7 +49,8 @@ use rocket::serde::{Deserialize, Serialize};
 /// For general information on configuration in Rocket, see [`rocket::config`].
 /// For higher-level details on configuring a database, see the [crate-level
 /// docs](crate#configuration).
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+// NOTE: Defaults provided by the figment created in the `Initializer` fairing.
+#[derive(Default, Serialize, Deserialize, Debug, Clone, PartialEq)]
 #[serde(crate = "rocket::serde")]
 pub struct Config {
     /// Database-specific connection and configuration URL.
@@ -80,4 +85,11 @@ pub struct Config {
     ///
     /// _Default:_ `None`.
     pub idle_timeout: Option<u64>,
+    /// A list of database extensions to load at run-time.
+    ///
+    /// **Note:** Only the `sqlx_sqlite` driver supports this option (for SQLite
+    /// extensions) at this time. All other drivers ignore this option.
+    ///
+    /// _Default:_ `None`.
+    pub extensions: Option<Vec<String>>,
 }
diff --git a/contrib/db_pools/lib/src/database.rs b/contrib/db_pools/lib/src/database.rs
index b8179be0..fa5cce34 100644
--- a/contrib/db_pools/lib/src/database.rs
+++ b/contrib/db_pools/lib/src/database.rs
@@ -261,8 +261,8 @@ impl<D: Database> Fairing for Initializer<D> {
 
         let figment = rocket.figment()
             .focus(&format!("databases.{}", D::NAME))
-            .merge(Serialized::default("max_connections", workers * 4))
-            .merge(Serialized::default("connect_timeout", 5));
+            .join(Serialized::default("max_connections", workers * 4))
+            .join(Serialized::default("connect_timeout", 5));
 
         match <D::Pool>::init(&figment).await {
             Ok(pool) => Ok(rocket.manage(D::from(pool))),
diff --git a/contrib/db_pools/lib/src/lib.rs b/contrib/db_pools/lib/src/lib.rs
index 0d9e5caf..8109b711 100644
--- a/contrib/db_pools/lib/src/lib.rs
+++ b/contrib/db_pools/lib/src/lib.rs
@@ -180,11 +180,14 @@
 //! [default.databases.db_name]
 //! url = "db.sqlite"
 //!
-//! # only `url` is required. the rest have defaults and are thus optional
+//! # Only `url` is required. These have sane defaults and are optional.
 //! min_connections = 64
 //! max_connections = 1024
 //! connect_timeout = 5
 //! idle_timeout = 120
+//!
+//! # This option is only supported by the `sqlx_sqlite` driver.
+//! extensions = ["memvfs", "rot13"]
 //! ```
 //!
 //! Or via environment variables:
diff --git a/contrib/db_pools/lib/src/pool.rs b/contrib/db_pools/lib/src/pool.rs
index 4e3ef35d..371c2f76 100644
--- a/contrib/db_pools/lib/src/pool.rs
+++ b/contrib/db_pools/lib/src/pool.rs
@@ -281,6 +281,12 @@ mod sqlx {
             *o = std::mem::take(o)
                 .busy_timeout(Duration::from_secs(__config.connect_timeout))
                 .create_if_missing(true);
+
+            if let Some(ref exts) = __config.extensions {
+                for ext in exts {
+                    *o = std::mem::take(o).extension(ext.clone());
+                }
+            }
         }
     }