diff --git a/src/main.rs b/src/main.rs index b895c4c..bf0764c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,50 +23,12 @@ async fn main() -> ExitCode { } async fn run() -> Result<(), Box> { - let database_url = std::env::var("DATABASE_URL") - .map_err(|_| "DATABASE_URL environment variable must be set")?; - - let max_connections: u32 = std::env::var("DATABASE_MAX_CONNECTIONS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(100); - - let min_connections: u32 = std::env::var("DATABASE_MIN_CONNECTIONS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(10); - - let acquire_timeout_secs: u64 = std::env::var("DATABASE_ACQUIRE_TIMEOUT_SECS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(10); - - info!( - "Configuring database pool: max={}, min={}, acquire_timeout={}s", - max_connections, min_connections, acquire_timeout_secs - ); - - let pool = sqlx::postgres::PgPoolOptions::new() - .max_connections(max_connections) - .min_connections(min_connections) - .acquire_timeout(std::time::Duration::from_secs(acquire_timeout_secs)) - .idle_timeout(std::time::Duration::from_secs(300)) - .max_lifetime(std::time::Duration::from_secs(1800)) - .connect(&database_url) - .await - .map_err(|e| format!("Failed to connect to Postgres: {}", e))?; - - sqlx::migrate!("./migrations") - .run(&pool) - .await - .map_err(|e| format!("Failed to run migrations: {}", e))?; - - let state = AppState::new(pool.clone()).await; + let state = AppState::new().await?; tranquil_pds::sync::listener::start_sequencer_listener(state.clone()).await; let (shutdown_tx, shutdown_rx) = watch::channel(false); - let mut comms_service = CommsService::new(pool); + let mut comms_service = CommsService::new(state.db.clone()); if let Some(email_sender) = EmailSender::from_env() { info!("Email comms enabled"); diff --git a/src/state.rs b/src/state.rs index c31de94..8020334 100644 --- a/src/state.rs +++ b/src/state.rs @@ -7,6 +7,7 @@ use crate::repo::PostgresBlockStore; use crate::storage::{BlobStorage, S3BlobStorage}; use crate::sync::firehose::SequencedEvent; use sqlx::PgPool; +use std::error::Error; use std::sync::Arc; use tokio::sync::broadcast; @@ -75,7 +76,51 @@ impl RateLimitKind { } impl AppState { - pub async fn new(db: PgPool) -> Self { + pub async fn new() -> Result> { + let database_url = std::env::var("DATABASE_URL") + .map_err(|_| "DATABASE_URL environment variable must be set")?; + + let max_connections: u32 = std::env::var("DATABASE_MAX_CONNECTIONS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(100); + + let min_connections: u32 = std::env::var("DATABASE_MIN_CONNECTIONS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(10); + + let acquire_timeout_secs: u64 = std::env::var("DATABASE_ACQUIRE_TIMEOUT_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(10); + + tracing::info!( + "Configuring database pool: max={}, min={}, acquire_timeout={}s", + max_connections, + min_connections, + acquire_timeout_secs + ); + + let db = sqlx::postgres::PgPoolOptions::new() + .max_connections(max_connections) + .min_connections(min_connections) + .acquire_timeout(std::time::Duration::from_secs(acquire_timeout_secs)) + .idle_timeout(std::time::Duration::from_secs(300)) + .max_lifetime(std::time::Duration::from_secs(1800)) + .connect(&database_url) + .await + .map_err(|e| format!("Failed to connect to Postgres: {}", e))?; + + sqlx::migrate!("./migrations") + .run(&db) + .await + .map_err(|e| format!("Failed to run migrations: {}", e))?; + + Ok(Self::from_db(db).await) + } + + pub async fn from_db(db: PgPool) -> Self { AuthConfig::init(); let block_store = PostgresBlockStore::new(db.clone()); diff --git a/tests/common/mod.rs b/tests/common/mod.rs index d22c663..dee51c0 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -258,7 +258,9 @@ async fn spawn_app(database_url: String) -> String { .with_email_update_limit(10000) .with_oauth_authorize_limit(10000) .with_oauth_token_limit(10000); - let state = AppState::new(pool).await.with_rate_limiters(rate_limiters); + let state = AppState::from_db(pool) + .await + .with_rate_limiters(rate_limiters); tranquil_pds::sync::listener::start_sequencer_listener(state.clone()).await; let app = tranquil_pds::app(state); tokio::spawn(async move {