codex_memory/database/
core.rs1use crate::error::{Error, Result};
2use sqlx::postgres::{PgPool, PgPoolOptions};
3use std::process::Command;
4use std::time::Duration;
5use url::Url;
6
7pub async fn create_pool(database_url: &str) -> Result<PgPool> {
9 let pool = PgPoolOptions::new()
10 .max_connections(50) .min_connections(2) .acquire_timeout(Duration::from_secs(30)) .idle_timeout(Duration::from_secs(600)) .max_lifetime(Duration::from_secs(1800)) .test_before_acquire(true) .after_connect(|conn, _| {
17 Box::pin(async move {
18 sqlx::query("SET statement_timeout = '30s'")
20 .execute(&mut *conn)
21 .await?;
22 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
35pub struct DatabaseSetup {
37 pub pool: PgPool,
38}
39
40impl DatabaseSetup {
41 pub async fn new(database_url: &str) -> Result<Self> {
43 let pool = create_pool(database_url).await?;
44 Ok(Self { pool })
45 }
46
47 pub fn pool(&self) -> &PgPool {
49 &self.pool
50 }
51}
52
53pub async fn setup_local_database() -> Result<()> {
55 println!("🔧 Setting up local PostgreSQL database...");
56
57 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 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 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 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 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 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
209pub async fn run_migrations(pool: &PgPool) -> Result<()> {
211 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 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 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 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 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 Ok(())
300}