athena_rs 3.4.7

//! Application bootstrap: shared `AppState`, Postgres registry, and pipeline registry.

mod postgres_init;

use actix_web::web::Data;
use anyhow::{Context, Result};
use chrono::{Duration as ChronoDuration, Utc};
use moka::future::Cache;
use reqwest::Client;
use serde_json::Value;
use std::collections::HashMap;

use std::sync::Arc;

use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::Semaphore;

use crate::AppState;
use crate::api::gateway::insert::{InsertWindowCoordinator, InsertWindowSettings};
use crate::api::metrics::MetricsState;
use crate::api::pipelines::{PipelineDefinition, load_registry_from_path};
use crate::config::Config;
use crate::config_validation::runtime_env_settings;
use crate::data::client_configs::ensure_athena_client_config_table;
use crate::data::clients::{
    ClientStatisticsRefreshMode, ClientStatisticsRefreshParams, SaveAthenaClientParams,
    list_athena_clients, refresh_client_statistics_with_params, upsert_athena_client,
};

use crate::config_validation::RuntimeEnvSettings;
#[cfg(feature = "deadpool_experimental")]
use crate::drivers::postgresql::deadpool_registry::DeadpoolPostgresRegistry;
use crate::drivers::postgresql::pool_manager::ConnectionPoolManager;
use crate::drivers::postgresql::sqlx_driver::{ClientConnectionTarget, PostgresClientRegistry};
use crate::utils::client_stats_batcher::{ClientStatsBatcher, ClientStatsBatcherConfig};
use crate::utils::request_logging::GatewayLogBatcher;
use postgres_init::connection_pool_manager_from_env;
use sqlx::PgPool;

pub use postgres_init::{
    CatalogClientStep, PostgresCatalogMergeReport, client_connection_targets_from_config,
    failed_config_keys_from_errors, merge_athena_clients_from_records,
    merge_catalog_targets_into_registry, plan_catalog_client_step,
    postgres_registry_entries_from_targets,
};

/// Shared state produced by the configuration bootstrap.
pub struct Bootstrap {
    /// Shared Actix state.
    pub app_state: Data<AppState>,
    /// Optional pipeline registry loaded from disk.
    pub pipeline_registry: Option<Arc<HashMap<String, PipelineDefinition>>>,
}

fn logging_task_limit_for_pool_max(pool_max_connections: usize) -> usize {
    (pool_max_connections.max(1) / 4).clamp(1, 16)
}

