systemprompt_database/admin/
query_executor.rs

1//! Query executor used by the CLI's `infra db query` and `db exec` commands.
2//!
3//! Part of the documented sqlx allowlist: SQL is supplied dynamically by
4//! the operator and validated through [`AdminSql`].
5//!
6//! Copyright (c) systemprompt.io — Business Source License 1.1.
7//! See <https://systemprompt.io> for licensing details.
8
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use sqlx::postgres::PgPool;
13use sqlx::{Column, Row};
14use thiserror::Error;
15
16use crate::admin::admin_sql::{AdminSql, AdminSqlError, DEFAULT_READONLY_ROW_LIMIT};
17use crate::models::QueryResult;
18
19#[derive(Error, Debug)]
20pub enum QueryExecutorError {
21    #[error(
22        "Write query not allowed in read-only mode: only SELECT, WITH, EXPLAIN, SHOW, TABLE, and \
23         VALUES queries are permitted"
24    )]
25    WriteQueryNotAllowed,
26
27    #[error("Invalid admin SQL: {0}")]
28    InvalidSql(#[from] AdminSqlError),
29
30    #[error("Query execution failed: {0}")]
31    ExecutionFailed(#[from] sqlx::Error),
32}
33
34#[derive(Debug)]
35pub struct QueryExecutor {
36    pool: Arc<PgPool>,
37}
38
39impl QueryExecutor {
40    pub const fn new(pool: Arc<PgPool>) -> Self {
41        Self { pool }
42    }
43
44    pub async fn execute_readonly(
45        &self,
46        raw_sql: &str,
47        row_limit: Option<usize>,
48    ) -> Result<QueryResult, QueryExecutorError> {
49        let sql = AdminSql::parse_readonly(raw_sql)?;
50        self.execute(sql, row_limit.unwrap_or(DEFAULT_READONLY_ROW_LIMIT))
51            .await
52    }
53
54    pub async fn execute_write(&self, raw_sql: &str) -> Result<QueryResult, QueryExecutorError> {
55        let sql = AdminSql::parse_unrestricted(raw_sql)?;
56        self.execute(sql, usize::MAX).await
57    }
58
59    pub(super) async fn execute(
60        &self,
61        sql: AdminSql,
62        row_limit: usize,
63    ) -> Result<QueryResult, QueryExecutorError> {
64        let start = std::time::Instant::now();
65
66        let rows = sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
67            .fetch_all(&*self.pool)
68            .await?;
69        let execution_time = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
70
71        let columns = rows.first().map_or_else(Vec::new, |first_row| {
72            first_row
73                .columns()
74                .iter()
75                .map(|c| c.name().to_owned())
76                .collect()
77        });
78
79        let total_rows = rows.len();
80        let capped_rows = rows.iter().take(row_limit);
81        let mut result_rows = Vec::with_capacity(total_rows.min(row_limit));
82
83        for row in capped_rows {
84            let mut row_map = HashMap::new();
85            for (i, column) in row.columns().iter().enumerate() {
86                row_map.insert(column.name().to_owned(), extract_value(row, i));
87            }
88            result_rows.push(row_map);
89        }
90
91        Ok(QueryResult {
92            columns,
93            rows: result_rows,
94            row_count: total_rows,
95            execution_time_ms: execution_time,
96        })
97    }
98}
99
100fn extract_value(row: &sqlx::postgres::PgRow, column_index: usize) -> serde_json::Value {
101    if let Ok(val) = row.try_get::<Option<chrono::DateTime<chrono::Utc>>, _>(column_index) {
102        return val.map_or(serde_json::Value::Null, |dt| {
103            serde_json::Value::String(dt.to_rfc3339())
104        });
105    }
106    if let Ok(val) = row.try_get::<Option<String>, _>(column_index) {
107        return val.map_or(serde_json::Value::Null, serde_json::Value::String);
108    }
109    if let Ok(val) = row.try_get::<Option<i64>, _>(column_index) {
110        return val.map_or(serde_json::Value::Null, |i| {
111            serde_json::Value::Number(i.into())
112        });
113    }
114    if let Ok(val) = row.try_get::<Option<i32>, _>(column_index) {
115        return val.map_or(serde_json::Value::Null, |i| {
116            serde_json::Value::Number(i.into())
117        });
118    }
119    if let Ok(val) = row.try_get::<Option<f64>, _>(column_index) {
120        return val.map_or(serde_json::Value::Null, |f| {
121            serde_json::Number::from_f64(f)
122                .map_or(serde_json::Value::Null, serde_json::Value::Number)
123        });
124    }
125    if let Ok(val) = row.try_get::<Option<bool>, _>(column_index) {
126        return val.map_or(serde_json::Value::Null, serde_json::Value::Bool);
127    }
128    if let Ok(val) = row.try_get::<Option<Vec<String>>, _>(column_index) {
129        return val.map_or(serde_json::Value::Null, |arr| {
130            serde_json::Value::Array(arr.into_iter().map(serde_json::Value::String).collect())
131        });
132    }
133    if let Ok(val) = row.try_get::<Option<serde_json::Value>, _>(column_index) {
134        return val.unwrap_or(serde_json::Value::Null);
135    }
136    serde_json::Value::Null
137}
systemprompt_database/admin/query_executor.rs

systemprompt_database/admin/
query_executor.rs
