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