sbol_db_postgres/repo/
sql_console.rs

1//! The Postgres SQL console: arbitrary read-write SQL for the lab bench.
2//!
3//! There's no read-only role and no statement-shape allowlist:
4//! authorization is the operator's responsibility (network ACL, reverse
5//! proxy, whatever fronts the lab). What this enforces per request:
6//!
7//! - A `statement_timeout` (the caller bounds it).
8//! - A `lock_timeout` so a slow `LOCK TABLE` can't pin the backend.
9//! - An `idle_in_transaction_session_timeout` so a misbehaving client
10//!   can't hold the txn open after the query returns.
11//! - A server-side row cap so a `SELECT *` doesn't OOM the process.
12//! - Client-disconnect cancellation via [`CancelGuard`], so a long-running
13//!   query stops eating a connection as soon as the browser tab closes.
14//!
15//! Validation parses with libpg_query (the real PostgreSQL parser) without
16//! executing, so the error message and position match what execution would
17//! report.
18
19use std::sync::atomic::{AtomicBool, Ordering};
20use std::sync::Arc;
21use std::time::Instant;
22
23use async_trait::async_trait;
24use sbol_db_core::DomainError;
25use sbol_db_storage::{
26    SqlConsole, SqlConsoleColumn, SqlExecuteRequest, SqlExecuteResult, SqlValidateError,
27};
28use serde_json::Value;
29use sqlx::postgres::{PgRow, PgValueRef};
30use sqlx::{Column, Row, TypeInfo, ValueRef};
31use tokio::runtime::Handle;
32
33use crate::PgPool;
34
35/// The lab SQL console over a Postgres pool.
36#[derive(Clone)]
37pub struct PgSqlConsole {
38    pool: PgPool,
39}
40
41impl PgSqlConsole {
42    pub fn new(pool: PgPool) -> Self {
43        Self { pool }
44    }
45}
46
47#[async_trait]
48impl SqlConsole for PgSqlConsole {
49    async fn execute(&self, req: SqlExecuteRequest) -> Result<SqlExecuteResult, DomainError> {
50        if req.query.trim().is_empty() {
51            return Err(DomainError::InvalidInput("empty query".into()));
52        }
53
54        // One transaction holds the captured pid, the SET LOCAL GUCs, and the
55        // query. The GUCs revert on commit; the pid is captured up front so a
56        // CancelGuard can target it if the client disconnects mid-query.
57        let mut tx = self.pool.begin().await.map_err(infra)?;
58
59        let (pid,): (i32,) = sqlx::query_as::<sqlx::Postgres, (i32,)>("SELECT pg_backend_pid()")
60            .fetch_one(&mut *tx)
61            .await
62            .map_err(infra)?;
63        let guard = CancelGuard::install(self.pool.clone(), pid);
64
65        sqlx::query::<sqlx::Postgres>(&format!("SET LOCAL statement_timeout = {}", req.timeout_ms))
66            .execute(&mut *tx)
67            .await
68            .map_err(infra)?;
69        sqlx::query::<sqlx::Postgres>("SET LOCAL lock_timeout = '5s'")
70            .execute(&mut *tx)
71            .await
72            .map_err(infra)?;
73        sqlx::query::<sqlx::Postgres>("SET LOCAL idle_in_transaction_session_timeout = '10s'")
74            .execute(&mut *tx)
75            .await
76            .map_err(infra)?;
77
78        let started = Instant::now();
79        let rows = sqlx::query::<sqlx::Postgres>(&req.query)
80            .fetch_all(&mut *tx)
81            .await
82            .map_err(query_err)?;
83        tx.commit().await.map_err(infra)?;
84        guard.disarm();
85
86        let elapsed_ms = started.elapsed().as_millis() as u64;
87        let (columns, rows, truncated, row_count) = rows_to_json(&rows, req.row_limit);
88
89        Ok(SqlExecuteResult {
90            columns,
91            rows,
92            row_count,
93            truncated,
94            elapsed_ms,
95            backend_pid: Some(pid),
96        })
97    }
98
99    async fn validate(&self, query: &str) -> Result<Option<SqlValidateError>, DomainError> {
100        if query.trim().is_empty() {
101            return Ok(None);
102        }
103        match pg_query::parse(query) {
104            Ok(_) => Ok(None),
105            Err(err) => Ok(Some(pg_query_to_validate(query, &err))),
106        }
107    }
108
109    fn dialect(&self) -> &'static str {
110        "pgsql"
111    }
112}
113
114fn pg_query_to_validate(source: &str, err: &pg_query::Error) -> SqlValidateError {
115    let message = err.to_string();
116    // pg_query::Error::Parse carries the libpg_query cursor position as a
117    // 1-indexed byte offset; other variants don't pin a location.
118    let cursor = match err {
119        pg_query::Error::Parse(_) => extract_cursor_pos(&message),
120        _ => None,
121    };
122    let (line, column) = match cursor {
123        Some(pos) if pos > 0 => offset_to_line_col(source, (pos - 1) as usize),
124        _ => (1, 1),
125    };
126    SqlValidateError {
127        message,
128        line,
129        column,
130        end_line: None,
131        end_column: None,
132    }
133}
134
135/// pg_query's `Error::Parse` Display format is `"<message> at <position>"`
136/// where `<position>` is the 1-indexed byte offset of the error. Pull it back
137/// out so we can translate to line/column. `None` for variants without one.
138fn extract_cursor_pos(msg: &str) -> Option<u32> {
139    let needle = " at ";
140    let idx = msg.rfind(needle)?;
141    msg[idx + needle.len()..].trim().parse::<u32>().ok()
142}
143
144/// Translate a byte offset into a 1-indexed (line, column) pair against the
145/// original query string. Counts characters past the last newline, not bytes,
146/// so multi-byte codepoints don't push columns out of register.
147fn offset_to_line_col(source: &str, byte_offset: usize) -> (u32, u32) {
148    let mut line: u32 = 1;
149    let mut col: u32 = 1;
150    let mut consumed_bytes = 0usize;
151    for ch in source.chars() {
152        if consumed_bytes >= byte_offset {
153            break;
154        }
155        if ch == '\n' {
156            line += 1;
157            col = 1;
158        } else {
159            col += 1;
160        }
161        consumed_bytes += ch.len_utf8();
162    }
163    (line, col)
164}
165
166/// Map a sqlx error from the query itself, distinguishing user-fault (bad
167/// syntax, constraint violation, statement timeout) from server-fault (pool
168/// exhaustion, lost connection). User-fault becomes `InvalidInput` so the
169/// server answers 400 without 5xx noise.
170fn query_err(err: sqlx::Error) -> DomainError {
171    if let Some(db) = err.as_database_error() {
172        let code = db.code().unwrap_or_default();
173        // Postgres SQLSTATE classes a user can cause:
174        //   22*: data exception   23*: integrity constraint
175        //   42*: syntax / access rule   57014: query canceled (timeout)
176        let class = code.chars().take(2).collect::<String>();
177        let is_user_fault = matches!(class.as_str(), "22" | "23" | "42") || code == "57014";
178        if is_user_fault {
179            return DomainError::InvalidInput(db.message().to_string());
180        }
181    }
182    infra(err)
183}
184
185fn infra(e: sqlx::Error) -> DomainError {
186    DomainError::Database(e.to_string())
187}
188
189/// Reshape `PgRow`s into the console wire shape. `limit` caps `rows`
190/// server-side; `truncated` reports whether any were dropped; `total` is the
191/// count we'd have returned without truncation.
192fn rows_to_json(rows: &[PgRow], limit: u32) -> (Vec<SqlConsoleColumn>, Vec<Vec<Value>>, bool, u64) {
193    let total = rows.len() as u64;
194    let truncated = total > limit as u64;
195    let take = rows.iter().take(limit as usize);
196
197    let columns: Vec<SqlConsoleColumn> = rows
198        .first()
199        .map(|row| {
200            row.columns()
201                .iter()
202                .map(|c| SqlConsoleColumn {
203                    name: c.name().to_string(),
204                    column_type: c.type_info().name().to_string(),
205                })
206                .collect()
207        })
208        .unwrap_or_default();
209
210    let json_rows = take
211        .map(|row| {
212            row.columns()
213                .iter()
214                .enumerate()
215                .map(|(i, col)| value_to_json(row, i, col.type_info().name()))
216                .collect()
217        })
218        .collect();
219
220    (columns, json_rows, truncated, total)
221}
222
223fn value_to_json(row: &PgRow, idx: usize, type_name: &str) -> Value {
224    let raw: PgValueRef<'_> = match row.try_get_raw(idx) {
225        Ok(v) => v,
226        Err(_) => return Value::Null,
227    };
228    if raw.is_null() {
229        return Value::Null;
230    }
231
232    // Dispatch on the Postgres type name (stable across protocol revisions
233    // and matching what we ship in the column metadata).
234    match type_name {
235        "BOOL" => decode::<bool>(row, idx)
236            .map(Value::Bool)
237            .unwrap_or(Value::Null),
238        "INT2" => decode::<i16>(row, idx)
239            .map(|n| Value::from(n as i64))
240            .unwrap_or(Value::Null),
241        "INT4" => decode::<i32>(row, idx)
242            .map(|n| Value::from(n as i64))
243            .unwrap_or(Value::Null),
244        "INT8" => decode::<i64>(row, idx)
245            .map(Value::from)
246            .unwrap_or(Value::Null),
247        "FLOAT4" => decode::<f32>(row, idx)
248            .and_then(|f| serde_json::Number::from_f64(f as f64).map(Value::Number))
249            .unwrap_or(Value::Null),
250        "FLOAT8" => decode::<f64>(row, idx)
251            .and_then(|f| serde_json::Number::from_f64(f).map(Value::Number))
252            .unwrap_or(Value::Null),
253        // sqlx feature-gates bigdecimal/rust_decimal; we depend on neither,
254        // so render NUMERIC as text.
255        "NUMERIC" => decode_via_text(row, idx),
256        "TEXT" | "VARCHAR" | "BPCHAR" | "NAME" | "CITEXT" => decode::<String>(row, idx)
257            .map(Value::String)
258            .unwrap_or(Value::Null),
259        "UUID" => decode::<uuid::Uuid>(row, idx)
260            .map(|u| Value::String(u.to_string()))
261            .unwrap_or(Value::Null),
262        "TIMESTAMPTZ" => decode::<chrono::DateTime<chrono::Utc>>(row, idx)
263            .map(|t| Value::String(t.to_rfc3339()))
264            .unwrap_or(Value::Null),
265        "TIMESTAMP" => decode::<chrono::NaiveDateTime>(row, idx)
266            .map(|t| Value::String(t.to_string()))
267            .unwrap_or(Value::Null),
268        "DATE" => decode::<chrono::NaiveDate>(row, idx)
269            .map(|t| Value::String(t.to_string()))
270            .unwrap_or(Value::Null),
271        "JSON" | "JSONB" => decode::<Value>(row, idx).unwrap_or(Value::Null),
272        "BYTEA" => decode::<Vec<u8>>(row, idx)
273            .map(|b| Value::String(hex::encode(b)))
274            .unwrap_or(Value::Null),
275        "TEXT[]" | "VARCHAR[]" | "_TEXT" | "_VARCHAR" | "_BPCHAR" | "_NAME" => {
276            decode::<Vec<String>>(row, idx)
277                .map(|v| Value::Array(v.into_iter().map(Value::String).collect()))
278                .unwrap_or(Value::Null)
279        }
280        "_INT4" => decode::<Vec<i32>>(row, idx)
281            .map(|v| Value::Array(v.into_iter().map(|i| Value::from(i as i64)).collect()))
282            .unwrap_or(Value::Null),
283        "_INT8" => decode::<Vec<i64>>(row, idx)
284            .map(|v| Value::Array(v.into_iter().map(Value::from).collect()))
285            .unwrap_or(Value::Null),
286        _ => {
287            // Custom Postgres domain (sbol_iri, sbol_ontology_term, …) lands
288            // here; most are text under the covers.
289            if let Some(s) = decode::<String>(row, idx) {
290                return Value::String(s);
291            }
292            decode_via_text(row, idx)
293        }
294    }
295}
296
297fn decode<'r, T>(row: &'r PgRow, idx: usize) -> Option<T>
298where
299    T: sqlx::Decode<'r, sqlx::Postgres> + sqlx::Type<sqlx::Postgres>,
300{
301    row.try_get(idx).ok()
302}
303
304fn decode_via_text(row: &PgRow, idx: usize) -> Value {
305    if let Ok(Some(s)) = row.try_get::<Option<String>, _>(idx) {
306        return Value::String(s);
307    }
308    let type_name = row
309        .columns()
310        .get(idx)
311        .map(|c| c.type_info().name().to_string())
312        .unwrap_or_else(|| "unknown".into());
313    Value::String(format!("<undecodable: {type_name}>"))
314}
315
316/// RAII guard that cancels backend PID `pid` on `pool` unless `disarm`ed
317/// before drop. On client disconnect axum drops the request future, which
318/// drops this guard and fires `pg_cancel_backend` from a fresh connection so
319/// the server doesn't sit on a slot for the full timeout. The cancel runs on
320/// a tokio task so `Drop` stays synchronous.
321struct CancelGuard {
322    inner: Arc<CancelInner>,
323}
324
325struct CancelInner {
326    pool: PgPool,
327    pid: i32,
328    disarmed: AtomicBool,
329    handle: Handle,
330}
331
332impl CancelGuard {
333    fn install(pool: PgPool, pid: i32) -> Self {
334        Self {
335            inner: Arc::new(CancelInner {
336                pool,
337                pid,
338                disarmed: AtomicBool::new(false),
339                handle: Handle::current(),
340            }),
341        }
342    }
343
344    fn disarm(self) {
345        self.inner.disarmed.store(true, Ordering::Release);
346    }
347}
348
349impl Drop for CancelGuard {
350    fn drop(&mut self) {
351        if self.inner.disarmed.load(Ordering::Acquire) {
352            return;
353        }
354        let inner = Arc::clone(&self.inner);
355        inner.handle.clone().spawn(async move {
356            let result = sqlx::query("SELECT pg_cancel_backend($1)")
357                .bind(inner.pid)
358                .execute(&inner.pool)
359                .await;
360            match result {
361                Ok(_) => tracing::debug!(pid = inner.pid, "lab sql query cancel sent"),
362                Err(e) => {
363                    tracing::warn!(pid = inner.pid, error = %e, "lab sql query cancel failed")
364                }
365            }
366        });
367    }
368}
sbol_db_postgres/repo/sql_console.rs

sbol_db_postgres/repo/
sql_console.rs
