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