omni_orchestrator/initialization/
setup_clickhouse.rs

1use anyhow::Result;
2use colored::Colorize;
3use std::env;
4
5/// Initializes and tests the connection to the ClickHouse database.
6///
7/// - Loads the ClickHouse URL from environment variables or defaults.
8/// - Creates a ClickHouse client and attempts a test query to verify connectivity.
9/// - Panics if the connection test fails.
10///
11/// # Returns
12/// Returns a configured `clickhouse::Client` ready for use.
13pub async fn setup_clickhouse() -> Result<clickhouse::Client> {
14    // Load ClickHouse URL from environment or .env file
15    let clickhouse_url = env::var("CLICKHOUSE_URL").unwrap_or_else(|_| {
16        dotenv::dotenv().ok();
17        env::var("DEFAULT_CLICKHOUSE_URL").unwrap_or_else(|_| "http://localhost:8123".to_string())
18    });
19    
20    log::info!("{}", format!("ClickHouse URL: {}", clickhouse_url).blue());
21    log::info!("{}", "Initializing ClickHouse connection...".blue());
22
23    // Build the ClickHouse client
24    let clickhouse_client = clickhouse::Client::default()
25        .with_url(&clickhouse_url)
26        .with_database("default")
27        .with_user("default")
28        .with_password("your_secure_password");
29
30    // Test the connection by executing a simple query
31    match clickhouse_client.query("SELECT 1").execute().await {
32        Ok(_) => log::info!("✓ ClickHouse connection test successful"),
33        Err(e) => {
34            log::error!("ClickHouse connection test failed: {:?}", e);
35            panic!("Cannot connect to ClickHouse");
36        }
37    }
38
39    log::info!("{}", "✓ ClickHouse connection established".green());
40    log::info!("{}", "✓ ClickHouse connection pool initialized".green());
41
42    Ok(clickhouse_client)
43}