/// Builds caches, HTTP clients, and the shared AppState that both the server and CLI use.
pub async fn build_shared_state(config: &Config, pipelines_path: &str) -> Result<Bootstrap> {
    let cache_ttl: u64 = config.get_cache_ttl_secs();
    let pool_idle_timeout: u64 = config.get_pool_idle_timeout_secs();
    let runtime_env: &RuntimeEnvSettings = runtime_env_settings();
    let request_cache_max_capacity: u64 = runtime_env.cache_max_capacity;
    let request_cache_max_entry_weight: usize = runtime_env.cache_max_entry_weight;

    let cache: Arc<Cache<String, Value>> = Arc::new(
        Cache::builder()
            .support_invalidation_closures()
            .max_capacity(request_cache_max_capacity)
            .weigher(move |key: &String, value: &Value| {
                let key_weight: usize = key.len();
                let value_weight: usize = serde_json::to_vec(value)
                    .map(|bytes| bytes.len())
                    .unwrap_or(0)
                    .min(request_cache_max_entry_weight);
                let total_weight: usize = key_weight.saturating_add(value_weight);
                u32::try_from(total_weight).unwrap_or(u32::MAX)
            })
            .time_to_live(Duration::from_secs(cache_ttl))
            .build(),
    );
    let immortal_cache: Arc<Cache<String, serde_json::Value>> = Arc::new(Cache::builder().build());

    let jdbc_pool_cache: Arc<Cache<String, sqlx::postgres::PgPool>> = Arc::new(
        Cache::builder()
            .max_capacity(64)
            .time_to_live(Duration::from_secs(1800))
            .build(),
    );
    #[cfg(feature = "deadpool_experimental")]
    let jdbc_deadpool_cache: Arc<
        Cache<String, Arc<tokio::sync::OnceCell<deadpool_postgres::Pool>>>,
    > = Arc::new(
        Cache::builder()
            .max_capacity(64)
            .time_to_live(Duration::from_secs(1800))
            .build(),
    );
    let client: Client = Client::builder()
        .pool_idle_timeout(Duration::from_secs(pool_idle_timeout))
        .build()
        .context("Failed to build HTTP client")?;

    let config_targets: Vec<ClientConnectionTarget> = client_connection_targets_from_config(config);
    let postgres_entries: Vec<(String, String)> =
        postgres_registry_entries_from_targets(&config_targets);

    let pool_manager: ConnectionPoolManager = connection_pool_manager_from_env();

    let (registry, failed_connections) =
        PostgresClientRegistry::from_entries(postgres_entries, pool_manager.clone())
            .await
            .context("Failed to build Postgres registry")?;

    #[cfg(feature = "deadpool_experimental")]
    let deadpool_registry: DeadpoolPostgresRegistry = {
        let max_size: usize = runtime_env.pg_pool_max_connections as usize;
        let warmup_timeout_ms: u64 = runtime_env.deadpool_warmup_timeout_ms;
        DeadpoolPostgresRegistry::from_entries(
            postgres_registry_entries_from_targets(&config_targets),
            max_size,
            Duration::from_millis(warmup_timeout_ms),
        )
        .await
    };

    let failed_config_client_keys = failed_config_keys_from_errors(&failed_connections);

    for (client_name, err) in &failed_connections {
        tracing::debug!(
            client = %client_name,
            error = %err,
            "Postgres client skipped at bootstrap (per-client details were logged when the connection was attempted)"
        );
    }

    if registry.is_empty() {
        tracing::warn!(
            "No Postgres clients connected; Athena will run without Postgres support \
             (set missing POSTGRES_* env vars referenced in config.yaml, or fix postgres_clients URIs)"
        );
    }

    for target in &config_targets {
        registry.remember_client(
            target.clone(),
            registry.get_pool(&target.client_name).is_some(),
        );
    }

    let logging_client_name: Option<String> = config.get_gateway_logging_client().cloned();

    if let (Some(logging_client), Some(logging_pg_uri)) = (
        logging_client_name.as_ref(),
        config.get_gateway_logging_pg_uri(),
    ) {
        let override_target: ClientConnectionTarget = ClientConnectionTarget {
            client_name: logging_client.clone(),
            source: "gateway_logging_override".to_string(),
            description: Some(
                "Configured from gateway.logging_pg_uri override in config".to_string(),
            ),
            pg_uri: Some(logging_pg_uri),
            pg_uri_env_var: None,
            config_uri_template: None,
            is_active: true,
            is_frozen: false,
        };

        match registry.upsert_client(override_target.clone()).await {
            Ok(()) => {
                tracing::info!(
                    client = %override_target.client_name,
                    "Connected logging client using dedicated gateway logging URI override"
                );
            }
            Err(err) => {
                tracing::warn!(
                    client = %override_target.client_name,
                    error = %err,
                    "Failed to connect dedicated gateway logging URI override; falling back to existing registry entry if available"
                );
            }
        }
    }

    let gateway_auth_client_name: Option<String> = config.get_gateway_auth_client().cloned();
    let gateway_benchmark_client_name: Option<String> =
        config.get_gateway_benchmark_client().cloned();

    if let Some(logging_client) = logging_client_name.as_ref() {
        if let Some(logging_pool) = registry.get_pool(logging_client) {
            if let Err(err) = ensure_athena_client_config_table(&logging_pool).await {
                tracing::warn!(
                    client = %logging_client,
                    error = %err,
                    "Failed to ensure athena_client_configs table"
                );
            }

            for target in &config_targets {
                if let Err(err) = upsert_athena_client(
                    &logging_pool,
                    SaveAthenaClientParams {
                        client_name: target.client_name.clone(),
                        description: target.description.clone(),
                        pg_uri: target.pg_uri.clone(),
                        pg_uri_env_var: target.pg_uri_env_var.clone(),
                        config_uri_template: target.config_uri_template.clone(),
                        source: "config".to_string(),
                        is_active: true,
                        is_frozen: false,
                        metadata: serde_json::json!({ "seeded_from": "config.yaml" }),
                    },
                )
                .await
                {
                    tracing::warn!(
                        client = %target.client_name,
                        error = %err,
                        "Failed to sync config client into athena_clients"
                    );
                }
            }

            if config.get_gateway_database_backed_client_loading_enabled() {
                match list_athena_clients(&logging_pool).await {
                    Ok(db_clients) => {
                        merge_athena_clients_from_records(
                            &registry,
                            db_clients,
                            &failed_config_client_keys,
                        )
                        .await;
                    }
                    Err(err) => {
                        tracing::warn!(
                            client = %logging_client,
                            error = %err,
                            "Failed to load athena_clients catalog; continuing with config-backed clients only"
                        );
                    }
                }
            } else {
                tracing::info!("Database-backed client catalog loading is disabled via config");
            }

            let pool_for_client_stats: PgPool = logging_pool.clone();
            let logging_client_for_stats: String = logging_client.clone();
            let startup_refresh_mode =
                if config.get_gateway_client_statistics_startup_refresh_full() {
                    ClientStatisticsRefreshMode::Full
                } else {
                    ClientStatisticsRefreshMode::Fast
                };
            let startup_refresh_cutoff = config
                .get_gateway_client_statistics_startup_refresh_lookback_days()
                .and_then(|days| {
                    Utc::now().checked_sub_signed(ChronoDuration::days(i64::from(days)))
                });
            let startup_refresh_params = ClientStatisticsRefreshParams {
                mode: startup_refresh_mode,
                cutoff: startup_refresh_cutoff,
            };
            tokio::spawn(async move {
                tracing::info!(
                    client = %logging_client_for_stats,
                    mode = ?startup_refresh_params.mode,
                    cutoff = ?startup_refresh_params.cutoff,
                    "Refreshing client_statistics / client_table_statistics in background (does not block API bind)"
                );
                if let Err(err) = refresh_client_statistics_with_params(
                    &pool_for_client_stats,
                    startup_refresh_params,
                )
                .await
                {
                    tracing::warn!(
                        client = %logging_client_for_stats,
                        error = %err,
                        "Background refresh of client statistics failed"
                    );
                }
            });
        } else {
            tracing::warn!(
                client = %logging_client,
                "Logging client is not connected; database-backed client catalog is unavailable"
            );
        }
    }

    registry.sync_connection_status();

    let pipeline_registry: Option<Arc<HashMap<String, PipelineDefinition>>> =
        match load_registry_from_path(pipelines_path) {
            Ok(map) => {
                tracing::info!(path = %pipelines_path, "Loaded pipeline registry");
                Some(Arc::new(map))
            }
            Err(err) => {
                tracing::warn!(
                    path = %pipelines_path,
                    error = %err,
                    "Failed to load pipelines registry"
                );
                None
            }
        };

    let insert_window_settings: InsertWindowSettings = InsertWindowSettings {
        max_batch: config.get_gateway_insert_window_max_batch(),
        max_queued: config.get_gateway_insert_window_max_queued(),
        deny_tables: config.get_gateway_insert_merge_deny_tables(),
    };
    let insert_window_coordinator: Arc<InsertWindowCoordinator> =
        InsertWindowCoordinator::new(insert_window_settings.clone());

    let client_stats_batcher: Option<Arc<ClientStatsBatcher>> = logging_client_name
        .as_ref()
        .and_then(|name| registry.get_pool(name))
        .map(|pool| {
            Arc::new(ClientStatsBatcher::spawn(
                pool,
                ClientStatsBatcherConfig::default(),
            ))
        });
    let gateway_log_batcher: Option<Arc<GatewayLogBatcher>> = logging_client_name
        .as_ref()
        .and_then(|name| registry.get_pool(name))
        .map(|pool| Arc::new(GatewayLogBatcher::spawn(pool, client_stats_batcher.clone())));
    let logging_task_limit: usize =
        logging_task_limit_for_pool_max(runtime_env.pg_pool_max_connections as usize);
    let logging_task_limiter: Option<Arc<Semaphore>> = logging_client_name
        .as_ref()
        .map(|_| Arc::new(Semaphore::new(logging_task_limit)));

    let app_state: Data<AppState> = Data::new(AppState {
        cache,
        immortal_cache,
        client,
        process_start_time_seconds: SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as i64,
        process_started_at: Instant::now(),
        pg_registry: Arc::new(registry),
        jdbc_pool_cache,
        #[cfg(feature = "deadpool_experimental")]
        deadpool_registry: Arc::new(deadpool_registry),
        #[cfg(feature = "deadpool_experimental")]
        jdbc_deadpool_cache,
        gateway_force_camel_case_to_snake_case: config.get_gateway_force_camel_case_to_snake_case(),
        gateway_auto_cast_uuid_filter_values_to_text: config
            .get_gateway_auto_cast_uuid_filter_values_to_text(),
        gateway_allow_schema_names_prefixed_as_table_name: config
            .get_gateway_allow_schema_names_prefixed_as_table_name(),
        pipeline_registry: pipeline_registry.clone(),
        logging_client_name,
        gateway_auth_client_name,
        gateway_benchmark_client_name,
        gateway_api_key_fail_mode: config.get_gateway_api_key_fail_mode(),
        gateway_jdbc_allow_private_hosts: config.get_gateway_jdbc_allow_private_hosts(),
        gateway_jdbc_allowed_hosts: config.get_gateway_jdbc_allowed_hosts(),
        gateway_resilience_timeout_secs: config.get_gateway_resilience_timeout_secs(),
        gateway_resilience_read_max_retries: config.get_gateway_resilience_read_max_retries(),
        gateway_resilience_initial_backoff_ms: config.get_gateway_resilience_initial_backoff_ms(),
        gateway_admission_store_backend: config.get_gateway_admission_store_backend(),
        gateway_admission_store_fail_mode: config.get_gateway_admission_store_fail_mode(),
        prometheus_metrics_enabled: config.get_prometheus_metrics_enabled(),
        metrics_state: Arc::new(MetricsState::new()),
        gateway_insert_execution_window_ms: config.get_gateway_insert_execution_window_ms(),
        gateway_insert_window_max_batch: insert_window_settings.max_batch,
        gateway_insert_window_max_queued: insert_window_settings.max_queued,
        gateway_insert_merge_deny_tables: insert_window_settings.deny_tables.clone(),
        insert_window_coordinator: insert_window_coordinator.clone(),
        client_stats_batcher,
        gateway_log_batcher,
        logging_task_limiter,
    });

    insert_window_coordinator.bind_app_state(app_state.clone());

    Ok(Bootstrap {
        app_state,
        pipeline_registry,
    })
}

#[cfg(test)]
mod tests {
    use super::logging_task_limit_for_pool_max;

    #[test]
    fn logging_task_limit_scales_with_pool_size() {
        assert_eq!(logging_task_limit_for_pool_max(0), 1);
        assert_eq!(logging_task_limit_for_pool_max(1), 1);
        assert_eq!(logging_task_limit_for_pool_max(4), 1);
        assert_eq!(logging_task_limit_for_pool_max(8), 2);
        assert_eq!(logging_task_limit_for_pool_max(40), 10);
    }

    #[test]
    fn logging_task_limit_is_capped() {
        assert_eq!(logging_task_limit_for_pool_max(64), 16);
        assert_eq!(logging_task_limit_for_pool_max(400), 16);
    }
}