systemprompt_cli/commands/infrastructure/db/
mod.rs1mod admin;
14mod admin_migrate;
15mod admin_migrate_down;
16mod admin_migrate_mark_applied;
17mod admin_migrate_plan;
18mod admin_migrate_repair;
19mod admin_migrate_status;
20mod admin_migrations;
21pub mod admin_squash;
22mod commands;
23mod dispatch;
24mod doctor;
25pub mod helpers;
26mod introspect;
27mod query;
28mod schema;
29mod types;
30
31use anyhow::{Context, Result, bail};
32use std::sync::Arc;
33use systemprompt_database::{DatabaseAdminService, DbPool, QueryExecutor};
34
35use crate::cli_settings::CliConfig;
36use crate::context::CommandContext;
37use crate::shared::render_result;
38use dispatch::{dispatch_profile_migration, dispatch_standalone_migration};
39
40pub use commands::{DbCommands, MigrationsCommands};
41pub use types::*;
42
43pub(crate) use helpers::format_bytes;
44
45pub async fn execute(cmd: DbCommands, ctx: &CommandContext) -> Result<()> {
46 let config = &ctx.cli;
47 let Some(cmd) = (match ctx.database_context() {
48 Some(db_ctx) => dispatch_standalone_migration(cmd, db_ctx, config).await?,
49 None => dispatch_profile_migration(cmd, config).await?,
50 }) else {
51 return Ok(());
52 };
53
54 let (pool, admin_service, query_executor) = connect_services(ctx).await?;
55
56 match cmd {
57 DbCommands::Query { sql, limit, offset } => {
58 run_query(&query_executor, &sql, limit, offset, config).await
59 },
60 DbCommands::Execute { sql } => run_write(&query_executor, &sql, config).await,
61 DbCommands::Tables { filter } => {
62 schema::execute_tables(&admin_service, filter, config).await
63 },
64 DbCommands::Describe { table_name } => {
65 schema::execute_describe(&admin_service, &table_name, config).await
66 },
67 DbCommands::Info => schema::execute_info(&admin_service, config).await,
68 DbCommands::Migrate { .. }
69 | DbCommands::MigrateDown { .. }
70 | DbCommands::MigrateSquash { .. }
71 | DbCommands::MigrateRepair { .. }
72 | DbCommands::MigrateMarkApplied { .. } => {
73 bail!("migration command was not consumed by the migration dispatcher")
74 },
75 DbCommands::Migrations { cmd } => {
76 admin::execute_migrations(ctx.app_context().await?, cmd, config).await
77 },
78 DbCommands::MigratePlan { extension, json } => {
79 admin::execute_migrate_plan(
80 ctx.app_context().await?,
81 extension.as_deref(),
82 json,
83 config,
84 )
85 .await
86 },
87 DbCommands::MigrateStatus { extension, json } => {
88 admin::execute_migrate_status(
89 ctx.app_context().await?,
90 extension.as_deref(),
91 json,
92 config,
93 )
94 .await
95 },
96 DbCommands::AssignAdmin { user } => {
97 if ctx.is_database_scoped() {
98 bail!("assign-admin requires full profile context");
99 }
100 admin::execute_assign_admin(ctx.app_context().await?, &user, config).await
101 },
102 DbCommands::Status => admin::execute_status(&admin_service, config).await,
103 DbCommands::Count { table_name } => {
104 schema::execute_count(&admin_service, &table_name, config).await
105 },
106 DbCommands::Indexes { table } => {
107 introspect::execute_indexes(&admin_service, table, config).await
108 },
109 DbCommands::Size => introspect::execute_size(&admin_service, config).await,
110 DbCommands::Doctor => doctor::execute_doctor(&pool, config).await,
111 }
112}
113
114async fn connect_services(
115 ctx: &CommandContext,
116) -> Result<(DbPool, DatabaseAdminService, QueryExecutor)> {
117 let pool = ctx
118 .db_pool()
119 .await
120 .context("Failed to connect to database. Check your profile configuration.")?;
121 let write_pool = pool
122 .write_pool_arc()
123 .context("Database must be PostgreSQL")?;
124 let admin_service = DatabaseAdminService::new(Arc::clone(&write_pool));
125 let query_executor = QueryExecutor::new(write_pool);
126 Ok((pool, admin_service, query_executor))
127}
128
129async fn run_query(
130 executor: &QueryExecutor,
131 sql: &str,
132 limit: Option<u32>,
133 offset: Option<u32>,
134 config: &CliConfig,
135) -> Result<()> {
136 let params = query::QueryParams { sql, limit, offset };
137 let result = query::execute_query(executor, ¶ms, config).await?;
138 render_result(&result, config);
139 Ok(())
140}
141
142async fn run_write(executor: &QueryExecutor, sql: &str, config: &CliConfig) -> Result<()> {
143 let result = query::execute_write(executor, sql, config).await?;
144 render_result(&result, config);
145 Ok(())
146}