objectiveai_cli/db/
query.rs

1//! `objectiveai db query` — read-only SQL executor.
2//!
3//! Wraps the user's single statement in a short-lived read-only
4//! transaction. When a timeout is provided, server-side
5//! `statement_timeout` and `lock_timeout` are set to its budget —
6//! enforcement is postgres's alone; there is no client-side timer.
7//! When the server cancels (SQLSTATE 57014) the error maps to
8//! [`Error::QueryTimeout`] and the transaction unwinds cleanly.
9//! Without a timeout the query runs uncapped.
10//!
11//! Each result cell is decoded to a `serde_json::Value` by
12//! [`pg_value_to_json`], which dispatches on
13//! `pg_type.typname`. Common Postgres types (text family,
14//! integer family, float family, bool, uuid, timestamps, dates,
15//! json/jsonb, bytea, inet/cidr, numeric) get a typed decode;
16//! arrays of those common element types recurse; anything else
17//! falls back to the text representation if it's available, or
18//! a `<unsupported $TYPE>` placeholder.
19
20use crate::error::Error;
21use sqlx::{
22    Column as _, Row as _, TypeInfo as _, ValueRef as _,
23    postgres::{PgRow, PgValueRef},
24};
25use std::time::Duration;
26
27pub struct RawQueryResult {
28    pub command_tag: String,
29    pub columns: Vec<Column>,
30    pub rows: Vec<Vec<serde_json::Value>>,
31}
32
33#[derive(Debug, serde::Serialize)]
34pub struct Column {
35    pub name: String,
36    pub r#type: String,
37}
38
39pub async fn run_readonly_query(
40    pool: &crate::db::Pool,
41    sql: &str,
42    timeout: Option<Duration>,
43) -> Result<RawQueryResult, Error> {
44    let mut tx = pool.begin().await.map_err(crate::db::Error::Sqlx)?;
45
46    sqlx::query("SET LOCAL TRANSACTION READ ONLY")
47        .execute(&mut *tx)
48        .await
49        .map_err(crate::db::Error::Sqlx)?;
50    if let Some(timeout) = timeout {
51        let server_ms = timeout.as_millis().max(1) as u64;
52        sqlx::query(&format!("SET LOCAL statement_timeout = {server_ms}"))
53            .execute(&mut *tx)
54            .await
55            .map_err(crate::db::Error::Sqlx)?;
56        sqlx::query(&format!("SET LOCAL lock_timeout = {server_ms}"))
57            .execute(&mut *tx)
58            .await
59            .map_err(crate::db::Error::Sqlx)?;
60    }
61
62    let rows: Vec<PgRow> = sqlx::query(sql)
63        .fetch_all(&mut *tx)
64        .await
65        .map_err(map_query_err)?;
66
67    let columns = rows
68        .first()
69        .map(|r| {
70            r.columns()
71                .iter()
72                .map(|c| Column {
73                    name: c.name().to_string(),
74                    r#type: c.type_info().name().to_string(),
75                })
76                .collect()
77        })
78        .unwrap_or_default();
79
80    let mut decoded_rows: Vec<Vec<serde_json::Value>> = Vec::with_capacity(rows.len());
81    for row in &rows {
82        let mut cells: Vec<serde_json::Value> = Vec::with_capacity(row.len());
83        for i in 0..row.len() {
84            let raw = row.try_get_raw(i).map_err(crate::db::Error::Sqlx)?;
85            cells.push(pg_value_to_json(raw)?);
86        }
87        decoded_rows.push(cells);
88    }
89
90    let kw = sql
91        .trim_start()
92        .split_whitespace()
93        .next()
94        .unwrap_or("")
95        .to_uppercase();
96    let command_tag = format!("{kw} {}", rows.len());
97
98    tx.commit().await.map_err(crate::db::Error::Sqlx)?;
99    Ok(RawQueryResult {
100        command_tag,
101        columns,
102        rows: decoded_rows,
103    })
104}
105
106fn map_query_err(err: sqlx::Error) -> Error {
107    if let sqlx::Error::Database(db) = &err {
108        match db.code().as_deref() {
109            Some("25006") => return Error::QueryReadOnlyViolation,
110            Some("57014") => return Error::QueryTimeout,
111            _ => {}
112        }
113    }
114    Error::Db(crate::db::Error::Sqlx(err))
115}
116
117fn pg_value_to_json(value: PgValueRef<'_>) -> Result<serde_json::Value, Error> {
118    use serde_json::{Number, Value};
119
120    if value.is_null() {
121        return Ok(Value::Null);
122    }
123
124    let type_name = value.type_info().name().to_string();
125    match type_name.as_str() {
126        "BOOL" => Ok(Value::Bool(decode::<bool>(value)?)),
127        "INT2" => Ok(Value::Number(decode::<i16>(value)?.into())),
128        "INT4" => Ok(Value::Number(decode::<i32>(value)?.into())),
129        "INT8" => Ok(Value::Number(decode::<i64>(value)?.into())),
130        "OID" => {
131            let v: sqlx::postgres::types::Oid = decode(value)?;
132            Ok(Value::Number(v.0.into()))
133        }
134        "FLOAT4" => {
135            let v: f32 = decode(value)?;
136            Ok(Number::from_f64(v as f64)
137                .map(Value::Number)
138                .unwrap_or_else(|| Value::String(v.to_string())))
139        }
140        "FLOAT8" => {
141            let v: f64 = decode(value)?;
142            Ok(Number::from_f64(v)
143                .map(Value::Number)
144                .unwrap_or_else(|| Value::String(v.to_string())))
145        }
146        "NUMERIC" => {
147            let s = value.as_str().map_err(|e| {
148                Error::Db(crate::db::Error::InvalidData(format!(
149                    "numeric decode (text): {e}"
150                )))
151            })?;
152            Ok(Value::String(s.to_string()))
153        }
154        "TEXT" | "VARCHAR" | "BPCHAR" | "CHAR" | "NAME" | "UNKNOWN" => {
155            Ok(Value::String(decode::<String>(value)?))
156        }
157        "UUID" => {
158            let v: sqlx::types::Uuid = decode(value)?;
159            Ok(Value::String(v.hyphenated().to_string()))
160        }
161        "DATE" => {
162            let v: chrono::NaiveDate = decode(value)?;
163            Ok(Value::String(v.to_string()))
164        }
165        "TIME" => {
166            let v: chrono::NaiveTime = decode(value)?;
167            Ok(Value::String(v.to_string()))
168        }
169        "TIMESTAMP" => {
170            let v: chrono::NaiveDateTime = decode(value)?;
171            Ok(Value::String(v.format("%Y-%m-%dT%H:%M:%S%.f").to_string()))
172        }
173        "TIMESTAMPTZ" => {
174            let v: chrono::DateTime<chrono::Utc> = decode(value)?;
175            Ok(Value::String(v.to_rfc3339()))
176        }
177        "BYTEA" => {
178            use base64::Engine;
179            let v: Vec<u8> = decode(value)?;
180            Ok(Value::String(
181                base64::engine::general_purpose::STANDARD.encode(&v),
182            ))
183        }
184        "JSON" | "JSONB" => {
185            let v: sqlx::types::Json<serde_json::Value> = decode(value)?;
186            Ok(v.0)
187        }
188        "INET" | "CIDR" => {
189            let v: sqlx::types::ipnetwork::IpNetwork = decode(value)?;
190            Ok(Value::String(v.to_string()))
191        }
192        "TEXT[]" | "VARCHAR[]" | "BPCHAR[]" | "NAME[]" => {
193            let v: Vec<String> = decode(value)?;
194            Ok(Value::Array(v.into_iter().map(Value::String).collect()))
195        }
196        "BOOL[]" => {
197            let v: Vec<bool> = decode(value)?;
198            Ok(Value::Array(v.into_iter().map(Value::Bool).collect()))
199        }
200        "INT2[]" => {
201            let v: Vec<i16> = decode(value)?;
202            Ok(Value::Array(
203                v.into_iter().map(|n| Value::Number(n.into())).collect(),
204            ))
205        }
206        "INT4[]" => {
207            let v: Vec<i32> = decode(value)?;
208            Ok(Value::Array(
209                v.into_iter().map(|n| Value::Number(n.into())).collect(),
210            ))
211        }
212        "INT8[]" => {
213            let v: Vec<i64> = decode(value)?;
214            Ok(Value::Array(
215                v.into_iter().map(|n| Value::Number(n.into())).collect(),
216            ))
217        }
218        "FLOAT4[]" => {
219            let v: Vec<f32> = decode(value)?;
220            Ok(Value::Array(
221                v.into_iter()
222                    .map(|n| {
223                        Number::from_f64(n as f64)
224                            .map(Value::Number)
225                            .unwrap_or_else(|| Value::String(n.to_string()))
226                    })
227                    .collect(),
228            ))
229        }
230        "FLOAT8[]" => {
231            let v: Vec<f64> = decode(value)?;
232            Ok(Value::Array(
233                v.into_iter()
234                    .map(|n| {
235                        Number::from_f64(n)
236                            .map(Value::Number)
237                            .unwrap_or_else(|| Value::String(n.to_string()))
238                    })
239                    .collect(),
240            ))
241        }
242        "UUID[]" => {
243            let v: Vec<sqlx::types::Uuid> = decode(value)?;
244            Ok(Value::Array(
245                v.into_iter()
246                    .map(|u| Value::String(u.hyphenated().to_string()))
247                    .collect(),
248            ))
249        }
250        _ => decode_text_fallback(value),
251    }
252}
253
254fn decode<'a, T>(value: PgValueRef<'a>) -> Result<T, Error>
255where
256    T: sqlx::Decode<'a, sqlx::Postgres>,
257{
258    T::decode(value).map_err(|e| {
259        Error::Db(crate::db::Error::InvalidData(format!("decode: {e}")))
260    })
261}
262
263fn decode_text_fallback(value: PgValueRef<'_>) -> Result<serde_json::Value, Error> {
264    let type_name = value.type_info().name().to_string();
265    match value.as_str() {
266        Ok(s) => Ok(serde_json::Value::String(s.to_string())),
267        Err(_) => Ok(serde_json::Value::String(format!(
268            "<unsupported {type_name}>"
269        ))),
270    }
271}
objectiveai_cli/db/query.rs

objectiveai_cli/db/
query.rs
