systemprompt_cli/commands/cloud/tenant/docker/
database.rs

1//! Per-tenant database operations inside the shared `PostgreSQL` container.
2//!
3//! Creates, drops, and authorises tenant databases by running `psql` via
4//! `docker exec`, sanitising identifiers before they reach the SQL text.
5
6use anyhow::{Context, Result, bail};
7use systemprompt_cloud::DockerCli;
8use systemprompt_logging::CliService;
9
10use super::config::{SHARED_ADMIN_USER, SHARED_CONTAINER_NAME};
11
12pub fn sanitize_database_name(name: &str) -> String {
13    name.chars()
14        .map(|c| {
15            if c.is_ascii_alphanumeric() || c == '_' {
16                c
17            } else {
18                '_'
19            }
20        })
21        .collect()
22}
23
24pub fn create_database_for_tenant(
25    docker: &DockerCli,
26    admin_password: &str,
27    port: u16,
28    db_name: &str,
29) -> Result<()> {
30    let database_url = format!(
31        "postgres://{}:{}@localhost:{}/postgres",
32        SHARED_ADMIN_USER, admin_password, port
33    );
34
35    let safe_db_name = sanitize_database_name(db_name);
36
37    let check_query = format!(
38        "SELECT 1 FROM pg_database WHERE datname = '{}'",
39        safe_db_name
40    );
41    let check_output = docker
42        .output(&[
43            "exec",
44            SHARED_CONTAINER_NAME,
45            "psql",
46            &database_url,
47            "-tAc",
48            &check_query,
49        ])
50        .with_context(|| {
51            format!(
52                "failed to run `docker exec {SHARED_CONTAINER_NAME} psql` checking for database \
53                 {safe_db_name}"
54            )
55        })?;
56
57    let exists = !String::from_utf8_lossy(&check_output.stdout)
58        .trim()
59        .is_empty();
60
61    if exists {
62        CliService::info(&format!("Database '{}' already exists", safe_db_name));
63        return Ok(());
64    }
65
66    let create_query = format!("CREATE DATABASE \"{}\"", safe_db_name);
67    let status = docker
68        .status(&[
69            "exec",
70            SHARED_CONTAINER_NAME,
71            "psql",
72            &database_url,
73            "-c",
74            &create_query,
75        ])
76        .with_context(|| {
77            format!(
78                "failed to run `docker exec {SHARED_CONTAINER_NAME} psql` creating database \
79                 {safe_db_name}"
80            )
81        })?;
82
83    if !status.success() {
84        bail!("Failed to create database '{}'", safe_db_name);
85    }
86
87    Ok(())
88}
89
90pub fn drop_database_for_tenant(
91    docker: &DockerCli,
92    admin_password: &str,
93    port: u16,
94    db_name: &str,
95) -> Result<()> {
96    let database_url = format!(
97        "postgres://{}:{}@localhost:{}/postgres",
98        SHARED_ADMIN_USER, admin_password, port
99    );
100
101    let safe_db_name = sanitize_database_name(db_name);
102
103    let terminate_query = format!(
104        "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}' AND pid <> \
105         pg_backend_pid()",
106        safe_db_name
107    );
108    if let Err(e) = docker.status(&[
109        "exec",
110        SHARED_CONTAINER_NAME,
111        "psql",
112        &database_url,
113        "-c",
114        &terminate_query,
115    ]) {
116        tracing::debug!(
117            error = %e,
118            db = %safe_db_name,
119            "failed to run `docker exec {SHARED_CONTAINER_NAME} psql` terminating existing connections",
120        );
121    }
122
123    let drop_query = format!("DROP DATABASE IF EXISTS \"{}\"", safe_db_name);
124    let status = docker
125        .status(&[
126            "exec",
127            SHARED_CONTAINER_NAME,
128            "psql",
129            &database_url,
130            "-c",
131            &drop_query,
132        ])
133        .with_context(|| {
134            format!(
135                "failed to run `docker exec {SHARED_CONTAINER_NAME} psql` dropping database \
136                 {safe_db_name}"
137            )
138        })?;
139
140    if !status.success() {
141        bail!("Failed to drop database '{}'", safe_db_name);
142    }
143
144    Ok(())
145}
146
147pub fn ensure_admin_role(docker: &DockerCli, admin_password: &str) -> Result<()> {
148    let role_check_query = format!(
149        "SELECT 1 FROM pg_roles WHERE rolname = '{}'",
150        SHARED_ADMIN_USER
151    );
152    let role_exists =
153        !admin_psql_capture(docker, &role_check_query, "checking admin role")?.is_empty();
154
155    if role_exists {
156        let alter_password_sql = format!(
157            "ALTER ROLE \"{}\" WITH PASSWORD '{}'",
158            SHARED_ADMIN_USER,
159            admin_password.replace('\'', "''")
160        );
161        if !admin_psql_execute(docker, &alter_password_sql, "updating admin role password")? {
162            bail!("Failed to update password for role '{}'", SHARED_ADMIN_USER);
163        }
164
165        return Ok(());
166    }
167
168    let create_role_sql = format!(
169        "CREATE ROLE \"{}\" WITH LOGIN CREATEDB SUPERUSER PASSWORD '{}'",
170        SHARED_ADMIN_USER,
171        admin_password.replace('\'', "''")
172    );
173    if !admin_psql_execute(docker, &create_role_sql, "creating admin role")? {
174        bail!("Failed to create role '{}'", SHARED_ADMIN_USER);
175    }
176
177    CliService::success(&format!("Created PostgreSQL role '{}'", SHARED_ADMIN_USER));
178    Ok(())
179}
180
181fn admin_psql_capture(docker: &DockerCli, sql: &str, action: &str) -> Result<String> {
182    let output = docker
183        .output(&[
184            "exec",
185            SHARED_CONTAINER_NAME,
186            "psql",
187            "-U",
188            SHARED_ADMIN_USER,
189            "-d",
190            "postgres",
191            "-tAc",
192            sql,
193        ])
194        .with_context(|| {
195            format!("failed to run `docker exec {SHARED_CONTAINER_NAME} psql` {action}")
196        })?;
197
198    Ok(String::from_utf8_lossy(&output.stdout).trim().to_owned())
199}
200
201fn admin_psql_execute(docker: &DockerCli, sql: &str, action: &str) -> Result<bool> {
202    let status = docker
203        .status(&[
204            "exec",
205            SHARED_CONTAINER_NAME,
206            "psql",
207            "-U",
208            SHARED_ADMIN_USER,
209            "-d",
210            "postgres",
211            "-c",
212            sql,
213        ])
214        .with_context(|| {
215            format!("failed to run `docker exec {SHARED_CONTAINER_NAME} psql` {action}")
216        })?;
217
218    Ok(status.success())
219}
systemprompt_cli/commands/cloud/tenant/docker/database.rs

systemprompt_cli/commands/cloud/tenant/docker/
database.rs
