Skip to main content

systemprompt_cli/commands/infrastructure/db/
mod.rs

1//! `db` CLI command group: schema inspection, queries, and migration tooling.
2//!
3//! [`execute`] runs commands against the invocation's
4//! [`CommandContext`](crate::context::CommandContext): migration variants are
5//! routed to the profile or standalone dispatcher depending on whether the
6//! invocation is database-scoped, and the remaining subcommands share the
7//! context's pool. Subcommands cover ad-hoc queries, schema introspection,
8//! migration apply/down/repair/squash, and the schema doctor.
9//!
10//! Copyright (c) systemprompt.io — Business Source License 1.1.
11//! See <https://systemprompt.io> for licensing details.
12
13mod 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;
21mod admin_squash;
22mod commands;
23mod dispatch;
24mod doctor;
25mod 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 {
58            sql,
59            limit,
60            offset,
61            format: _,
62        } => run_query(&query_executor, &sql, limit, offset, config).await,
63        DbCommands::Execute { sql, format: _ } => run_write(&query_executor, &sql, config).await,
64        DbCommands::Tables { filter } => {
65            schema::execute_tables(&admin_service, filter, config).await
66        },
67        DbCommands::Describe { table_name } => {
68            schema::execute_describe(&admin_service, &table_name, config).await
69        },
70        DbCommands::Info => schema::execute_info(&admin_service, config).await,
71        DbCommands::Migrate { .. }
72        | DbCommands::MigrateDown { .. }
73        | DbCommands::MigrateSquash { .. }
74        | DbCommands::MigrateRepair { .. }
75        | DbCommands::MigrateMarkApplied { .. } => {
76            bail!("migration command was not consumed by the migration dispatcher")
77        },
78        DbCommands::Migrations { cmd } => {
79            admin::execute_migrations(ctx.app_context().await?, cmd, config).await
80        },
81        DbCommands::MigratePlan { extension, json } => {
82            admin::execute_migrate_plan(
83                ctx.app_context().await?,
84                extension.as_deref(),
85                json,
86                config,
87            )
88            .await
89        },
90        DbCommands::MigrateStatus { extension, json } => {
91            admin::execute_migrate_status(
92                ctx.app_context().await?,
93                extension.as_deref(),
94                json,
95                config,
96            )
97            .await
98        },
99        DbCommands::AssignAdmin { user } => {
100            if ctx.is_database_scoped() {
101                bail!("assign-admin requires full profile context");
102            }
103            admin::execute_assign_admin(ctx.app_context().await?, &user, config).await
104        },
105        DbCommands::Status => admin::execute_status(&admin_service, config).await,
106        DbCommands::Validate => schema::execute_validate(&admin_service, config).await,
107        DbCommands::Count { table_name } => {
108            schema::execute_count(&admin_service, &table_name, config).await
109        },
110        DbCommands::Indexes { table } => {
111            introspect::execute_indexes(&admin_service, table, config).await
112        },
113        DbCommands::Size => introspect::execute_size(&admin_service, config).await,
114        DbCommands::Doctor => doctor::execute_doctor(&pool, config).await,
115    }
116}
117
118async fn connect_services(
119    ctx: &CommandContext,
120) -> Result<(DbPool, DatabaseAdminService, QueryExecutor)> {
121    let pool = ctx
122        .db_pool()
123        .await
124        .context("Failed to connect to database. Check your profile configuration.")?;
125    let write_pool = pool
126        .write_pool_arc()
127        .context("Database must be PostgreSQL")?;
128    let admin_service = DatabaseAdminService::new(Arc::clone(&write_pool));
129    let query_executor = QueryExecutor::new(write_pool);
130    Ok((pool, admin_service, query_executor))
131}
132
133async fn run_query(
134    executor: &QueryExecutor,
135    sql: &str,
136    limit: Option<u32>,
137    offset: Option<u32>,
138    config: &CliConfig,
139) -> Result<()> {
140    let params = query::QueryParams { sql, limit, offset };
141    let result = query::execute_query(executor, &params, config).await?;
142    render_result(&result, config);
143    Ok(())
144}
145
146async fn run_write(executor: &QueryExecutor, sql: &str, config: &CliConfig) -> Result<()> {
147    let result = query::execute_write(executor, sql, config).await?;
148    render_result(&result, config);
149    Ok(())
150}