kabu_storage_db/
pool.rs

1use diesel::Connection;
2use diesel::pg::PgConnection;
3use diesel_async::AsyncPgConnection;
4use diesel_async::pooled_connection::AsyncDieselConnectionManager;
5use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
6use thiserror::Error;
7use tracing::info;
8
9#[derive(Debug, Error)]
10pub enum PoolError {
11    #[error("Failed to create connection pool: {0}")]
12    PoolError(String),
13    #[error("Failed to run migrations: {0}")]
14    MigrationError(String),
15}
16
17pub type DbPool = bb8::Pool<AsyncDieselConnectionManager<AsyncPgConnection>>;
18
19pub async fn init_db_pool_with_migrations(db_url: String, migrations: EmbeddedMigrations) -> Result<DbPool, PoolError> {
20    // First, check and run migrations using a synchronous connection
21    check_and_run_migrations(&db_url, migrations)?;
22
23    // set up connection pool using bb8 with diesel-async manager
24    let manager = AsyncDieselConnectionManager::<AsyncPgConnection>::new(db_url);
25    let pool = bb8::Pool::builder().build(manager).await.map_err(|e| PoolError::PoolError(e.to_string()))?;
26    Ok(pool)
27}
28
29pub fn check_and_run_migrations(db_url: &str, migrations: EmbeddedMigrations) -> Result<(), PoolError> {
30    let mut conn =
31        PgConnection::establish(db_url).map_err(|e| PoolError::MigrationError(format!("Failed to establish connection: {e}")))?;
32
33    info!("Checking and applying database migrations...");
34
35    // Run migrations - this will check for pending and apply them
36    match conn.run_pending_migrations(migrations) {
37        Ok(applied) => {
38            if applied.is_empty() {
39                info!("Database is up to date - no pending migrations");
40            } else {
41                info!("Successfully applied {} migration(s):", applied.len());
42                for (i, migration) in applied.iter().enumerate() {
43                    info!("  {}. {}", i + 1, migration);
44                }
45            }
46        }
47        Err(e) => {
48            return Err(PoolError::MigrationError(format!("Failed to run migrations: {e}")));
49        }
50    }
51
52    // Always show the latest migration
53    if let Ok(applied_migrations) = conn.applied_migrations() {
54        if let Some(latest) = applied_migrations.last() {
55            info!("Latest migration: {}", latest);
56        } else {
57            info!("No migrations have been applied yet");
58        }
59    }
60
61    Ok(())
62}