r2_data2/
state.rs

1use crate::{AppConfig, DbPool, db::PoolHandler, error::AppError, handlers::FullSchema};
2use moka::future::Cache;
3use papaya::HashMap;
4use rig::providers::openai as rig_openai;
5use std::{ops::Deref, sync::Arc, time::Duration};
6use tracing::{error, info}; // Import with alias
7
8#[derive(Clone)]
9pub struct AppState(Arc<AppStateInner>);
10
11pub struct AppStateInner {
12    pub config: AppConfig,
13    pub pools: Arc<HashMap<String, DbPool>>,
14    // Cache for the full schema, storing the Result wrapped in Arc
15    pub schema_cache: Cache<String, Arc<Result<FullSchema, AppError>>>,
16    // Add OpenAI client from rig-core
17    pub openai_client: rig_openai::Client,
18}
19
20// Manual Debug implementation because sqlx Pools don't implement Debug
21impl std::fmt::Debug for AppStateInner {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.debug_struct("AppStateInner")
24            .field("config", &self.config)
25            .field("db_pools_count", &self.pools.len()) // Only show count
26            // Do not display the cache content
27            // Do not display the openai client details
28            .finish_non_exhaustive()
29    }
30}
31
32// Manual Debug implementation for AppState wrapper
33impl std::fmt::Debug for AppState {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        f.debug_tuple("AppState").field(&self.0).finish()
36    }
37}
38
39impl Deref for AppState {
40    type Target = AppStateInner;
41
42    fn deref(&self) -> &Self::Target {
43        &self.0
44    }
45}
46
47impl AppState {
48    pub async fn new(config: AppConfig) -> Result<Self, anyhow::Error> {
49        let pools = HashMap::new();
50
51        for db_config in &config.databases {
52            info!(
53                "Connecting to database '{}' (type: {})",
54                db_config.name, db_config.db_type
55            );
56            match DbPool::try_new(db_config).await {
57                Ok(pool) => {
58                    pools.pin().insert(db_config.name.clone(), pool);
59                }
60                Err(e) => {
61                    error!("Failed to connect to database '{}': {}", db_config.name, e);
62                }
63            }
64        }
65        info!("Database connections established.");
66
67        // Create the schema cache
68        let schema_cache = Cache::builder()
69            // Time to live (TTL): 10 minutes
70            .time_to_live(Duration::from_secs(10 * 60))
71            // Max capacity (optional, e.g., only 1 entry needed)
72            .max_capacity(1)
73            .build();
74
75        // Initialize OpenAI client using environment variable
76        // This will panic if OPENAI_API_KEY is not set.
77        // Consider adding error handling or configuration check earlier.
78        info!("Initializing OpenAI client from environment...");
79        let openai_client = rig_openai::Client::from_env();
80        info!("OpenAI client initialized.");
81
82        let inner = AppStateInner {
83            config,
84            pools: Arc::new(pools),
85            schema_cache,
86            openai_client, // Add client to state
87        };
88        Ok(Self(Arc::new(inner)))
89    }
90
91    #[cfg(test)]
92    pub fn new_for_test(config: AppConfig) -> Self {
93        // Create empty/dummy versions of fields not needed for config-only tests
94        let pools = Arc::new(HashMap::new());
95        let schema_cache = Cache::builder().build();
96        // Initialize client from env - it won't be used in config-only tests.
97        // This might panic if OPENAI_API_KEY is *required* and *not set* during init,
98        // but typically `from_env` reads it lazily or handles its absence until first use.
99        let openai_client = rig_openai::Client::from_env();
100
101        let inner = AppStateInner {
102            config,
103            pools,
104            schema_cache,
105            openai_client,
106        };
107        Self(Arc::new(inner))
108    }
109}
r2_data2/state.rs

r2_data2/
state.rs
