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

systemprompt_database/admin/
query_executor.rs
