Skip to main content

systemprompt_cli/commands/cloud/db/
mod.rs

1mod backup;
2mod restore;
3
4use anyhow::{Context, Result, anyhow, bail};
5use clap::{Subcommand, ValueEnum};
6use std::path::PathBuf;
7use std::process::Command;
8use systemprompt_cloud::{ProfilePath, ProjectContext};
9use systemprompt_runtime::DatabaseContext;
10
11use crate::cli_settings::CliConfig;
12use crate::commands::infrastructure::db;
13
14#[derive(Debug, Subcommand)]
15pub enum CloudDbCommands {
16    #[command(about = "Run migrations on cloud database")]
17    Migrate {
18        #[arg(long, help = "Profile name")]
19        profile: String,
20    },
21
22    #[command(about = "Execute SQL query (read-only) on cloud database")]
23    Query {
24        #[arg(long, help = "Profile name")]
25        profile: String,
26        sql: String,
27        #[arg(long)]
28        limit: Option<u32>,
29        #[arg(long)]
30        offset: Option<u32>,
31        #[arg(long)]
32        format: Option<String>,
33    },
34
35    #[command(about = "Execute write operation on cloud database")]
36    Execute {
37        #[arg(long, help = "Profile name")]
38        profile: String,
39        sql: String,
40        #[arg(long)]
41        format: Option<String>,
42    },
43
44    #[command(about = "Validate cloud database schema")]
45    Validate {
46        #[arg(long, help = "Profile name")]
47        profile: String,
48    },
49
50    #[command(about = "Show cloud database connection status")]
51    Status {
52        #[arg(long, help = "Profile name")]
53        profile: String,
54    },
55
56    #[command(about = "Show cloud database info")]
57    Info {
58        #[arg(long, help = "Profile name")]
59        profile: String,
60    },
61
62    #[command(about = "List all tables in cloud database")]
63    Tables {
64        #[arg(long, help = "Profile name")]
65        profile: String,
66        #[arg(long, help = "Filter tables by pattern")]
67        filter: Option<String>,
68    },
69
70    #[command(about = "Describe table schema in cloud database")]
71    Describe {
72        #[arg(long, help = "Profile name")]
73        profile: String,
74        table_name: String,
75    },
76
77    #[command(about = "Get row count for a table in cloud database")]
78    Count {
79        #[arg(long, help = "Profile name")]
80        profile: String,
81        table_name: String,
82    },
83
84    #[command(about = "List all indexes in cloud database")]
85    Indexes {
86        #[arg(long, help = "Profile name")]
87        profile: String,
88        #[arg(long, help = "Filter by table name")]
89        table: Option<String>,
90    },
91
92    #[command(about = "Show cloud database and table sizes")]
93    Size {
94        #[arg(long, help = "Profile name")]
95        profile: String,
96    },
97
98    #[command(about = "Backup cloud database using pg_dump")]
99    Backup {
100        #[arg(long, help = "Profile name")]
101        profile: String,
102
103        #[arg(
104            long,
105            default_value = "custom",
106            help = "Backup format: custom, sql, directory"
107        )]
108        format: BackupFormat,
109
110        #[arg(
111            long,
112            help = "Output file path (default: backups/<profile>-<timestamp>.<ext>)"
113        )]
114        output: Option<String>,
115    },
116
117    #[command(about = "Restore cloud database from a backup file")]
118    Restore {
119        #[arg(long, help = "Profile name")]
120        profile: String,
121
122        #[arg(help = "Path to backup file")]
123        file: String,
124
125        #[arg(short = 'y', long, help = "Skip confirmation prompt")]
126        yes: bool,
127    },
128}
129
130#[derive(Debug, Clone, Copy, ValueEnum)]
131pub enum BackupFormat {
132    #[value(help = "pg_dump custom format (-Fc), supports parallel restore")]
133    Custom,
134    #[value(help = "Plain SQL text format (-Fp), human-readable")]
135    Sql,
136    #[value(help = "Directory format (-Fd), supports parallel dump and restore")]
137    Directory,
138}
139
140impl CloudDbCommands {
141    fn profile_name(&self) -> &str {
142        match self {
143            Self::Migrate { profile }
144            | Self::Query { profile, .. }
145            | Self::Execute { profile, .. }
146            | Self::Validate { profile }
147            | Self::Status { profile }
148            | Self::Info { profile }
149            | Self::Tables { profile, .. }
150            | Self::Describe { profile, .. }
151            | Self::Count { profile, .. }
152            | Self::Indexes { profile, .. }
153            | Self::Size { profile }
154            | Self::Backup { profile, .. }
155            | Self::Restore { profile, .. } => profile,
156        }
157    }
158
159    fn into_db_command(self) -> Option<db::DbCommands> {
160        match self {
161            Self::Migrate { .. } => Some(db::DbCommands::Migrate),
162            Self::Query {
163                sql,
164                limit,
165                offset,
166                format,
167                ..
168            } => Some(db::DbCommands::Query {
169                sql,
170                limit,
171                offset,
172                format,
173            }),
174            Self::Execute { sql, format, .. } => Some(db::DbCommands::Execute { sql, format }),
175            Self::Validate { .. } => Some(db::DbCommands::Validate),
176            Self::Status { .. } => Some(db::DbCommands::Status),
177            Self::Info { .. } => Some(db::DbCommands::Info),
178            Self::Tables { filter, .. } => Some(db::DbCommands::Tables { filter }),
179            Self::Describe { table_name, .. } => Some(db::DbCommands::Describe { table_name }),
180            Self::Count { table_name, .. } => Some(db::DbCommands::Count { table_name }),
181            Self::Indexes { table, .. } => Some(db::DbCommands::Indexes { table }),
182            Self::Size { .. } => Some(db::DbCommands::Size),
183            Self::Backup { .. } | Self::Restore { .. } => None,
184        }
185    }
186}
187
188pub async fn execute(cmd: CloudDbCommands, config: &CliConfig) -> Result<()> {
189    let profile_name = cmd.profile_name().to_string();
190    let db_url = load_cloud_database_url(&profile_name)?;
191    execute_inner(cmd, &profile_name, &db_url, config).await
192}
193
194pub async fn execute_with_database_url(
195    cmd: CloudDbCommands,
196    database_url: &str,
197    config: &CliConfig,
198) -> Result<()> {
199    let profile_name = cmd.profile_name().to_string();
200    execute_inner(cmd, &profile_name, database_url, config).await
201}
202
203async fn execute_inner(
204    cmd: CloudDbCommands,
205    profile_name: &str,
206    db_url: &str,
207    config: &CliConfig,
208) -> Result<()> {
209    match &cmd {
210        CloudDbCommands::Backup { format, output, .. } => {
211            return backup::execute(profile_name, db_url, *format, output.as_deref());
212        },
213        CloudDbCommands::Restore { file, yes, .. } => {
214            return restore::execute(profile_name, db_url, file, *yes, config);
215        },
216        _ => {},
217    }
218
219    let db_ctx = DatabaseContext::from_url(db_url).await?;
220    let db_cmd = cmd
221        .into_db_command()
222        .ok_or_else(|| anyhow!("Unexpected command variant"))?;
223
224    db::execute_with_db(db_cmd, &db_ctx, config).await
225}
226
227fn load_cloud_database_url(profile_name: &str) -> Result<String> {
228    let ctx = ProjectContext::discover();
229    let profile_dir = ctx.profile_dir(profile_name);
230
231    if !profile_dir.exists() {
232        return Err(anyhow!("Profile '{}' not found", profile_name));
233    }
234
235    let secrets_path = ProfilePath::Secrets.resolve(&profile_dir);
236    let secrets = systemprompt_models::Secrets::load_from_path(&secrets_path)
237        .with_context(|| format!("Failed to load secrets for profile '{}'", profile_name))?;
238
239    Ok(secrets.effective_database_url(true).to_string())
240}
241
242fn ensure_pg_tool(tool: &str) -> Result<()> {
243    match Command::new(tool).arg("--version").output() {
244        Ok(output) if output.status.success() => Ok(()),
245        _ => bail!(
246            "'{}' not found. Install PostgreSQL client tools:\n  apt install postgresql-client",
247            tool
248        ),
249    }
250}
251
252fn find_pg_dump() -> Result<PathBuf> {
253    for version in [17, 16, 15, 14] {
254        let path = PathBuf::from(format!("/usr/lib/postgresql/{}/bin/pg_dump", version));
255        if path.exists() {
256            return Ok(path);
257        }
258    }
259    ensure_pg_tool("pg_dump")?;
260    Ok(PathBuf::from("pg_dump"))
261}
262
263fn find_pg_restore() -> Result<PathBuf> {
264    for version in [17, 16, 15, 14] {
265        let path = PathBuf::from(format!("/usr/lib/postgresql/{}/bin/pg_restore", version));
266        if path.exists() {
267            return Ok(path);
268        }
269    }
270    ensure_pg_tool("pg_restore")?;
271    Ok(PathBuf::from("pg_restore"))
272}
273
274fn adjust_ssl_mode(database_url: &str) -> String {
275    database_url.replace("sslmode=require", "sslmode=prefer")
276}