1use anyhow::{anyhow, bail, Context, Result};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::fs;
5use std::path::{Path, PathBuf};
6use std::process::Command;
7use std::time::Duration;
8use systemprompt_cloud::ProjectContext;
9use systemprompt_logging::CliService;
10
11pub const SHARED_CONTAINER_NAME: &str = "systemprompt-postgres-shared";
12pub const SHARED_ADMIN_USER: &str = "systemprompt_admin";
13pub const SHARED_VOLUME_NAME: &str = "systemprompt-postgres-shared-data";
14pub const SHARED_PORT: u16 = 5432;
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct SharedContainerConfig {
18 pub admin_password: String,
19 pub port: u16,
20 pub created_at: DateTime<Utc>,
21 pub tenant_databases: Vec<TenantDatabaseMapping>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct TenantDatabaseMapping {
26 pub tenant_id: String,
27 pub database_name: String,
28}
29
30impl SharedContainerConfig {
31 pub fn new(admin_password: String, port: u16) -> Self {
32 Self {
33 admin_password,
34 port,
35 created_at: Utc::now(),
36 tenant_databases: Vec::new(),
37 }
38 }
39
40 pub fn add_tenant(&mut self, tenant_id: String, database_name: String) {
41 self.tenant_databases.push(TenantDatabaseMapping {
42 tenant_id,
43 database_name,
44 });
45 }
46
47 pub fn remove_tenant(&mut self, tenant_id: &str) -> Option<TenantDatabaseMapping> {
48 self.tenant_databases
49 .iter()
50 .position(|t| t.tenant_id == tenant_id)
51 .map(|pos| self.tenant_databases.remove(pos))
52 }
53}
54
55pub fn shared_config_path() -> PathBuf {
56 let ctx = ProjectContext::discover();
57 ctx.docker_dir().join("shared_config.json")
58}
59
60pub fn load_shared_config() -> Result<Option<SharedContainerConfig>> {
61 let path = shared_config_path();
62 if !path.exists() {
63 return Ok(None);
64 }
65 let content =
66 fs::read_to_string(&path).with_context(|| format!("Failed to read {}", path.display()))?;
67 let config: SharedContainerConfig = serde_json::from_str(&content)
68 .with_context(|| format!("Failed to parse {}", path.display()))?;
69 Ok(Some(config))
70}
71
72pub fn save_shared_config(config: &SharedContainerConfig) -> Result<()> {
73 let path = shared_config_path();
74 if let Some(parent) = path.parent() {
75 fs::create_dir_all(parent)?;
76 }
77 let content = serde_json::to_string_pretty(config)?;
78 fs::write(&path, content)?;
79
80 #[cfg(unix)]
81 {
82 use std::os::unix::fs::PermissionsExt;
83 let mut perms = fs::metadata(&path)?.permissions();
84 perms.set_mode(0o600);
85 fs::set_permissions(&path, perms)?;
86 }
87
88 Ok(())
89}
90
91pub fn is_shared_container_running() -> bool {
92 let output = Command::new("docker")
93 .args(["ps", "-q", "-f", &format!("name={}", SHARED_CONTAINER_NAME)])
94 .output();
95
96 match output {
97 Ok(out) => !String::from_utf8_lossy(&out.stdout).trim().is_empty(),
98 Err(e) => {
99 tracing::debug!(error = %e, "Failed to check shared container status");
100 false
101 },
102 }
103}
104
105pub fn get_container_password() -> Option<String> {
106 let output = Command::new("docker")
107 .args([
108 "inspect",
109 SHARED_CONTAINER_NAME,
110 "--format",
111 "{{range .Config.Env}}{{println .}}{{end}}",
112 ])
113 .output();
114
115 match output {
116 Ok(out) if out.status.success() => {
117 let env_vars = String::from_utf8_lossy(&out.stdout);
118 for line in env_vars.lines() {
119 if let Some(password) = line.strip_prefix("POSTGRES_PASSWORD=") {
120 return Some(password.to_string());
121 }
122 }
123 None
124 },
125 Ok(_out) => {
126 tracing::debug!("Docker inspect returned non-success exit code");
127 None
128 },
129 Err(e) => {
130 tracing::debug!(error = %e, "Failed to inspect container");
131 None
132 },
133 }
134}
135
136pub fn check_volume_exists() -> bool {
137 let output = Command::new("docker")
138 .args([
139 "volume",
140 "ls",
141 "-q",
142 "-f",
143 &format!("name={}", SHARED_VOLUME_NAME),
144 ])
145 .output();
146
147 match output {
148 Ok(out) => !String::from_utf8_lossy(&out.stdout).trim().is_empty(),
149 Err(e) => {
150 tracing::debug!(error = %e, "Failed to check volume existence");
151 false
152 },
153 }
154}
155
156pub fn remove_shared_volume() -> Result<()> {
157 let status = Command::new("docker")
158 .args(["volume", "rm", SHARED_VOLUME_NAME])
159 .status()
160 .context("Failed to remove PostgreSQL volume")?;
161
162 if !status.success() {
163 bail!(
164 "Failed to remove volume '{}'. Is a container still using it?",
165 SHARED_VOLUME_NAME
166 );
167 }
168
169 Ok(())
170}
171
172pub fn generate_shared_postgres_compose(password: &str, port: u16) -> String {
173 format!(
174 r#"# systemprompt.io Shared PostgreSQL Container
175# Generated by: systemprompt cloud tenant create
176# Manage: docker compose -f .systemprompt/docker/shared.yaml up/down
177
178services:
179 postgres:
180 image: postgres:18-alpine
181 container_name: {container_name}
182 restart: unless-stopped
183 environment:
184 POSTGRES_USER: {admin_user}
185 POSTGRES_PASSWORD: {password}
186 POSTGRES_DB: postgres
187 ports:
188 - "{port}:5432"
189 volumes:
190 - {volume_name}:/var/lib/postgresql
191 healthcheck:
192 test: ["CMD-SHELL", "pg_isready -U {admin_user}"]
193 interval: 5s
194 timeout: 5s
195 retries: 5
196
197volumes:
198 {volume_name}:
199 name: {volume_name}
200"#,
201 container_name = SHARED_CONTAINER_NAME,
202 admin_user = SHARED_ADMIN_USER,
203 password = password,
204 port = port,
205 volume_name = SHARED_VOLUME_NAME
206 )
207}
208
209pub fn generate_admin_password() -> String {
210 use std::time::{SystemTime, UNIX_EPOCH};
211 let timestamp = SystemTime::now()
212 .duration_since(UNIX_EPOCH)
213 .map(|d| d.as_nanos())
214 .unwrap_or(1);
215 let random_part = format!("{:x}{:x}", timestamp, timestamp.wrapping_mul(31337));
216 random_part.chars().take(32).collect()
217}
218
219pub async fn create_database_for_tenant(
220 admin_password: &str,
221 port: u16,
222 db_name: &str,
223) -> Result<()> {
224 let database_url = format!(
225 "postgres://{}:{}@localhost:{}/postgres",
226 SHARED_ADMIN_USER, admin_password, port
227 );
228
229 let safe_db_name = sanitize_database_name(db_name);
230
231 let check_query = format!(
232 "SELECT 1 FROM pg_database WHERE datname = '{}'",
233 safe_db_name
234 );
235 let check_output = Command::new("docker")
236 .args([
237 "exec",
238 SHARED_CONTAINER_NAME,
239 "psql",
240 &database_url,
241 "-tAc",
242 &check_query,
243 ])
244 .output()
245 .context("Failed to check if database exists")?;
246
247 let exists = !String::from_utf8_lossy(&check_output.stdout)
248 .trim()
249 .is_empty();
250
251 if exists {
252 CliService::info(&format!("Database '{}' already exists", safe_db_name));
253 return Ok(());
254 }
255
256 let create_query = format!("CREATE DATABASE \"{}\"", safe_db_name);
257 let status = Command::new("docker")
258 .args([
259 "exec",
260 SHARED_CONTAINER_NAME,
261 "psql",
262 &database_url,
263 "-c",
264 &create_query,
265 ])
266 .status()
267 .context("Failed to create database")?;
268
269 if !status.success() {
270 bail!("Failed to create database '{}'", safe_db_name);
271 }
272
273 Ok(())
274}
275
276pub async fn drop_database_for_tenant(
277 admin_password: &str,
278 port: u16,
279 db_name: &str,
280) -> Result<()> {
281 let database_url = format!(
282 "postgres://{}:{}@localhost:{}/postgres",
283 SHARED_ADMIN_USER, admin_password, port
284 );
285
286 let safe_db_name = sanitize_database_name(db_name);
287
288 let terminate_query = format!(
289 "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}' AND pid <> \
290 pg_backend_pid()",
291 safe_db_name
292 );
293 if let Err(e) = Command::new("docker")
294 .args([
295 "exec",
296 SHARED_CONTAINER_NAME,
297 "psql",
298 &database_url,
299 "-c",
300 &terminate_query,
301 ])
302 .status()
303 {
304 tracing::debug!(error = %e, "Failed to terminate existing connections");
305 }
306
307 let drop_query = format!("DROP DATABASE IF EXISTS \"{}\"", safe_db_name);
308 let status = Command::new("docker")
309 .args([
310 "exec",
311 SHARED_CONTAINER_NAME,
312 "psql",
313 &database_url,
314 "-c",
315 &drop_query,
316 ])
317 .status()
318 .context("Failed to drop database")?;
319
320 if !status.success() {
321 bail!("Failed to drop database '{}'", safe_db_name);
322 }
323
324 Ok(())
325}
326
327pub fn stop_shared_container() -> Result<()> {
328 let ctx = ProjectContext::discover();
329 let compose_path = ctx.docker_dir().join("shared.yaml");
330
331 if compose_path.exists() {
332 let compose_path_str = compose_path
333 .to_str()
334 .ok_or_else(|| anyhow!("Invalid compose path"))?;
335
336 CliService::info("Stopping shared PostgreSQL container...");
337 let status = Command::new("docker")
338 .args(["compose", "-f", compose_path_str, "down", "-v"])
339 .status()
340 .context("Failed to stop shared container")?;
341
342 if !status.success() {
343 CliService::warning("Failed to stop container via compose, trying direct stop");
344 }
345 }
346
347 let output = Command::new("docker")
348 .args([
349 "ps",
350 "-aq",
351 "-f",
352 &format!("name={}", SHARED_CONTAINER_NAME),
353 ])
354 .output()?;
355
356 let container_id = String::from_utf8_lossy(&output.stdout).trim().to_string();
357 if !container_id.is_empty() {
358 Command::new("docker")
359 .args(["stop", &container_id])
360 .status()?;
361 Command::new("docker")
362 .args(["rm", &container_id])
363 .status()?;
364 }
365
366 let config_path = shared_config_path();
367 if config_path.exists() {
368 fs::remove_file(&config_path)?;
369 }
370
371 CliService::success("Shared PostgreSQL container removed");
372 Ok(())
373}
374
375fn sanitize_database_name(name: &str) -> String {
376 name.chars()
377 .map(|c| {
378 if c.is_ascii_alphanumeric() || c == '_' {
379 c
380 } else {
381 '_'
382 }
383 })
384 .collect()
385}
386
387pub fn nanoid() -> String {
388 use std::time::{SystemTime, UNIX_EPOCH};
389 let timestamp = SystemTime::now()
390 .duration_since(UNIX_EPOCH)
391 .map(|d| d.as_millis())
392 .unwrap_or(1);
393 format!("{:x}", timestamp)
394}
395
396pub async fn wait_for_postgres_healthy(compose_path: &Path, timeout_secs: u64) -> Result<()> {
397 let start = std::time::Instant::now();
398 let compose_path_str = compose_path
399 .to_str()
400 .ok_or_else(|| anyhow!("Invalid compose path"))?;
401
402 loop {
403 let output = Command::new("docker")
404 .args([
405 "compose",
406 "-f",
407 compose_path_str,
408 "ps",
409 "--format",
410 "{{.Health}}",
411 ])
412 .output()
413 .context("Failed to check container health")?;
414
415 let health = String::from_utf8_lossy(&output.stdout).trim().to_string();
416
417 if health.contains("healthy") {
418 return Ok(());
419 }
420
421 if start.elapsed().as_secs() > timeout_secs {
422 bail!(
423 "Timeout waiting for PostgreSQL to become healthy.\nCheck logs with: docker \
424 compose -f {} logs",
425 compose_path.display()
426 );
427 }
428
429 tokio::time::sleep(Duration::from_secs(2)).await;
430 }
431}
432
433pub fn ensure_admin_role(admin_password: &str) -> Result<()> {
434 let role_check_query = format!(
435 "SELECT 1 FROM pg_roles WHERE rolname = '{}'",
436 SHARED_ADMIN_USER
437 );
438 let check_output = Command::new("docker")
439 .args([
440 "exec",
441 SHARED_CONTAINER_NAME,
442 "psql",
443 "-U",
444 SHARED_ADMIN_USER,
445 "-d",
446 "postgres",
447 "-tAc",
448 &role_check_query,
449 ])
450 .output()
451 .context("Failed to check if admin role exists")?;
452
453 let role_exists = !String::from_utf8_lossy(&check_output.stdout)
454 .trim()
455 .is_empty();
456
457 if role_exists {
458 let alter_password_sql = format!(
459 "ALTER ROLE \"{}\" WITH PASSWORD '{}'",
460 SHARED_ADMIN_USER,
461 admin_password.replace('\'', "''")
462 );
463 let status = Command::new("docker")
464 .args([
465 "exec",
466 SHARED_CONTAINER_NAME,
467 "psql",
468 "-U",
469 SHARED_ADMIN_USER,
470 "-d",
471 "postgres",
472 "-c",
473 &alter_password_sql,
474 ])
475 .status()
476 .context("Failed to update admin role password")?;
477
478 if !status.success() {
479 bail!("Failed to update password for role '{}'", SHARED_ADMIN_USER);
480 }
481
482 return Ok(());
483 }
484
485 let create_role_sql = format!(
486 "CREATE ROLE \"{}\" WITH LOGIN CREATEDB SUPERUSER PASSWORD '{}'",
487 SHARED_ADMIN_USER,
488 admin_password.replace('\'', "''")
489 );
490 let status = Command::new("docker")
491 .args([
492 "exec",
493 SHARED_CONTAINER_NAME,
494 "psql",
495 "-U",
496 SHARED_ADMIN_USER,
497 "-d",
498 "postgres",
499 "-c",
500 &create_role_sql,
501 ])
502 .status()
503 .context("Failed to create admin role")?;
504
505 if !status.success() {
506 bail!("Failed to create role '{}'", SHARED_ADMIN_USER);
507 }
508
509 CliService::success(&format!("Created PostgreSQL role '{}'", SHARED_ADMIN_USER));
510 Ok(())
511}