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                allow_checksum_drift: false,
163            }),
164            Self::Query {
165                sql,
166                limit,
167                offset,
168                format,
169                ..
170            } => Some(db::DbCommands::Query {
171                sql,
172                limit,
173                offset,
174                format,
175            }),
176            Self::Execute { sql, format, .. } => Some(db::DbCommands::Execute { sql, format }),
177            Self::Validate { .. } => Some(db::DbCommands::Validate),
178            Self::Status { .. } => Some(db::DbCommands::Status),
179            Self::Info { .. } => Some(db::DbCommands::Info),
180            Self::Tables { filter, .. } => Some(db::DbCommands::Tables { filter }),
181            Self::Describe { table_name, .. } => Some(db::DbCommands::Describe { table_name }),
182            Self::Count { table_name, .. } => Some(db::DbCommands::Count { table_name }),
183            Self::Indexes { table, .. } => Some(db::DbCommands::Indexes { table }),
184            Self::Size { .. } => Some(db::DbCommands::Size),
185            Self::Backup { .. } | Self::Restore { .. } => None,
186        }
187    }
188}
189
190pub async fn execute(cmd: CloudDbCommands, config: &CliConfig) -> Result<()> {
191    let profile_name = cmd.profile_name().to_string();
192    let db_url = load_cloud_database_url(&profile_name)?;
193    execute_inner(cmd, &profile_name, &db_url, config).await
194}
195
196pub async fn execute_with_database_url(
197    cmd: CloudDbCommands,
198    database_url: &str,
199    config: &CliConfig,
200) -> Result<()> {
201    let profile_name = cmd.profile_name().to_string();
202    execute_inner(cmd, &profile_name, database_url, config).await
203}
204
205async fn execute_inner(
206    cmd: CloudDbCommands,
207    profile_name: &str,
208    db_url: &str,
209    config: &CliConfig,
210) -> Result<()> {
211    match &cmd {
212        CloudDbCommands::Backup { format, output, .. } => {
213            return backup::execute(profile_name, db_url, *format, output.as_deref());
214        },
215        CloudDbCommands::Restore { file, yes, .. } => {
216            return restore::execute(profile_name, db_url, file, *yes, config);
217        },
218        _ => {},
219    }
220
221    let db_ctx = DatabaseContext::from_url(db_url).await?;
222    let db_cmd = cmd
223        .into_db_command()
224        .ok_or_else(|| anyhow!("Unexpected command variant"))?;
225
226    db::execute_with_db(db_cmd, &db_ctx, config).await
227}
228
229fn load_cloud_database_url(profile_name: &str) -> Result<String> {
230    let ctx = ProjectContext::discover();
231    let profile_dir = ctx.profile_dir(profile_name);
232
233    if !profile_dir.exists() {
234        return Err(anyhow!("Profile '{}' not found", profile_name));
235    }
236
237    let secrets_path = ProfilePath::Secrets.resolve(&profile_dir);
238    let secrets = systemprompt_config::load_secrets_from_path(&secrets_path)
239        .with_context(|| format!("Failed to load secrets for profile '{}'", profile_name))?;
240
241    Ok(secrets.effective_database_url(true).to_string())
242}
243
244fn ensure_pg_tool(tool: &str) -> Result<()> {
245    match Command::new(tool).arg("--version").output() {
246        Ok(output) if output.status.success() => Ok(()),
247        _ => bail!(
248            "'{}' not found. Install PostgreSQL client tools:\n  apt install postgresql-client",
249            tool
250        ),
251    }
252}
253
254fn find_pg_dump() -> Result<PathBuf> {
255    for version in [17, 16, 15, 14] {
256        let path = PathBuf::from(format!("/usr/lib/postgresql/{}/bin/pg_dump", version));
257        if path.exists() {
258            return Ok(path);
259        }
260    }
261    ensure_pg_tool("pg_dump")?;
262    Ok(PathBuf::from("pg_dump"))
263}
264
265fn find_pg_restore() -> Result<PathBuf> {
266    for version in [17, 16, 15, 14] {
267        let path = PathBuf::from(format!("/usr/lib/postgresql/{}/bin/pg_restore", version));
268        if path.exists() {
269            return Ok(path);
270        }
271    }
272    ensure_pg_tool("pg_restore")?;
273    Ok(PathBuf::from("pg_restore"))
274}
275
276fn adjust_ssl_mode(database_url: &str) -> String {
277    database_url.replace("sslmode=require", "sslmode=prefer")
278}
systemprompt_cli/commands/cloud/db/mod.rs

systemprompt_cli/commands/cloud/db/
mod.rs
