objectiveai_cli/command/db/
query.rs

1//! `db query` — pre-flight validate, then run under the request's
2//! optional timeout (threaded to postgres as `statement_timeout`
3//! when set; uncapped when omitted).
4//!
5//! Pre-flight rejects shapes the response variant can't represent:
6//! - multi-statement input (any unquoted `;` followed by non-empty
7//!   non-comment content),
8//! - `COPY ... TO STDOUT|STDIN` (no FE/BE protocol path for it),
9//! - transaction control verbs (`BEGIN`, `START`, `COMMIT`, `END`,
10//!   `ROLLBACK`, `SAVEPOINT`, `RELEASE`) — the handler runs the
11//!   user query inside its own read-only tx and these would
12//!   collide.
13
14use std::time::Duration;
15
16use objectiveai_sdk::cli::command::db::query::{Column, Request, Response};
17
18use crate::context::Context;
19use crate::db::query::{Column as RawColumn, RawQueryResult};
20use crate::error::Error;
21
22pub async fn execute(ctx: &Context, request: Request) -> Result<Response, Error> {
23    pre_flight_validate(&request.query)?;
24
25    let timeout = request.base.timeout_seconds.map(Duration::from_secs);
26    let raw = crate::db::query::run_readonly_query(ctx.db_client().await?, &request.query, timeout).await?;
27
28    let RawQueryResult { command_tag, columns, rows } = raw;
29    Ok(Response {
30        command_tag,
31        columns: columns
32            .into_iter()
33            .map(|RawColumn { name, r#type }| Column { name, r#type })
34            .collect(),
35        rows,
36        truncated: false,
37    })
38}
39
40/// Cheap leading-token / non-quoted-semicolon scan. We don't
41/// parse SQL — we just reject shapes the response variant can't
42/// represent. Anything else is forwarded to Postgres, which will
43/// reject writes via the read-only transaction.
44fn pre_flight_validate(sql: &str) -> Result<(), Error> {
45    let trimmed = sql.trim();
46    if trimmed.is_empty() {
47        return Err(Error::InvalidQuery("empty query".to_string()));
48    }
49
50    if has_trailing_statement(sql) {
51        return Err(Error::InvalidQuery(
52            "only one statement per call".to_string(),
53        ));
54    }
55
56    let leading = leading_keyword(trimmed);
57    match leading.as_str() {
58        "COPY" => {
59            // COPY ... TO STDOUT / FROM STDIN both need the
60            // FE/BE COPY protocol which we don't expose.
61            if has_copy_stdin_stdout(trimmed) {
62                return Err(Error::InvalidQuery(
63                    "COPY ... TO STDOUT / FROM STDIN is not supported".to_string(),
64                ));
65            }
66        }
67        "BEGIN" | "START" | "COMMIT" | "END" | "ROLLBACK" | "SAVEPOINT" | "RELEASE" => {
68            return Err(Error::InvalidQuery(
69                "transaction control is not permitted".to_string(),
70            ));
71        }
72        _ => {}
73    }
74
75    Ok(())
76}
77
78fn leading_keyword(sql: &str) -> String {
79    sql.split(|c: char| !c.is_ascii_alphabetic())
80        .next()
81        .unwrap_or("")
82        .to_ascii_uppercase()
83}
84
85/// `true` if there's a non-quoted `;` followed by any non-
86/// whitespace, non-comment content. A single trailing `;` after
87/// the last statement is fine.
88fn has_trailing_statement(sql: &str) -> bool {
89    let mut in_single = false;
90    let mut in_double = false;
91    let mut in_line_comment = false;
92    let mut in_block_comment = false;
93    let mut last_semi: Option<usize> = None;
94    let bytes = sql.as_bytes();
95    let mut i = 0;
96    while i < bytes.len() {
97        let c = bytes[i] as char;
98        if in_line_comment {
99            if c == '\n' {
100                in_line_comment = false;
101            }
102            i += 1;
103            continue;
104        }
105        if in_block_comment {
106            if c == '*' && i + 1 < bytes.len() && bytes[i + 1] as char == '/' {
107                in_block_comment = false;
108                i += 2;
109                continue;
110            }
111            i += 1;
112            continue;
113        }
114        if in_single {
115            if c == '\'' {
116                in_single = false;
117            }
118            i += 1;
119            continue;
120        }
121        if in_double {
122            if c == '"' {
123                in_double = false;
124            }
125            i += 1;
126            continue;
127        }
128        match c {
129            '\'' => in_single = true,
130            '"' => in_double = true,
131            '-' if i + 1 < bytes.len() && bytes[i + 1] as char == '-' => {
132                in_line_comment = true;
133                i += 2;
134                continue;
135            }
136            '/' if i + 1 < bytes.len() && bytes[i + 1] as char == '*' => {
137                in_block_comment = true;
138                i += 2;
139                continue;
140            }
141            ';' => last_semi = Some(i),
142            _ => {}
143        }
144        i += 1;
145    }
146    let Some(idx) = last_semi else { return false };
147    sql[idx + 1..].chars().any(|c| !c.is_whitespace())
148}
149
150fn has_copy_stdin_stdout(sql: &str) -> bool {
151    let upper = sql.to_ascii_uppercase();
152    upper.contains("STDIN") || upper.contains("STDOUT")
153}
154
155pub mod request_schema {
156    use objectiveai_sdk::cli::command::db::query as sdk;
157    use objectiveai_sdk::cli::command::db::query::request_schema::{Request, Response};
158
159    use crate::context::Context;
160    use crate::error::Error;
161
162    pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
163        Ok(objectiveai_sdk::cli::command::ResponseSchema(
164            schemars::schema_for!(sdk::Request),
165        ))
166    }
167}
168
169pub mod response_schema {
170    use objectiveai_sdk::cli::command::db::query as sdk;
171    use objectiveai_sdk::cli::command::db::query::response_schema::{Request, Response};
172
173    use crate::context::Context;
174    use crate::error::Error;
175
176    pub async fn execute(_ctx: &Context, _request: Request) -> Result<Response, Error> {
177        Ok(objectiveai_sdk::cli::command::ResponseSchema(
178            schemars::schema_for!(sdk::Response),
179        ))
180    }
181}
objectiveai_cli/command/db/query.rs

objectiveai_cli/command/db/
query.rs
