codex_memory/database/
core.rs

1use crate::error::{Error, Result};
2use sqlx::postgres::{PgPool, PgPoolOptions};
3use std::process::Command;
4use std::time::Duration;
5use url::Url;
6
7/// Create an optimized database connection pool for MCP server operations
8pub async fn create_pool(database_url: &str) -> Result<PgPool> {
9    let pool = PgPoolOptions::new()
10        .max_connections(5) // Reduced to prevent connection pool exhaustion from multiple MCP instances
11        .min_connections(0) // No minimum connections for fastest startup
12        .acquire_timeout(Duration::from_secs(10)) // Faster startup for Claude Code compatibility
13        .idle_timeout(Duration::from_secs(300)) // 5 minute idle timeout (reduced)
14        .max_lifetime(Duration::from_secs(900)) // 15 minute max lifetime (reduced for better rotation)
15        .test_before_acquire(true) // Validate connections before use
16        .connect(database_url)
17        .await?;
18
19    Ok(pool)
20}
21
22/// Simple database setup structure
23pub struct DatabaseSetup {
24    pub pool: PgPool,
25}
26
27impl DatabaseSetup {
28    /// Create new database setup
29    pub async fn new(database_url: &str) -> Result<Self> {
30        let pool = create_pool(database_url).await?;
31        Ok(Self { pool })
32    }
33
34    /// Get the pool reference
35    pub fn pool(&self) -> &PgPool {
36        &self.pool
37    }
38}
39
40/// Setup local PostgreSQL database (create database and user)
41pub async fn setup_local_database() -> Result<()> {
42    println!("🔧 Setting up local PostgreSQL database...");
43
44    // Parse the database URL from environment
45    let database_url = std::env::var("DATABASE_URL")
46        .map_err(|_| Error::Config("DATABASE_URL environment variable must be set".to_string()))?;
47
48    let parsed_url = Url::parse(&database_url)?;
49    let username = parsed_url.username();
50    let password = parsed_url.password().unwrap_or("codex_pass");
51    let database = parsed_url.path().trim_start_matches('/');
52    let host = parsed_url
53        .host_str()
54        .ok_or_else(|| Error::Config("DATABASE_URL must contain a host".to_string()))?;
55    let port = parsed_url.port().unwrap_or(5432);
56
57    println!("📍 Connecting to PostgreSQL at {}:{}", host, port);
58    println!("👤 Creating user '{}' if not exists...", username);
59
60    // Create user if not exists
61    let create_user_sql = format!(
62        "DO $$ BEGIN IF NOT EXISTS (SELECT FROM pg_user WHERE usename = '{}') THEN CREATE USER {} WITH PASSWORD '{}'; END IF; END $$;",
63        username, username, password
64    );
65
66    let mut cmd = Command::new("psql");
67    cmd.arg("-h")
68        .arg(host)
69        .arg("-p")
70        .arg(port.to_string())
71        .arg("-U")
72        .arg("postgres")
73        .arg("-c")
74        .arg(&create_user_sql);
75
76    let output = cmd.output()?;
77    if !output.status.success() {
78        let stderr = String::from_utf8_lossy(&output.stderr);
79        if !stderr.contains("already exists") {
80            eprintln!("⚠️  Failed to create user: {}", stderr);
81        }
82    } else {
83        println!("✅ User '{}' ready", username);
84    }
85
86    println!("📚 Creating database '{}' if not exists...", database);
87
88    // Create database if not exists
89    let create_db_sql = format!("CREATE DATABASE {} OWNER {};", database, username);
90
91    let mut cmd = Command::new("psql");
92    cmd.arg("-h")
93        .arg(host)
94        .arg("-p")
95        .arg(port.to_string())
96        .arg("-U")
97        .arg("postgres")
98        .arg("-c")
99        .arg(&create_db_sql);
100
101    let output = cmd.output()?;
102    if !output.status.success() {
103        let stderr = String::from_utf8_lossy(&output.stderr);
104        if !stderr.contains("already exists") {
105            eprintln!("⚠️  Failed to create database: {}", stderr);
106        }
107    } else {
108        println!("✅ Database '{}' ready", database);
109    }
110
111    // Grant privileges
112    let grant_sql = format!(
113        "GRANT ALL PRIVILEGES ON DATABASE {} TO {};",
114        database, username
115    );
116
117    let mut cmd = Command::new("psql");
118    cmd.arg("-h")
119        .arg(host)
120        .arg("-p")
121        .arg(port.to_string())
122        .arg("-U")
123        .arg("postgres")
124        .arg("-c")
125        .arg(&grant_sql);
126
127    let output = cmd.output()?;
128    if !output.status.success() {
129        let stderr = String::from_utf8_lossy(&output.stderr);
130        eprintln!("⚠️  Failed to grant privileges: {}", stderr);
131    } else {
132        println!("✅ Privileges granted to '{}'", username);
133    }
134
135    // Enable UUID extension
136    let enable_uuid_sql = "CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";";
137
138    let mut cmd = Command::new("psql");
139    cmd.arg("-h")
140        .arg(host)
141        .arg("-p")
142        .arg(port.to_string())
143        .arg("-U")
144        .arg(username)
145        .arg("-d")
146        .arg(database)
147        .arg("-c")
148        .arg(enable_uuid_sql);
149
150    if !password.is_empty() {
151        cmd.env("PGPASSWORD", password);
152    }
153
154    let output = cmd.output()?;
155    if !output.status.success() {
156        let stderr = String::from_utf8_lossy(&output.stderr);
157        if !stderr.contains("already exists") {
158            eprintln!("⚠️  Failed to enable UUID extension: {}", stderr);
159        }
160    } else {
161        println!("✅ UUID extension enabled");
162    }
163
164    // Enable pgvector extension
165    let enable_vector_sql = "CREATE EXTENSION IF NOT EXISTS vector;";
166
167    let mut cmd = Command::new("psql");
168    cmd.arg("-h")
169        .arg(host)
170        .arg("-p")
171        .arg(port.to_string())
172        .arg("-U")
173        .arg(username)
174        .arg("-d")
175        .arg(database)
176        .arg("-c")
177        .arg(enable_vector_sql);
178
179    if !password.is_empty() {
180        cmd.env("PGPASSWORD", password);
181    }
182
183    let output = cmd.output()?;
184    if !output.status.success() {
185        let stderr = String::from_utf8_lossy(&output.stderr);
186        if !stderr.contains("already exists") && !stderr.contains("could not open extension") {
187            eprintln!("⚠️  pgvector extension not available (optional for basic functionality)");
188        }
189    } else {
190        println!("✅ pgvector extension enabled (for future use)");
191    }
192
193    Ok(())
194}
195
196/// Run database migrations
197pub async fn run_migrations(pool: &PgPool) -> Result<()> {
198    // Create the memories table if it doesn't exist
199    sqlx::query(
200        r#"
201        CREATE TABLE IF NOT EXISTS memories (
202            id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
203            content TEXT NOT NULL,
204            content_hash VARCHAR(64) NOT NULL UNIQUE,
205            context TEXT NOT NULL,
206            summary TEXT NOT NULL,
207            metadata JSONB DEFAULT '{}',
208            tags TEXT[] DEFAULT '{}',
209            chunk_index INTEGER DEFAULT NULL,
210            total_chunks INTEGER DEFAULT NULL,
211            parent_id UUID DEFAULT NULL,
212            created_at TIMESTAMPTZ DEFAULT NOW(),
213            updated_at TIMESTAMPTZ DEFAULT NOW()
214        )
215    "#,
216    )
217    .execute(pool)
218    .await?;
219
220    // Create indexes for performance
221    sqlx::query("CREATE UNIQUE INDEX IF NOT EXISTS idx_content_hash ON memories(content_hash)")
222        .execute(pool)
223        .await
224        .ok();
225
226    sqlx::query("CREATE INDEX IF NOT EXISTS idx_created_at ON memories(created_at DESC)")
227        .execute(pool)
228        .await
229        .ok();
230
231    sqlx::query("CREATE INDEX IF NOT EXISTS idx_tags ON memories USING GIN(tags)")
232        .execute(pool)
233        .await
234        .ok();
235
236    sqlx::query("CREATE INDEX IF NOT EXISTS idx_parent_id ON memories(parent_id)")
237        .execute(pool)
238        .await
239        .ok();
240
241    sqlx::query("CREATE INDEX IF NOT EXISTS idx_parent_chunk ON memories(parent_id, chunk_index)")
242        .execute(pool)
243        .await
244        .ok();
245
246    sqlx::query("CREATE INDEX IF NOT EXISTS idx_summary ON memories USING GIN(to_tsvector('english', summary))")
247        .execute(pool)
248        .await
249        .ok();
250
251    // Create memory_stats view for statistics
252    sqlx::query(
253        r#"
254        CREATE OR REPLACE VIEW memory_stats AS
255        SELECT
256            (SELECT COUNT(*) FROM memories) AS total_memories,
257            (SELECT pg_size_pretty(pg_total_relation_size('memories'))) AS table_size,
258            (SELECT MAX(created_at) FROM memories) AS last_memory_created
259        "#,
260    )
261    .execute(pool)
262    .await
263    .ok();
264
265    // Note: Advanced cognitive features have been moved to the codex-dreams project
266    // tier system removed as of migration 004
267    // See: https://github.com/Ladvien/codex-dreams for memory tiering and cognitive processing
268
269    Ok(())
270}
codex_memory/database/core.rs

codex_memory/database/
core.rs
