codex_memory/memory/
connection.rs

1use anyhow::Result;
2use sqlx::postgres::{PgPool, PgPoolOptions};
3use std::time::Duration;
4use tracing::info;
5
6#[derive(Debug, Clone)]
7pub struct ConnectionConfig {
8    pub host: String,
9    pub port: u16,
10    pub database: String,
11    pub username: String,
12    pub password: String,
13    pub max_connections: u32,
14    pub min_connections: u32,
15    pub connection_timeout_seconds: u64,
16    pub idle_timeout_seconds: u64,
17    pub max_lifetime_seconds: u64,
18}
19
20impl Default for ConnectionConfig {
21    fn default() -> Self {
22        Self {
23            host: "localhost".to_string(),
24            port: 5432,
25            database: "codex_memory".to_string(),
26            username: "postgres".to_string(),
27            password: "postgres".to_string(),
28            max_connections: 100,
29            min_connections: 10,
30            connection_timeout_seconds: 30,
31            idle_timeout_seconds: 600,
32            max_lifetime_seconds: 1800,
33        }
34    }
35}
36
37pub struct ConnectionPool {
38    pool: PgPool,
39    config: ConnectionConfig,
40}
41
42impl ConnectionPool {
43    pub async fn new(config: ConnectionConfig) -> Result<Self> {
44        let connection_string = format!(
45            "postgres://{}:{}@{}:{}/{}",
46            config.username, config.password, config.host, config.port, config.database
47        );
48
49        let pool = PgPoolOptions::new()
50            .max_connections(config.max_connections)
51            .min_connections(config.min_connections)
52            .acquire_timeout(Duration::from_secs(config.connection_timeout_seconds))
53            .idle_timeout(Duration::from_secs(config.idle_timeout_seconds))
54            .max_lifetime(Duration::from_secs(config.max_lifetime_seconds))
55            .connect(&connection_string)
56            .await?;
57
58        // Test the connection
59        sqlx::query("SELECT 1").fetch_one(&pool).await?;
60
61        info!(
62            "Connected to PostgreSQL at {}:{}/{} with {} max connections",
63            config.host, config.port, config.database, config.max_connections
64        );
65
66        Ok(Self { pool, config })
67    }
68
69    pub fn pool(&self) -> &PgPool {
70        &self.pool
71    }
72
73    pub async fn check_health(&self) -> Result<bool> {
74        match sqlx::query("SELECT 1").fetch_one(&self.pool).await {
75            Ok(_) => Ok(true),
76            Err(_) => Ok(false),
77        }
78    }
79
80    pub async fn get_pool_stats(&self) -> PoolStats {
81        PoolStats {
82            size: self.pool.size(),
83            idle: self.pool.num_idle() as u32,
84            max_size: self.config.max_connections,
85        }
86    }
87
88    pub async fn close(&self) {
89        self.pool.close().await;
90        info!("Connection pool closed");
91    }
92}
93
94#[derive(Debug, Clone)]
95pub struct PoolStats {
96    pub size: u32,
97    pub idle: u32,
98    pub max_size: u32,
99}
100
101impl PoolStats {
102    pub fn utilization_percentage(&self) -> f32 {
103        if self.max_size == 0 {
104            return 0.0;
105        }
106        ((self.size - self.idle) as f32 / self.max_size as f32) * 100.0
107    }
108
109    pub fn is_saturated(&self, threshold: f32) -> bool {
110        self.utilization_percentage() >= threshold
111    }
112}
113
114pub async fn create_connection_pool(config: ConnectionConfig) -> Result<PgPool> {
115    let connection_string = format!(
116        "postgres://{}:{}@{}:{}/{}",
117        config.username, config.password, config.host, config.port, config.database
118    );
119
120    let pool = PgPoolOptions::new()
121        .max_connections(config.max_connections)
122        .min_connections(config.min_connections)
123        .acquire_timeout(Duration::from_secs(config.connection_timeout_seconds))
124        .idle_timeout(Duration::from_secs(config.idle_timeout_seconds))
125        .max_lifetime(Duration::from_secs(config.max_lifetime_seconds))
126        .connect(&connection_string)
127        .await?;
128
129    Ok(pool)
130}
131
132// Simple functions for the main.rs
133pub async fn create_pool(database_url: &str, max_connections: u32) -> Result<PgPool> {
134    let pool = PgPoolOptions::new()
135        .max_connections(max_connections)
136        .connect(database_url)
137        .await?;
138
139    // Test the connection
140    sqlx::query("SELECT 1").fetch_one(&pool).await?;
141
142    info!(
143        "Connected to PostgreSQL with {} max connections",
144        max_connections
145    );
146    Ok(pool)
147}
148
149pub fn get_pool(pool: &PgPool) -> &PgPool {
150    pool
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[test]
158    fn test_pool_stats_utilization() {
159        let stats = PoolStats {
160            size: 50,
161            idle: 20,
162            max_size: 100,
163        };
164
165        assert!((stats.utilization_percentage() - 30.0).abs() < 0.01);
166        assert!(!stats.is_saturated(70.0));
167        assert!(stats.is_saturated(30.0));
168    }
169
170    #[test]
171    fn test_default_config() {
172        let config = ConnectionConfig::default();
173        assert_eq!(config.host, "localhost");
174        assert_eq!(config.port, 5432);
175        assert_eq!(config.max_connections, 100);
176    }
177}
codex_memory/memory/connection.rs

codex_memory/memory/
connection.rs
