systemprompt_cli/commands/infrastructure/db/
mod.rs

1mod admin;
2mod helpers;
3mod introspect;
4mod query;
5mod schema;
6mod types;
7
8use anyhow::{bail, Context, Result};
9use clap::Subcommand;
10use std::sync::Arc;
11use systemprompt_database::{DatabaseAdminService, QueryExecutor};
12use systemprompt_runtime::{AppContext, DatabaseContext};
13
14use crate::cli_settings::CliConfig;
15
16pub use types::*;
17
18#[derive(Debug, Subcommand)]
19pub enum DbCommands {
20    #[command(about = "Execute SQL query (read-only)")]
21    Query {
22        sql: String,
23        #[arg(long)]
24        limit: Option<u32>,
25        #[arg(long)]
26        offset: Option<u32>,
27        #[arg(long)]
28        format: Option<String>,
29    },
30    #[command(about = "Execute write operation (INSERT, UPDATE, DELETE)")]
31    Execute {
32        sql: String,
33        #[arg(long)]
34        format: Option<String>,
35    },
36    #[command(about = "List all tables with row counts and sizes")]
37    Tables {
38        #[arg(long, help = "Filter tables by pattern")]
39        filter: Option<String>,
40    },
41    #[command(about = "Describe table schema with columns and indexes")]
42    Describe { table_name: String },
43    #[command(about = "Show database information")]
44    Info,
45    #[command(about = "Run database migrations")]
46    Migrate,
47    #[command(about = "Show migration status and history")]
48    Migrations {
49        #[command(subcommand)]
50        cmd: MigrationsCommands,
51    },
52    #[command(about = "Assign admin role to a user")]
53    AssignAdmin { user: String },
54    #[command(about = "Show database connection status")]
55    Status,
56    #[command(about = "Validate database schema against expected tables")]
57    Validate,
58    #[command(about = "Get row count for a table")]
59    Count { table_name: String },
60    #[command(about = "List all indexes")]
61    Indexes {
62        #[arg(long, help = "Filter by table name")]
63        table: Option<String>,
64    },
65    #[command(about = "Show database and table sizes")]
66    Size,
67}
68
69#[derive(Debug, Subcommand)]
70pub enum MigrationsCommands {
71    #[command(about = "Show migration status for all extensions")]
72    Status,
73    #[command(about = "Show migration history for an extension")]
74    History {
75        #[arg(help = "Extension ID")]
76        extension: String,
77    },
78}
79
80struct DatabaseTool {
81    ctx: AppContext,
82    admin_service: DatabaseAdminService,
83    query_executor: QueryExecutor,
84}
85
86impl DatabaseTool {
87    async fn new() -> Result<Self> {
88        let ctx = AppContext::new()
89            .await
90            .context("Failed to connect to database. Check your profile configuration.")?;
91        let pool = ctx.db_pool().pool_arc()?;
92        let admin_service = DatabaseAdminService::new(Arc::clone(&pool));
93        let query_executor = QueryExecutor::new(pool);
94        Ok(Self {
95            ctx,
96            admin_service,
97            query_executor,
98        })
99    }
100}
101
102pub async fn execute(cmd: DbCommands, config: &CliConfig) -> Result<()> {
103    if matches!(cmd, DbCommands::Migrate) {
104        return admin::execute_migrate(config).await;
105    }
106
107    let db = DatabaseTool::new().await?;
108
109    match cmd {
110        DbCommands::Query {
111            sql,
112            limit,
113            offset,
114            format,
115        } => {
116            let params = query::QueryParams {
117                sql: &sql,
118                limit,
119                offset,
120                format: format.as_deref(),
121            };
122            query::execute_query(&db.query_executor, &params, config).await
123        },
124        DbCommands::Execute { sql, format } => {
125            query::execute_write(&db.query_executor, &sql, format.as_deref(), config).await
126        },
127        DbCommands::Tables { filter } => {
128            schema::execute_tables(&db.admin_service, filter, config).await
129        },
130        DbCommands::Describe { table_name } => {
131            schema::execute_describe(&db.admin_service, &table_name, config).await
132        },
133        DbCommands::Info => schema::execute_info(&db.admin_service, config).await,
134        DbCommands::Migrate => unreachable!(),
135        DbCommands::Migrations { cmd } => admin::execute_migrations(&db.ctx, cmd, config).await,
136        DbCommands::AssignAdmin { user } => {
137            admin::execute_assign_admin(&db.ctx, &user, config).await
138        },
139        DbCommands::Status => admin::execute_status(&db.admin_service, config).await,
140        DbCommands::Validate => schema::execute_validate(&db.admin_service, config).await,
141        DbCommands::Count { table_name } => {
142            schema::execute_count(&db.admin_service, &table_name, config).await
143        },
144        DbCommands::Indexes { table } => {
145            introspect::execute_indexes(&db.admin_service, table, config).await
146        },
147        DbCommands::Size => introspect::execute_size(&db.admin_service, config).await,
148    }
149}
150
151pub async fn execute_with_db(
152    cmd: DbCommands,
153    db_ctx: &DatabaseContext,
154    config: &CliConfig,
155) -> Result<()> {
156    let pool = db_ctx
157        .db_pool()
158        .pool_arc()
159        .context("Database must be PostgreSQL")?;
160    let admin_service = DatabaseAdminService::new(Arc::clone(&pool));
161    let query_executor = QueryExecutor::new(pool);
162
163    match cmd {
164        DbCommands::Query {
165            sql,
166            limit,
167            offset,
168            format,
169        } => {
170            let params = query::QueryParams {
171                sql: &sql,
172                limit,
173                offset,
174                format: format.as_deref(),
175            };
176            query::execute_query(&query_executor, &params, config).await
177        },
178        DbCommands::Execute { sql, format } => {
179            query::execute_write(&query_executor, &sql, format.as_deref(), config).await
180        },
181        DbCommands::Tables { filter } => {
182            schema::execute_tables(&admin_service, filter, config).await
183        },
184        DbCommands::Describe { table_name } => {
185            schema::execute_describe(&admin_service, &table_name, config).await
186        },
187        DbCommands::Info => schema::execute_info(&admin_service, config).await,
188        DbCommands::Migrate => admin::execute_migrate_standalone(db_ctx, config).await,
189        DbCommands::Migrations { cmd } => {
190            admin::execute_migrations_standalone(db_ctx, cmd, config).await
191        },
192        DbCommands::AssignAdmin { .. } => {
193            bail!("assign-admin requires full profile context")
194        },
195        DbCommands::Status => admin::execute_status(&admin_service, config).await,
196        DbCommands::Validate => schema::execute_validate(&admin_service, config).await,
197        DbCommands::Count { table_name } => {
198            schema::execute_count(&admin_service, &table_name, config).await
199        },
200        DbCommands::Indexes { table } => {
201            introspect::execute_indexes(&admin_service, table, config).await
202        },
203        DbCommands::Size => introspect::execute_size(&admin_service, config).await,
204    }
205}
systemprompt_cli/commands/infrastructure/db/mod.rs

systemprompt_cli/commands/infrastructure/db/
mod.rs
