[WIP] using db-pool library to create a pool of databases

Cargo.toml

@@ -143,7 +143,7 @@ lemmy_db_views_vote = { version = "=1.0.0-alpha.5", path = "./crates/db_views/vo
 activitypub_federation = { version = "0.7.0-beta.4", default-features = false, features = [
   "actix-web",
 ] }
-diesel = { version = "2.2.10", features = [
+diesel = { version = "2.2.11", features = [
   "chrono",
   "postgres",
   "serde_json",

crates/api/api_utils/src/context.rs

@@ -66,7 +66,7 @@ impl LemmyContext {
   #[allow(clippy::expect_used)]
   pub async fn init_test_federation_config() -> FederationConfig<LemmyContext> {
     // call this to run migrations
-    let pool = build_db_pool_for_tests();
+    let pool = build_db_pool_for_tests().await;
 
     let client = client_builder(&SETTINGS).build().expect("build client");
 

crates/db_schema/Cargo.toml

@@ -57,6 +57,7 @@ diesel-derive-newtype = { workspace = true, optional = true }
 diesel-async = { workspace = true, features = [
   "deadpool",
   "postgres",
+  "async-connection-wrapper",
 ], optional = true }
 regex = { workspace = true, optional = true }
 diesel_ltree = { workspace = true, optional = true }
@@ -73,6 +74,7 @@ i-love-jesus = { workspace = true, optional = true }
 derive-new.workspace = true
 tuplex = { workspace = true, optional = true }
 moka = { workspace = true, optional = true }
+db-pool = { git = "https://github.com/momentary-lapse/db-pool.git", branch = "edition2021", features = ["diesel-async-postgres", "diesel-async-deadpool"] }
 
 
 [dev-dependencies]

crates/db_schema/src/utils.rs

@@ -3,6 +3,12 @@ pub mod uplete;
 
 use crate::newtypes::DbUrl;
 use chrono::TimeDelta;
+use db_pool::{
+  r#async::{
+    DatabasePool, DatabasePoolBuilderTrait, DieselAsyncPostgresBackend, DieselDeadpool,
+  },
+  PrivilegedPostgresConfig,
+};
 use deadpool::Runtime;
 use diesel::{
   dsl,
@@ -21,6 +27,7 @@ use diesel::{
   IntoSql,
 };
 use diesel_async::{
+  async_connection_wrapper::AsyncConnectionWrapper,
   pg::AsyncPgConnection,
   pooled_connection::{
     deadpool::{Hook, HookError, Object as PooledConnection, Pool},
@@ -57,6 +64,7 @@ use std::{
   sync::{Arc, LazyLock, OnceLock},
   time::Duration,
 };
+use tokio::sync::OnceCell;
 use tracing::error;
 use url::Url;
 
@@ -518,8 +526,50 @@ pub fn build_db_pool() -> LemmyResult<ActualDbPool> {
 }
 
 #[allow(clippy::expect_used)]
-pub fn build_db_pool_for_tests() -> ActualDbPool {
-  build_db_pool().expect("db pool missing")
+pub async fn build_db_pool_for_tests() -> ActualDbPool {
+  static POOL: OnceCell<DatabasePool<DieselAsyncPostgresBackend<DieselDeadpool>>> =
+    OnceCell::const_new();
+  let db_pool = POOL
+    .get_or_init(|| async {
+      let conn_string = SETTINGS.get_database_url();
+
+      let db_url = Url::parse(conn_string.as_str()).unwrap();
+
+      let config = PrivilegedPostgresConfig::new()
+        .host(db_url.host().unwrap().to_string())
+        .port(db_url.port().unwrap())
+        .username(db_url.username().to_string())
+        .password(Some(db_url.password().unwrap().to_string()));
+
+      let backend = DieselAsyncPostgresBackend::new(
+        config,
+        |manager| Pool::builder(manager).max_size(SETTINGS.database.pool_size),
+        |manager| Pool::builder(manager).max_size(2),
+        None,
+        move |conn| {
+          Box::pin(async {
+            let async_wrapper: AsyncConnectionWrapper<AsyncPgConnection> =
+              AsyncConnectionWrapper::from(conn);
+
+            schema_setup::run_with_connection(
+              schema_setup::Options::default().run(),
+              async_wrapper,
+            )
+            .expect("run migrations");
+
+            None
+          })
+        },
+      )
+      .await
+      .unwrap();
+
+      backend.create_database_pool().await.unwrap()
+    })
+    .await;
+
+  // TODO make compatible with ActualDbPool
+  db_pool.pull_immutable().await
 }
 
 #[allow(clippy::expect_used)]

crates/db_schema_file/src/schema_setup.rs

@@ -4,7 +4,6 @@ use crate::schema::previously_run_sql;
 use anyhow::{anyhow, Context};
 use chrono::TimeDelta;
 use diesel::{
-  connection::SimpleConnection,
   dsl::exists,
   migration::{Migration, MigrationVersion},
   pg::Pg,
@@ -20,6 +19,7 @@ use diesel::{
 use diesel_migrations::MigrationHarness;
 use lemmy_utils::{error::LemmyResult, settings::SETTINGS};
 use std::time::Instant;
+use diesel::connection::LoadConnection;
 use tracing::debug;
 
 diesel::table! {
@@ -48,14 +48,20 @@ fn replaceable_schema() -> String {
 
 const REPLACEABLE_SCHEMA_PATH: &str = "crates/db_schema/replaceable_schema";
 
-struct MigrationHarnessWrapper<'a> {
-  conn: &'a mut PgConnection,
+struct MigrationHarnessWrapper<'a, Conn>
+where
+  Conn: MigrationHarness<Pg>,
+{
+  conn: &'a mut Conn,
   #[cfg(test)]
   enable_diff_check: bool,
   options: &'a Options,
 }
 
-impl MigrationHarnessWrapper<'_> {
+impl<Conn> MigrationHarnessWrapper<'_, Conn>
+where
+  Conn: MigrationHarness<Pg>,
+{
   fn run_migration_inner(
     &mut self,
     migration: &dyn Migration<Pg>,
@@ -74,7 +80,10 @@ impl MigrationHarnessWrapper<'_> {
   }
 }
 
-impl MigrationHarness<Pg> for MigrationHarnessWrapper<'_> {
+impl<Conn> MigrationHarness<Pg> for MigrationHarnessWrapper<'_, Conn>
+where
+  Conn: MigrationHarness<Pg>,
+{
   fn run_migration(
     &mut self,
     migration: &dyn Migration<Pg>,
@@ -178,12 +187,10 @@ pub enum Branch {
   ReplaceableSchemaNotRebuilt,
 }
 
-pub fn run(options: Options) -> LemmyResult<Branch> {
-  let db_url = SETTINGS.get_database_url();
-
-  // Migrations don't support async connection, and this function doesn't need to be async
-  let mut conn = PgConnection::establish(&db_url)?;
-
+pub fn run_with_connection<Conn>(options: Options, mut conn: Conn) -> LemmyResult<Branch>
+where
+  Conn: Connection<Backend = Pg> + MigrationHarness<Pg> + LoadConnection,
+{
   // If possible, skip getting a lock and recreating the "r" schema, so
   // lemmy_server processes in a horizontally scaled setup can start without causing locks
   if !options.revert
@@ -250,7 +257,17 @@ pub fn run(options: Options) -> LemmyResult<Branch> {
   Ok(output)
 }
 
-fn run_replaceable_schema(conn: &mut PgConnection) -> LemmyResult<()> {
+pub fn run(options: Options) -> LemmyResult<Branch> {
+  let db_url = SETTINGS.get_database_url();
+
+  // Migrations don't support async connection, and this function doesn't need to be async
+  run_with_connection(options, PgConnection::establish(&db_url)?)
+}
+
+fn run_replaceable_schema<Conn>(conn: &mut Conn) -> LemmyResult<()>
+where
+  Conn: Connection<Backend = Pg>,
+{
   conn.transaction(|conn| {
     conn
       .batch_execute(&replaceable_schema())
@@ -266,7 +283,10 @@ fn run_replaceable_schema(conn: &mut PgConnection) -> LemmyResult<()> {
   })
 }
 
-fn revert_replaceable_schema(conn: &mut PgConnection) -> LemmyResult<()> {
+fn revert_replaceable_schema<Conn>(conn: &mut Conn) -> LemmyResult<()>
+where
+  Conn: Connection<Backend = Pg>,
+{
   conn
     .batch_execute("DROP SCHEMA IF EXISTS r CASCADE;")
     .with_context(|| format!("Failed to revert SQL files in {REPLACEABLE_SCHEMA_PATH}"))?;
@@ -277,10 +297,13 @@ fn revert_replaceable_schema(conn: &mut PgConnection) -> LemmyResult<()> {
   Ok(())
 }
 
-fn run_selected_migrations(
-  conn: &mut PgConnection,
+fn run_selected_migrations<Conn>(
+  conn: &mut Conn,
   options: &Options,
-) -> diesel::migration::Result<()> {
+) -> diesel::migration::Result<()>
+where
+  Conn: MigrationHarness<Pg>,
+{
   let mut wrapper = MigrationHarnessWrapper {
     conn,
     options,