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

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