taler-rust

db.rs (6636B)

---

 /*
   This file is part of TALER
   Copyright (C) 2025-2026 Taler Systems SA
 
   TALER is free software; you can redistribute it and/or modify it under the
   terms of the GNU Affero General Public License as published by the Free Software
   Foundation; either version 3, or (at your option) any later version.
 
   TALER is distributed in the hope that it will be useful, but WITHOUT ANY
   WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
   A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more details.
 
   You should have received a copy of the GNU Affero General Public License along with
   TALER; see the file COPYING.  If not, see <http://www.gnu.org/licenses/>
 */
 
 use std::{path::Path, str::FromStr as _, time::Duration};
 
 use sqlx::{
     Connection, PgConnection, PgPool,
     postgres::{PgConnectOptions, PgPoolOptions},
 };
 use taler_api::config::DbCfg;
 use taler_common::{
     config::{Config, parser::ConfigSource},
     db::{dbinit, pool_with_options},
 };
 use tracing::info;
 
 use crate::setup_tracing;
 
 /// An exclusive reservation of a reusable test database.
 ///
 /// The administrative connection holds a session advisory lock. Dropping this
 /// reservation closes that connection and releases the slot. Keep it alive
 /// until every task and connection using the database has been dropped.
 pub struct TestDb {
     _lock: PgConnection,
     options: PgConnectOptions,
 }
 
 impl TestDb {
     /// Connection options for the reserved database.
     pub fn connect_options(&self) -> PgConnectOptions {
         self.options.clone()
     }
 
     /// Reset the component schema and create a pool using its local config.
     pub async fn setup(&self, src: ConfigSource) -> PgPool {
         let cfg = Config::load(src, None::<&str>).unwrap();
         let name = format!("{}db-postgres", src.component_name);
         let db_cfg = DbCfg::parse(cfg.section(&name)).unwrap();
         self.setup_manual(db_cfg.sql_dir.as_ref(), src.component_name)
             .await
     }
 
     /// Reset the component schema using an explicit SQL directory.
     pub async fn setup_manual(&self, sql_dir: &Path, component_name: &str) -> PgPool {
         let pool = pool_with_options(
             self.connect_options(),
             &component_name.replace('-', "_"),
             PgPoolOptions::new().max_connections(3),
         )
         .await
         .unwrap();
         let mut conn = pool.acquire().await.unwrap();
         dbinit(&mut conn, sql_dir, component_name, true)
             .await
             .unwrap();
         drop(conn);
         pool
     }
 }
 
 /// Run a database test using a component's local configuration.
 ///
 /// The test gets its own current-thread Tokio runtime. On return or panic the
 /// runtime drops its remaining tasks before the database reservation is released.
 /// Return values must not retain database resources; use `()` or `Result<(), E>`.
 /// Keep all work using the database on this runtime so teardown can stop it.
 pub fn run_db_test(src: ConfigSource, test_name: &str, test: impl AsyncFnOnce(PgPool)) {
     run(test_name, async |db| test(db.setup(src).await).await)
 }
 
 /// Run a database test using an explicit SQL directory and component name.
 pub fn run_db_test_manual(
     sql_dir: &Path,
     component_name: &str,
     test_name: &str,
     test: impl AsyncFnOnce(PgPool),
 ) {
     run(test_name, async |db| {
         test(db.setup_manual(sql_dir, component_name).await).await
     })
 }
 
 /// Run a test with reserved connection options and no schema initialization.
 ///
 /// Use this for migration tests, custom fixtures, and tests sharing a database
 /// across components. As with `run_db_test`, the reservation outlives the test
 /// runtime, and all database work must remain on that runtime.
 pub fn run_db_test_raw(test_name: &str, test: impl AsyncFnOnce(PgConnectOptions)) {
     run(test_name, async |db| test(db.connect_options()).await)
 }
 
 fn run(test_name: &str, test: impl AsyncFnOnce(&TestDb)) {
     setup_tracing();
     // Locals drop in reverse declaration order, including during unwinding.
     // The reservation MUST outlive the runtime: detached tasks may own pools,
     // checked-out connections, or listeners until runtime shutdown drops them.
     let mut reservation = None;
     let runtime = tokio::runtime::Builder::new_current_thread()
         .enable_all()
         .build()
         .expect("create database test runtime");
     runtime.block_on(async {
         reservation = Some(
             reserve_db(test_name)
                 .await
                 .unwrap_or_else(|e| panic!("{test_name}: reserve test database: {e}")),
         );
         let db = reservation.as_ref().unwrap();
         info!(target: "test", test_name, database = ?db.options.get_database(), "Reserved test database");
         test(db).await
     })
 }
 
 /// Reserve one of 25 reusable test databases through a session advisory lock.
 /// Wait until a database is available if all are currently reserved.
 ///
 /// Connection defaults are `postgres:///taler_rust_check`. PostgreSQL's usual
 /// `PGHOST`, `PGPORT`, and `PGUSER` settings can select another test server.
 pub async fn test_db() -> sqlx::Result<TestDb> {
     reserve_db("taler-test-utils").await
 }
 
 async fn reserve_db(test_name: &str) -> sqlx::Result<TestDb> {
     let admin =
         PgConnectOptions::from_str("postgres:///taler_rust_check")?.application_name(test_name);
     let mut conn = PgConnection::connect_with(&admin).await?;
 
     // Advisory locks enforce the limit across test processes as well as threads.
     let slot = loop {
         let slot: Option<i32> = sqlx::query_scalar(
             "SELECT id FROM generate_series(0, 24) AS id
              WHERE pg_try_advisory_lock(id) LIMIT 1",
         )
         .fetch_optional(&mut conn)
         .await?;
         if let Some(slot) = slot {
             break slot;
         }
         tokio::time::sleep(Duration::from_millis(100)).await;
     };
     let name = match slot {
         0 => "taler_rust_test".to_owned(),
         id => format!("taler_rust_test_{id}"),
     };
     // Check in a fresh statement after acquiring the lock. A previous holder
     // may have created this database after the allocation query's snapshot.
     let exists: bool = sqlx::query_scalar(
         "SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_database WHERE datname = $1)",
     )
     .bind(&name)
     .fetch_one(&mut conn)
     .await?;
     if !exists {
         sqlx::raw_sql(sqlx::AssertSqlSafe(format!("CREATE DATABASE {name}")))
             .execute(&mut conn)
             .await?;
     }
     let options = admin.database(&name);
     Ok(TestDb {
         _lock: conn,
         options,
     })
 }