diff --git a/backend/windmill-common/src/worker.rs b/backend/windmill-common/src/worker.rs index c03856f8d5..ca141bfebd 100644 --- a/backend/windmill-common/src/worker.rs +++ b/backend/windmill-common/src/worker.rs @@ -765,6 +765,11 @@ pub struct SqlAnnotations { pub return_last_result: bool, // deprecated, use result_collection instead pub result_collection: SqlResultCollectionStrategy, pub prepare: bool, // Used to prepare datatable queries without executing + // Emit {columns: [{name, oid, type_name}], rows: [[text|null]]} from the + // last statement instead of the default `[{col: val}]` JSON shape. Used by + // `wmill datatable serve` to map Postgres results onto the wire protocol + // without re-stringifying every JSON value. + pub raw_output: bool, } #[annotations("#")] @@ -2328,6 +2333,20 @@ mod tests { assert!(annotations.sandbox); } + #[test] + fn test_sql_raw_output_annotation() { + let with_flag = SqlAnnotations::parse("-- raw_output\nSELECT 1"); + assert!(with_flag.raw_output); + assert!(!with_flag.prepare); + + let without_flag = SqlAnnotations::parse("SELECT 1"); + assert!(!without_flag.raw_output); + + let combined = SqlAnnotations::parse("-- prepare\n-- raw_output\nSELECT 1"); + assert!(combined.raw_output); + assert!(combined.prepare); + } + #[test] fn test_python_sandbox_annotation_with_other_annotations() { let content = "# no_cache\n# sandbox\ndef main():\n pass"; diff --git a/backend/windmill-worker/src/lib.rs b/backend/windmill-worker/src/lib.rs index 0d5df81a83..08f9381205 100644 --- a/backend/windmill-worker/src/lib.rs +++ b/backend/windmill-worker/src/lib.rs @@ -59,6 +59,7 @@ mod otel_oss; mod otel_tracing_proxy_ee; mod otel_tracing_proxy_oss; pub mod pg_executor; +mod pg_raw_output; #[cfg(feature = "php")] mod php_executor; mod prepare_deps; diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs index 16a555cc7e..a1a26d0630 100644 --- a/backend/windmill-worker/src/pg_executor.rs +++ b/backend/windmill-worker/src/pg_executor.rs @@ -337,6 +337,7 @@ fn do_postgresql_inner<'a>( job_id: Uuid, workspace_id: &'a str, log_conn: &'a Connection, + raw_output: bool, ) -> error::Result>>>> { let mut query_params = vec![]; let mut param_types: Vec = vec![]; @@ -522,26 +523,31 @@ fn do_postgresql_inner<'a>( ); } - for row in rows.into_iter() { - let r = postgres_row_to_json_value_with_state(row, &format_state); - if let Ok(v) = r.as_ref() { - let size = sizeof_val(v); - siz.fetch_add(size, Ordering::Relaxed); - } - if *CLOUD_HOSTED { - let siz = siz.load(Ordering::Relaxed); - if siz > MAX_RESULT_SIZE * 4 { - return Err(Error::ExecutionErr(format!( - "Query result too large for cloud (size = {} > {})", - siz, - MAX_RESULT_SIZE & 4, - ))); + if raw_output { + let envelope = crate::pg_raw_output::build_envelope(rows, &format_state, siz)?; + res.push(to_raw_value(&envelope)); + } else { + for row in rows.into_iter() { + let r = postgres_row_to_json_value_with_state(row, &format_state); + if let Ok(v) = r.as_ref() { + let size = sizeof_val(v); + siz.fetch_add(size, Ordering::Relaxed); + } + if *CLOUD_HOSTED { + let siz = siz.load(Ordering::Relaxed); + if siz > MAX_RESULT_SIZE * 4 { + return Err(Error::ExecutionErr(format!( + "Query result too large for cloud (size = {} > {})", + siz, + MAX_RESULT_SIZE * 4, + ))); + } + } + if let Ok(v) = r { + res.push(to_raw_value(&v)); + } else { + return Err(to_anyhow(r.err().unwrap()).into()); } - } - if let Ok(v) = r { - res.push(to_raw_value(&v)); - } else { - return Err(to_anyhow(r.err().unwrap()).into()); } } } @@ -610,7 +616,9 @@ pub async fn do_postgresql( }; let annotations = windmill_common::worker::SqlAnnotations::parse(query); - let collection_strategy = if annotations.return_last_result { + let collection_strategy = if annotations.raw_output || annotations.return_last_result { + // raw_output emits a single envelope from the last statement, so the + // multi-statement collection modes don't apply. SqlResultCollectionStrategy::LastStatementAllRows } else { annotations.result_collection @@ -784,11 +792,14 @@ pub async fn do_postgresql( continue; } + let skip_collect = collection_strategy.collect_last_statement_only(queries.len()) + && i < queries.len() - 1; + let is_last = i == queries.len() - 1; let result = do_postgresql_inner( query.to_string(), ¶m_idx_to_arg_and_value, client, - if i == queries.len() - 1 + if is_last && s3.is_none() && collection_strategy.collect_last_statement_only(queries.len()) && !collection_strategy.collect_scalar() @@ -798,19 +809,26 @@ pub async fn do_postgresql( None }, size_ref, - collection_strategy.collect_last_statement_only(queries.len()) - && i < queries.len() - 1, + skip_collect, collection_strategy.collect_first_row_only(), s3.clone(), job.id, &job.workspace_id, conn, + annotations.raw_output && is_last && !skip_collect, )? .await?; results.push(result); } - collection_strategy.collect(results) + if annotations.raw_output { + // The raw_output envelope already aggregates the last statement's + // result; skip the collection_strategy reshape that wraps rows in + // a JSON array. + Ok(crate::pg_raw_output::extract_envelope_or_empty(results)) + } else { + collection_strategy.collect(results) + } }; let result = if run_inline { diff --git a/backend/windmill-worker/src/pg_raw_output.rs b/backend/windmill-worker/src/pg_raw_output.rs new file mode 100644 index 0000000000..85f8b6d1aa --- /dev/null +++ b/backend/windmill-worker/src/pg_raw_output.rs @@ -0,0 +1,205 @@ +//! Postgres `-- raw_output` envelope encoder. +//! +//! When a postgres script carries `-- raw_output`, `pg_executor` returns a +//! single JSON object of shape `{columns: [{name, oid, type_name}], rows: +//! [[text|null]]}` from the last statement, instead of the default +//! `[{col: val}]` row array. The shape gives a Postgres-protocol client (today +//! that's `wmill datatable serve`) everything it needs to build +//! `RowDescription` + `DataRow` messages without re-stringifying every JSON +//! value. +//! +//! All extra types and the per-row text conversion live here so +//! `pg_executor.rs` only has to dispatch on the annotation. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +use serde::Serialize; +use serde_json::value::RawValue; +use tokio_postgres::Row; +use windmill_common::error::{self, to_anyhow, Error}; +use windmill_common::worker::CLOUD_HOSTED; + +use crate::pg_executor::{postgres_row_to_row_data_with_state, JSONValue, ResultFormatState}; +use crate::MAX_RESULT_SIZE; + +#[derive(Serialize, Clone, Debug)] +pub struct RawOutputColumn { + pub name: String, + pub oid: u32, + pub type_name: String, +} + +#[derive(Serialize, Debug)] +pub struct RawOutputEnvelope { + pub columns: Vec, + pub rows: Vec>>, +} + +impl RawOutputEnvelope { + /// Empty envelope returned when a raw_output query had no statements that + /// produced rows (e.g. only `SET`). Surfaces to the client as + /// `CommandComplete "SELECT 0"`. + pub fn empty_raw_value() -> Box { + RawValue::from_string("{\"columns\":[],\"rows\":[]}".to_string()).unwrap() + } +} + +/// Build a `RawOutputEnvelope` for one statement's worth of rows. Reuses the +/// existing JSON-cell formatter (numeric precision warning, interval/timetz +/// coercion, JSON columns…) and only post-processes its output into the text +/// form Postgres's wire protocol expects. +pub fn build_envelope( + rows: Vec, + format_state: &ResultFormatState, + siz: &AtomicUsize, +) -> error::Result { + let columns: Vec = rows + .first() + .map(|r| { + r.columns() + .iter() + .map(|c| RawOutputColumn { + name: c.name().to_string(), + oid: c.type_().oid(), + type_name: c.type_().name().to_string(), + }) + .collect() + }) + .unwrap_or_default(); + + let mut text_rows: Vec>> = Vec::with_capacity(rows.len()); + for row in rows { + let row_data = postgres_row_to_row_data_with_state(row, format_state).map_err(to_anyhow)?; + let text_row: Vec> = columns + .iter() + .map(|c| { + json_value_to_pg_text(row_data.get(&c.name).cloned().unwrap_or(JSONValue::Null)) + }) + .collect(); + siz.fetch_add( + text_row + .iter() + .filter_map(|v| v.as_deref().map(str::len)) + .sum(), + Ordering::Relaxed, + ); + if *CLOUD_HOSTED { + let total = siz.load(Ordering::Relaxed); + if total > MAX_RESULT_SIZE * 4 { + return Err(Error::ExecutionErr(format!( + "Query result too large for cloud (size = {} > {})", + total, + MAX_RESULT_SIZE * 4, + ))); + } + } + text_rows.push(text_row); + } + + Ok(RawOutputEnvelope { columns, rows: text_rows }) +} + +/// Extract the single envelope `Box` from the per-statement results +/// produced by `do_postgresql`, falling back to an empty envelope if no +/// statement contributed one (multi-statement query whose last item didn't +/// return rows). +pub fn extract_envelope_or_empty(results: Vec>>) -> Box { + results + .into_iter() + .last() + .and_then(|v| v.into_iter().last()) + .unwrap_or_else(RawOutputEnvelope::empty_raw_value) +} + +/// Convert a JSON value (as produced by `postgres_row_to_row_data_with_state`) +/// to the text shape Postgres's wire protocol expects in `DataRow`. +/// - SQL NULL → None (encoded as -1 length on the wire) +/// - JSON strings already carry the server's text repr +/// - Numbers/bools render with Postgres conventions (`t`/`f`) +/// - Arrays/objects fall back to JSON serialization (psql will display them +/// as JSON literals; not a valid pg array literal, but readable) +fn json_value_to_pg_text(v: JSONValue) -> Option { + match v { + JSONValue::Null => None, + JSONValue::String(s) => Some(s), + JSONValue::Number(n) => Some(n.to_string()), + JSONValue::Bool(true) => Some("t".to_string()), + JSONValue::Bool(false) => Some("f".to_string()), + JSONValue::Array(_) | JSONValue::Object(_) => Some(v.to_string()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn json_value_to_pg_text_covers_every_kind() { + // NULLs survive as None so the wire layer emits length -1. + assert_eq!(json_value_to_pg_text(JSONValue::Null), None); + // Strings already carry the server's text repr (e.g. "2024-01-15"); + // we deliberately strip the JSON quotes. + assert_eq!( + json_value_to_pg_text(JSONValue::String("hello".into())), + Some("hello".to_string()), + ); + // Numbers render without JSON wrapping. + assert_eq!( + json_value_to_pg_text(JSONValue::Number(serde_json::Number::from(42))), + Some("42".to_string()), + ); + // Booleans use Postgres's `t`/`f` text convention. + assert_eq!( + json_value_to_pg_text(JSONValue::Bool(true)), + Some("t".to_string()) + ); + assert_eq!( + json_value_to_pg_text(JSONValue::Bool(false)), + Some("f".to_string()) + ); + // Arrays/objects fall back to JSON — psql will display them verbatim. + assert_eq!( + json_value_to_pg_text(serde_json::json!([1, 2, 3])), + Some("[1,2,3]".to_string()), + ); + assert_eq!( + json_value_to_pg_text(serde_json::json!({"k": "v"})), + Some("{\"k\":\"v\"}".to_string()), + ); + } + + #[test] + fn envelope_serializes_to_the_shape_the_cli_consumes() { + // Lock in the JSON shape that the CLI's pg-wire bridge consumes. + let envelope = RawOutputEnvelope { + columns: vec![ + RawOutputColumn { name: "id".into(), oid: 23, type_name: "int4".into() }, + RawOutputColumn { name: "name".into(), oid: 25, type_name: "text".into() }, + ], + rows: vec![ + vec![Some("1".into()), Some("alice".into())], + vec![Some("2".into()), None], + ], + }; + let json = serde_json::to_value(&envelope).unwrap(); + assert_eq!( + json, + serde_json::json!({ + "columns": [ + {"name": "id", "oid": 23, "type_name": "int4"}, + {"name": "name", "oid": 25, "type_name": "text"}, + ], + "rows": [["1", "alice"], ["2", null]], + }), + ); + } + + #[test] + fn extract_envelope_falls_back_to_empty_when_no_statement_produced_one() { + let raw = extract_envelope_or_empty(vec![]); + assert_eq!(raw.get(), "{\"columns\":[],\"rows\":[]}"); + + let raw = extract_envelope_or_empty(vec![vec![]]); + assert_eq!(raw.get(), "{\"columns\":[],\"rows\":[]}"); + } +} diff --git a/cli/bun.lock b/cli/bun.lock index de02781822..c94551beb3 100644 --- a/cli/bun.lock +++ b/cli/bun.lock @@ -16,6 +16,7 @@ "jszip": "3.8.0", "minimatch": "^10.0.0", "open": "^10.0.0", + "pg-gateway": "0.3.0-beta.4", "svelte": "^5.45.2", "tar-stream": "^3.1.7", "windmill-parser-wasm-csharp": "1.510.1", @@ -253,6 +254,8 @@ "pako": ["pako@1.0.11", "", {}, "sha512-4hLB8Py4zZce5s4yd9XzopqwVv/yGNhV1Bl8NTmCq1763HeK2+EwVTv+leGeL13Dnh2wfbqowVPXCIO0z4taYw=="], + "pg-gateway": ["pg-gateway@0.3.0-beta.4", "", {}, "sha512-CTjsM7Z+0Nx2/dyZ6r8zRsc3f9FScoD5UAOlfUx1Fdv/JOIWvRbF7gou6l6vP+uypXQVoYPgw8xZDXgMGvBa4Q=="], + "process-nextick-args": ["process-nextick-args@2.0.1", "", {}, "sha512-3ouUOpQhtgrbOa17J7+uxOTpITYWaGP7/AhoR3+A+/1e9skrzelGi/dXzEYyvbxubEF6Wn2ypscTKiKJFFn1ag=="], "readable-stream": ["readable-stream@2.3.8", "", { "dependencies": { "core-util-is": "~1.0.0", "inherits": "~2.0.3", "isarray": "~1.0.0", "process-nextick-args": "~2.0.0", "safe-buffer": "~5.1.1", "string_decoder": "~1.1.1", "util-deprecate": "~1.0.1" } }, "sha512-8p0AUk4XODgIewSi0l8Epjs+EVnWiK7NoDIEGU0HhE7+ZyY8D1IMY7odu5lRrFXGg71L15KG8QrPmum45RTtdA=="], diff --git a/cli/package.json b/cli/package.json index 5afe6f04bb..2f46e3829b 100644 --- a/cli/package.json +++ b/cli/package.json @@ -25,6 +25,7 @@ "jszip": "3.8.0", "minimatch": "^10.0.0", "open": "^10.0.0", + "pg-gateway": "0.3.0-beta.4", "svelte": "^5.45.2", "tar-stream": "^3.1.7", "windmill-parser-wasm-csharp": "1.510.1", diff --git a/cli/src/commands/datatable/datatable.ts b/cli/src/commands/datatable/datatable.ts index 36ae612b5c..c86c971a90 100644 --- a/cli/src/commands/datatable/datatable.ts +++ b/cli/src/commands/datatable/datatable.ts @@ -7,6 +7,8 @@ import { resolveWorkspace } from "../../core/context.ts"; import * as log from "../../core/log.ts"; import { GlobalOptions } from "../../types.ts"; import { runCatalogQuery } from "../../utils/catalog.ts"; +import { psql as psqlDatatable } from "./psql.ts"; +import { serve as serveDatatable } from "./serve.ts"; const DEFAULT_DATATABLE_NAME = "main"; @@ -39,6 +41,18 @@ async function run( await runCatalogQuery(opts, "datatable", name, sql); } +async function serve( + opts: GlobalOptions & { port?: number; host?: string; password?: string }, +) { + await serveDatatable(opts); +} + +async function psql( + opts: GlobalOptions & { name?: string; port?: number; host?: string; password?: string }, +) { + await psqlDatatable(opts); +} + const command = new Command() .description("datatable related commands") .command("list", "list all datatables in the workspace") @@ -54,6 +68,44 @@ const command = new Command() "-s --silent", "Output only the final result as JSON. Useful for scripting.", ) - .action(run as any); + .action(run as any) + .command( + "serve", + "Serve all datatables as a Postgres-wire endpoint (psql, DBeaver, pgAdmin); the client picks the datatable via the database name in its connection string", + ) + .option( + "--port ", + "Port to listen on (default: first free port in 5433-5500)", + ) + .option( + "--host ", + "Bind address (default: 127.0.0.1)", + ) + .option( + "--password ", + "Password for Postgres clients (default: generate a random password at startup)", + ) + .action(serve as any) + .command( + "psql", + "Start a serve listener and launch psql connected to it", + ) + .option( + "-n --name ", + "Datatable to connect psql to (default: main)", + ) + .option( + "--port ", + "Port the proxy listens on (default: first free port in 5433-5500)", + ) + .option( + "--host ", + "Bind address for the proxy (default: 127.0.0.1)", + ) + .option( + "--password ", + "Password for the temporary Postgres proxy (default: generate a random password at startup)", + ) + .action(psql as any); export default command; diff --git a/cli/src/commands/datatable/pg_wire.ts b/cli/src/commands/datatable/pg_wire.ts new file mode 100644 index 0000000000..b4d82fc8fb --- /dev/null +++ b/cli/src/commands/datatable/pg_wire.ts @@ -0,0 +1,251 @@ +// Postgres wire-protocol helpers used by `wmill datatable serve`. +// +// pg-gateway handles the startup/auth handshake for us, but it leaves +// query-response framing (RowDescription / DataRow / CommandComplete / +// ReadyForQuery) to the caller. The functions below build those four +// messages from a `RawOutputEnvelope` returned by the backend. +// +// References: +// https://www.postgresql.org/docs/current/protocol-message-formats.html + +const TEXT_OID = 25; + +export interface RawOutputColumn { + name: string; + oid: number; + type_name: string; +} + +export interface RawOutputEnvelope { + columns: RawOutputColumn[]; + rows: (string | null)[][]; +} + +export interface ExecuteResponse { + message: Uint8Array; + nextRowOffset: number; + suspended: boolean; +} + +const encoder = new TextEncoder(); + +/** + * Build a single backend message: one type byte, a 4-byte length (covering + * the length field itself but not the type byte), then the payload. + */ +function buildMessage(type: number, payload: Uint8Array): Uint8Array { + const out = new Uint8Array(1 + 4 + payload.byteLength); + out[0] = type; + const len = 4 + payload.byteLength; + out[1] = (len >>> 24) & 0xff; + out[2] = (len >>> 16) & 0xff; + out[3] = (len >>> 8) & 0xff; + out[4] = len & 0xff; + out.set(payload, 5); + return out; +} + +function writeInt16(buf: Uint8Array, offset: number, v: number): void { + buf[offset] = (v >>> 8) & 0xff; + buf[offset + 1] = v & 0xff; +} + +function writeInt32(buf: Uint8Array, offset: number, v: number): void { + buf[offset] = (v >>> 24) & 0xff; + buf[offset + 1] = (v >>> 16) & 0xff; + buf[offset + 2] = (v >>> 8) & 0xff; + buf[offset + 3] = v & 0xff; +} + +/** + * RowDescription ('T'): for each column, emit + * name CSTRING, tableOid INT32, attnum INT16, + * typeOid INT32, typeSize INT16, typeMod INT32, formatCode INT16 + * + * We don't know the underlying table/attnum (the value travels through JSON, + * not directly from the prepared statement), so we send 0 — psql is fine with + * that. Format code 0 means "text" — we'll send DataRow values as UTF-8 bytes. + */ +export function buildRowDescription(columns: RawOutputColumn[]): Uint8Array { + const fields = columns.map((col) => { + const nameBytes = encoder.encode(col.name); + const buf = new Uint8Array(nameBytes.byteLength + 1 + 4 + 2 + 4 + 2 + 4 + 2); + let off = 0; + buf.set(nameBytes, off); + off += nameBytes.byteLength; + buf[off++] = 0; // CSTRING terminator + writeInt32(buf, off, 0); off += 4; // tableOid + writeInt16(buf, off, 0); off += 2; // attnum + writeInt32(buf, off, col.oid || TEXT_OID); off += 4; // typeOid (fall back to text) + writeInt16(buf, off, -1); off += 2; // typeSize (-1 = variable) + writeInt32(buf, off, -1); off += 4; // atttypmod + writeInt16(buf, off, 0); off += 2; // format = text + return buf; + }); + const totalLen = 2 + fields.reduce((s, f) => s + f.byteLength, 0); + const payload = new Uint8Array(totalLen); + writeInt16(payload, 0, columns.length); + let off = 2; + for (const f of fields) { + payload.set(f, off); + off += f.byteLength; + } + return buildMessage(0x54 /* 'T' */, payload); +} + +/** + * DataRow ('D'): INT16 column count, then for each column either + * INT32 -1 (NULL) + * or + * INT32 length + UTF-8 bytes + */ +export function buildDataRow(row: (string | null)[]): Uint8Array { + const encoded = row.map((v) => (v === null ? null : encoder.encode(v))); + const total = + 2 + encoded.reduce((s, e) => s + 4 + (e === null ? 0 : e.byteLength), 0); + const payload = new Uint8Array(total); + writeInt16(payload, 0, row.length); + let off = 2; + for (const e of encoded) { + if (e === null) { + writeInt32(payload, off, -1); + off += 4; + } else { + writeInt32(payload, off, e.byteLength); + off += 4; + payload.set(e, off); + off += e.byteLength; + } + } + return buildMessage(0x44 /* 'D' */, payload); +} + +/** + * CommandComplete ('C'): null-terminated tag like "SELECT 12". + * We always send SELECT-tagged completions — psql doesn't validate the tag + * against the actual statement kind, so SET/CREATE/etc. just show as "SELECT 0". + */ +export function buildCommandComplete(tag: string): Uint8Array { + const tagBytes = encoder.encode(tag); + const payload = new Uint8Array(tagBytes.byteLength + 1); + payload.set(tagBytes, 0); + payload[tagBytes.byteLength] = 0; + return buildMessage(0x43 /* 'C' */, payload); +} + +/** ReadyForQuery ('Z') with the transaction-status indicator byte. */ +export function buildReadyForQuery(status: "I" | "T" | "E" = "I"): Uint8Array { + const payload = new Uint8Array(1); + payload[0] = status.charCodeAt(0); + return buildMessage(0x5a /* 'Z' */, payload); +} + +/** EmptyQueryResponse ('I') with no payload. */ +export function buildEmptyQueryResponse(): Uint8Array { + return buildMessage(0x49 /* 'I' */, new Uint8Array(0)); +} + +/** ParseComplete ('1') with no payload. */ +export function buildParseComplete(): Uint8Array { + return buildMessage(0x31 /* '1' */, new Uint8Array(0)); +} + +/** BindComplete ('2') with no payload. */ +export function buildBindComplete(): Uint8Array { + return buildMessage(0x32 /* '2' */, new Uint8Array(0)); +} + +/** CloseComplete ('3') with no payload. */ +export function buildCloseComplete(): Uint8Array { + return buildMessage(0x33 /* '3' */, new Uint8Array(0)); +} + +/** NoData ('n') with no payload. */ +export function buildNoData(): Uint8Array { + return buildMessage(0x6e /* 'n' */, new Uint8Array(0)); +} + +/** PortalSuspended ('s') with no payload. */ +export function buildPortalSuspended(): Uint8Array { + return buildMessage(0x73 /* 's' */, new Uint8Array(0)); +} + +/** ParameterDescription ('t'): INT16 count + one INT32 OID per parameter. */ +export function buildParameterDescription(typeOids: number[]): Uint8Array { + const payload = new Uint8Array(2 + typeOids.length * 4); + writeInt16(payload, 0, typeOids.length); + let off = 2; + for (const oid of typeOids) { + writeInt32(payload, off, oid); + off += 4; + } + return buildMessage(0x74 /* 't' */, payload); +} + +/** + * Translate a `RawOutputEnvelope` into the byte stream a Postgres client + * expects in response to a Query message: RowDescription, one DataRow per + * row, CommandComplete, ReadyForQuery. + * + * If there are no columns AND no rows, we treat the statement as a no-op + * (e.g. `SET client_encoding`) and emit only CommandComplete + ReadyForQuery, + * skipping RowDescription. psql is fine with either shape. + */ +export function buildQueryResponse(envelope: RawOutputEnvelope): Uint8Array { + const parts: Uint8Array[] = []; + if (envelope.columns.length > 0) { + parts.push(buildRowDescription(envelope.columns)); + for (const row of envelope.rows) { + parts.push(buildDataRow(row)); + } + } + parts.push(buildCommandComplete(`SELECT ${envelope.rows.length}`)); + parts.push(buildReadyForQuery("I")); + return concat(parts); +} + +export function buildExecuteResponse( + envelope: RawOutputEnvelope, + opts?: { + includeRowDescription?: boolean; + maxRows?: number; + rowOffset?: number; + }, +): ExecuteResponse { + const includeRowDescription = opts?.includeRowDescription ?? true; + const maxRows = opts?.maxRows ?? 0; + const rowOffset = opts?.rowOffset ?? 0; + const remainingRows = envelope.rows.slice(rowOffset); + const suspended = maxRows > 0 && remainingRows.length > maxRows; + const rowsToSend = suspended ? remainingRows.slice(0, maxRows) : remainingRows; + + const parts: Uint8Array[] = []; + if (includeRowDescription && envelope.columns.length > 0) { + parts.push(buildRowDescription(envelope.columns)); + } + for (const row of rowsToSend) { + parts.push(buildDataRow(row)); + } + if (suspended) { + parts.push(buildPortalSuspended()); + } else { + parts.push(buildCommandComplete(`SELECT ${envelope.rows.length}`)); + } + + return { + message: concat(parts), + nextRowOffset: rowOffset + rowsToSend.length, + suspended, + }; +} + +export function concat(parts: Uint8Array[]): Uint8Array { + const total = parts.reduce((s, p) => s + p.byteLength, 0); + const out = new Uint8Array(total); + let off = 0; + for (const p of parts) { + out.set(p, off); + off += p.byteLength; + } + return out; +} diff --git a/cli/src/commands/datatable/psql.ts b/cli/src/commands/datatable/psql.ts new file mode 100644 index 0000000000..095c6e3fed --- /dev/null +++ b/cli/src/commands/datatable/psql.ts @@ -0,0 +1,66 @@ +import { colors } from "@cliffy/ansi/colors"; +import { spawn } from "node:child_process"; + +import * as log from "../../core/log.ts"; +import { startServe, type ServeOpts } from "./serve.ts"; + +export interface PsqlOpts extends ServeOpts { + name?: string; +} + +const DEFAULT_DATATABLE_NAME = "main"; + +export async function psql(opts: PsqlOpts): Promise { + const datatableName = opts.name ?? DEFAULT_DATATABLE_NAME; + const handle = await startServe(opts); + const connectionString = handle.connectionString(datatableName); + + log.info( + colors.gray( + `Launching psql against datatable '${datatableName}' (proxy on ${handle.host}:${handle.port})`, + ), + ); + + const child = spawn("psql", [connectionString], { + stdio: "inherit", + }); + + // If psql isn't installed, surface that cleanly instead of a stack trace. + child.on("error", (err) => { + const code = (err as NodeJS.ErrnoException).code; + if (code === "ENOENT") { + log.error( + "psql not found in PATH. Install the postgresql client (e.g. `apt install postgresql-client`, `brew install libpq`) and try again.", + ); + } else { + log.error(`failed to launch psql: ${err.message}`); + } + handle.close().finally(() => process.exit(1)); + }); + + const exitCode = await new Promise((resolve) => { + child.on("exit", (code, signal) => { + // Mirror shell convention: signal N → 128 + N. + if (signal && typeof signal === "string") { + const num = signalToNumber(signal); + resolve(num !== undefined ? 128 + num : 1); + } else { + resolve(code ?? 0); + } + }); + }); + + await handle.close(); + process.exit(exitCode); +} + +function signalToNumber(name: string): number | undefined { + const table: Record = { + SIGINT: 2, + SIGQUIT: 3, + SIGTERM: 15, + SIGHUP: 1, + SIGKILL: 9, + }; + return table[name]; +} diff --git a/cli/src/commands/datatable/serve.ts b/cli/src/commands/datatable/serve.ts new file mode 100644 index 0000000000..4988f7a385 --- /dev/null +++ b/cli/src/commands/datatable/serve.ts @@ -0,0 +1,844 @@ +import { colors } from "@cliffy/ansi/colors"; +import { createServer, type Socket } from "node:net"; +import { randomBytes } from "node:crypto"; +import * as getPort from "get-port"; +import { BackendError, createPreHashedPassword } from "pg-gateway"; +import { fromNodeSocket } from "pg-gateway/node"; + +import * as wmill from "../../../gen/services.gen.ts"; +import { requireLogin } from "../../core/auth.ts"; +import { resolveWorkspace } from "../../core/context.ts"; +import * as log from "../../core/log.ts"; +import { GlobalOptions } from "../../types.ts"; +import { pollJobWithQueueLogging } from "../../utils/job_polling.ts"; +import { + buildBindComplete, + buildCloseComplete, + buildEmptyQueryResponse, + buildExecuteResponse, + buildNoData, + buildParameterDescription, + buildParseComplete, + buildQueryResponse, + buildReadyForQuery, + buildRowDescription, + concat, + type RawOutputEnvelope, +} from "./pg_wire.ts"; +import { + queryTouchesVirtualDatabaseCatalog, + virtualizeDatabaseCatalogQuery, +} from "./virtual_catalog.ts"; + +const DEFAULT_USER = "wmill"; +const DEFAULT_PORT_RANGE_START = 5433; +const DEFAULT_PORT_RANGE_END = 5500; + +export interface ServeOpts extends GlobalOptions { + port?: number; + host?: string; + password?: string; +} + +export interface ServeHandle { + host: string; + port: number; + user: string; + password: string; + /** Build a `postgresql://…/` URL for a given datatable name. */ + connectionString: (datatableName: string) => string; + close: () => Promise; +} + +interface PreparedStatementState { + query: string; + parameterTypeOids: number[]; +} + +interface PortalState { + query: string; + resultFormatCodes: number[]; + parameterTypeOids: number[]; + parameterSqlExpressions: string[]; + envelope?: RawOutputEnvelope; + rowOffset: number; + sentRowDescription: boolean; +} + +/** + * Bind the Postgres-wire proxy and return once it's accepting connections. + * Each client connection picks its target datatable via the standard + * Postgres `database` startup parameter (the `…/dbname` segment of the URL), + * so the same listener can multiplex any number of datatables — and clients + * like pgAdmin list them as separate databases instead of one shared "main". + * + * Caller owns `handle.close()`. Used by both `serve` (block on signal) and + * `psql` (block on child process exit). + */ +export async function startServe(opts: ServeOpts): Promise { + const host = opts.host ?? "127.0.0.1"; + + const workspace = await resolveWorkspace(opts); + await requireLogin(opts); + + const port = + opts.port ?? + (await (getPort as any).default({ + port: (getPort as any).portNumbers( + DEFAULT_PORT_RANGE_START, + DEFAULT_PORT_RANGE_END, + ), + })); + + const password = opts.password ?? randomBytes(12).toString("hex"); + const preHashedPassword = await createPreHashedPassword( + DEFAULT_USER, + password, + ); + + const server = createServer((socket) => { + handleConnection(socket, workspace.workspaceId, preHashedPassword) + .catch((err) => { + log.warn( + colors.yellow( + `connection error: ${err instanceof Error ? err.message : err}`, + ), + ); + try { + socket.destroy(); + } catch { + // ignore + } + }); + }); + + server.on("error", (err) => { + log.error(`server error: ${err.message}`); + process.exitCode = 1; + }); + + await new Promise((resolve) => server.listen(port, host, () => resolve())); + + return { + host, + port, + user: DEFAULT_USER, + password, + connectionString: (datatableName: string) => + `postgresql://${DEFAULT_USER}:${encodeURIComponent(password)}@${host}:${port}/${encodeURIComponent(datatableName)}`, + close: () => + new Promise((resolve) => { + server.close(() => resolve()); + // Force-resolve if open client connections keep the server from closing. + setTimeout(resolve, 1000).unref(); + }), + }; +} + +export async function serve(opts: ServeOpts): Promise { + const handle = await startServe(opts); + + // List available datatables so the user sees what they can `\c` into. + const workspace = await resolveWorkspace(opts); + let datatableNames: string[] = []; + try { + const items = await wmill.listDataTables({ + workspace: workspace.workspaceId, + }); + datatableNames = items.map((x) => x.name); + } catch { + // Listing is best-effort — the proxy still serves whatever the client asks for. + } + + log.info(colors.gray(`Serving datatables on ${handle.host}:${handle.port} via Postgres wire protocol`)); + log.info(""); + if (datatableNames.length > 0) { + log.info(colors.gray("Available datatables:")); + for (const name of datatableNames) { + log.info(` ${colors.bold("psql")} '${handle.connectionString(name)}'`); + } + } else { + log.info( + ` ${colors.bold("psql")} '${handle.connectionString("")}'`, + ); + } + log.info(""); + log.info(colors.gray("Press Ctrl+C to stop.")); + + const shutdown = () => { + log.info(colors.gray("\nShutting down...")); + handle.close().finally(() => process.exit(0)); + // Force-exit if connections linger. + setTimeout(() => process.exit(0), 1000).unref(); + }; + process.once("SIGINT", shutdown); + process.once("SIGTERM", shutdown); + + await new Promise(() => {}); +} + +async function handleConnection( + socket: Socket, + workspaceId: string, + preHashedPassword: string, +): Promise { + const preparedStatements = new Map(); + const portals = new Map(); + let waitingForSyncAfterError = false; + + await fromNodeSocket(socket, { + auth: { + method: "md5", + // Same pre-hashed value regardless of the username the client sends. + // Postgres's md5 auth derives preHashedPassword from `password + username`, + // so a client connecting as anything other than DEFAULT_USER will produce + // a different client-side hash and fail validation — effectively locking + // the listener to DEFAULT_USER. + getPreHashedPassword: async () => preHashedPassword, + }, + // pg-gateway 0.3.0-beta.4 declares `onQuery` in its TypeScript types but + // never actually invokes it — every frontend message goes through + // `onMessage`, and unhandled codes fall through to a default handler that + // returns "Message code not yet implemented". We handle the Query ('Q') + // message ourselves and only intercept post-auth messages so the library + // still drives the startup/auth handshake. + onMessage: async (data, state) => { + if (state.step !== "ReadyForQuery") return undefined; + const code = data[0]; + if (code === FE_TERMINATE) { + return undefined; // let pg-gateway close the connection + } + // Standard Postgres startup includes a `database` parameter — that's the + // path segment in the URL (`postgresql://…/`). pg-gateway exposes + // it under clientParams; we use it to route the query to the matching + // datatable. Falling back to the user name matches libpq's own default. + const datatableName = + state.clientParams?.database ?? state.clientParams?.user ?? "main"; + if (waitingForSyncAfterError) { + if (code === FE_SYNC) { + waitingForSyncAfterError = false; + return buildReadyForQuery("I"); + } + return new Uint8Array(0); + } + + try { + switch (code) { + case FE_QUERY: + return await runQuery(workspaceId, datatableName, readSimpleQuery(data)); + case FE_PARSE: + return handleParseMessage(data, preparedStatements, portals); + case FE_BIND: + return handleBindMessage(data, preparedStatements, portals); + case FE_DESCRIBE: + return await handleDescribeMessage( + data, + preparedStatements, + portals, + (query) => runQueryEnvelope(workspaceId, datatableName, query), + ); + case FE_EXECUTE: + return await handleExecuteMessage( + data, + portals, + (query) => runQueryEnvelope(workspaceId, datatableName, query), + ); + case FE_CLOSE: + return handleCloseMessage(data, preparedStatements, portals); + case FE_FLUSH: + return new Uint8Array(0); + case FE_SYNC: + return buildReadyForQuery("I"); + default: + throw createPgProtocolError( + `Postgres message code 0x${code.toString(16).padStart(2, "0")} is not supported by 'wmill datatable serve'`, + "0A000", + ); + } + } catch (err) { + const includeReadyForQuery = code === FE_QUERY; + waitingForSyncAfterError = !includeReadyForQuery; + return buildPgErrorResponse(err, includeReadyForQuery); + } + }, + }); +} + +const FE_QUERY = 0x51; // 'Q' +const FE_PARSE = 0x50; // 'P' +const FE_BIND = 0x42; // 'B' +const FE_CLOSE = 0x43; // 'C' +const FE_DESCRIBE = 0x44; // 'D' +const FE_EXECUTE = 0x45; // 'E' +const FE_FLUSH = 0x48; // 'H' +const FE_SYNC = 0x53; // 'S' +const FE_TERMINATE = 0x58; // 'X' + +type PgProtocolError = Error & { pgCode?: string }; + +const textDecoder = new TextDecoder(); + +class FrontendMessageReader { + constructor( + private readonly data: Uint8Array, + private offset = 5, + ) {} + + byte(): number { + return this.data[this.offset++] ?? 0; + } + + int16(): number { + const view = new DataView( + this.data.buffer, + this.data.byteOffset, + this.data.byteLength, + ); + const value = view.getInt16(this.offset); + this.offset += 2; + return value; + } + + int32(): number { + const view = new DataView( + this.data.buffer, + this.data.byteOffset, + this.data.byteLength, + ); + const value = view.getInt32(this.offset); + this.offset += 4; + return value; + } + + cstring(): string { + const start = this.offset; + while (this.offset < this.data.length && this.data[this.offset] !== 0) { + this.offset += 1; + } + const value = textDecoder.decode(this.data.subarray(start, this.offset)); + this.offset += 1; + return value; + } + + bytes(length: number): Uint8Array { + const value = this.data.subarray(this.offset, this.offset + length); + this.offset += length; + return value; + } +} + +/** + * Decode a simple-query message: 'Q', INT32 length (incl. itself), then a + * null-terminated UTF-8 string. + */ +function readSimpleQuery(data: Uint8Array): string { + const length = + (data[1] << 24) | (data[2] << 16) | (data[3] << 8) | data[4]; + // Strip the trailing null byte. + const queryBytes = data.subarray(5, 1 + length - 1); + return textDecoder.decode(queryBytes); +} + +function readParseMessage(data: Uint8Array): { + statementName: string; + query: string; + parameterTypeOids: number[]; +} { + const reader = new FrontendMessageReader(data); + const statementName = reader.cstring(); + const query = reader.cstring(); + const parameterCount = reader.int16(); + const parameterTypeOids = Array.from({ length: parameterCount }, () => + reader.int32(), + ); + return { statementName, query, parameterTypeOids }; +} + +function readBindMessage(data: Uint8Array): { + portalName: string; + statementName: string; + parameterFormatCodes: number[]; + parameters: (Uint8Array | null)[]; + resultFormatCodes: number[]; +} { + const reader = new FrontendMessageReader(data); + const portalName = reader.cstring(); + const statementName = reader.cstring(); + const parameterFormatCodeCount = reader.int16(); + const parameterFormatCodes = Array.from( + { length: parameterFormatCodeCount }, + () => reader.int16(), + ); + const parameterCount = reader.int16(); + const parameters = Array.from({ length: parameterCount }, () => { + const length = reader.int32(); + return length === -1 ? null : reader.bytes(length); + }); + const resultFormatCodeCount = reader.int16(); + const resultFormatCodes = Array.from( + { length: resultFormatCodeCount }, + () => reader.int16(), + ); + return { + portalName, + statementName, + parameterFormatCodes, + parameters, + resultFormatCodes, + }; +} + +function readDescribeMessage(data: Uint8Array): { + targetType: "P" | "S"; + name: string; +} { + const reader = new FrontendMessageReader(data); + const targetType = String.fromCharCode(reader.byte()) as "P" | "S"; + return { targetType, name: reader.cstring() }; +} + +function readExecuteMessage(data: Uint8Array): { + portalName: string; + maxRows: number; +} { + const reader = new FrontendMessageReader(data); + return { portalName: reader.cstring(), maxRows: reader.int32() }; +} + +function readCloseMessage(data: Uint8Array): { + targetType: "P" | "S"; + name: string; +} { + const reader = new FrontendMessageReader(data); + const targetType = String.fromCharCode(reader.byte()) as "P" | "S"; + return { targetType, name: reader.cstring() }; +} + +function handleParseMessage( + data: Uint8Array, + preparedStatements: Map, + portals: Map, +): Uint8Array { + const parse = readParseMessage(data); + preparedStatements.set(parse.statementName, { + query: parse.query, + parameterTypeOids: parse.parameterTypeOids, + }); + if (parse.statementName === "") { + portals.delete(""); + } + return buildParseComplete(); +} + +function handleBindMessage( + data: Uint8Array, + preparedStatements: Map, + portals: Map, +): Uint8Array { + const bind = readBindMessage(data); + const statement = preparedStatements.get(bind.statementName); + if (!statement) { + throw createPgProtocolError( + `prepared statement ${formatProtocolName(bind.statementName, "statement")} does not exist`, + "26000", + ); + } + for (let i = 0; i < bind.parameters.length; i += 1) { + if (bind.parameters[i] !== null && resolveParameterFormatCode(bind.parameterFormatCodes, i) !== 0) { + throw createPgProtocolError( + "binary parameter formats are not supported by 'wmill datatable serve'", + "0A000", + ); + } + } + if (bind.resultFormatCodes.some((code) => code !== 0)) { + throw createPgProtocolError( + "binary result formats are not supported by 'wmill datatable serve'", + "0A000", + ); + } + portals.set(bind.portalName, { + query: statement.query, + resultFormatCodes: bind.resultFormatCodes, + parameterTypeOids: statement.parameterTypeOids, + parameterSqlExpressions: bind.parameters.map(sqlExpressionForBoundParameter), + rowOffset: 0, + sentRowDescription: false, + }); + return buildBindComplete(); +} + +async function handleDescribeMessage( + data: Uint8Array, + preparedStatements: Map, + portals: Map, + executeQuery: (query: string) => Promise, +): Promise { + const describe = readDescribeMessage(data); + if (describe.targetType === "S") { + const statement = preparedStatements.get(describe.name); + if (!statement) { + throw createPgProtocolError( + `prepared statement ${formatProtocolName(describe.name, "statement")} does not exist`, + "26000", + ); + } + if ( + statement.parameterTypeOids.length === 0 && + isQueryLikelySafeToDescribe(statement.query) + ) { + const envelope = await executeQuery(statement.query); + return concat([ + buildParameterDescription(statement.parameterTypeOids), + envelope.columns.length > 0 + ? buildRowDescription(envelope.columns) + : buildNoData(), + ]); + } + return concat([ + buildParameterDescription(statement.parameterTypeOids), + buildNoData(), + ]); + } + + const portal = portals.get(describe.name); + if (!portal) { + throw createPgProtocolError( + `portal ${formatProtocolName(describe.name, "portal")} does not exist`, + "34000", + ); + } + + if (!portal.envelope && isQueryLikelySafeToDescribe(portal.query)) { + portal.envelope = await executeQuery(buildPortalQuery(portal)); + } + if (portal.envelope?.columns.length) { + portal.sentRowDescription = true; + return buildRowDescription(portal.envelope.columns); + } + return buildNoData(); +} + +async function handleExecuteMessage( + data: Uint8Array, + portals: Map, + executeQuery: (query: string) => Promise, +): Promise { + const execute = readExecuteMessage(data); + const portal = portals.get(execute.portalName); + if (!portal) { + throw createPgProtocolError( + `portal ${formatProtocolName(execute.portalName, "portal")} does not exist`, + "34000", + ); + } + if (portal.resultFormatCodes.some((code) => code !== 0)) { + throw createPgProtocolError( + "binary result formats are not supported by 'wmill datatable serve'", + "0A000", + ); + } + if (portal.query.trim().length === 0) { + return buildEmptyQueryResponse(); + } + if (!portal.envelope) { + portal.envelope = await executeQuery(buildPortalQuery(portal)); + } + + const response = buildExecuteResponse(portal.envelope, { + includeRowDescription: !portal.sentRowDescription, + maxRows: execute.maxRows, + rowOffset: portal.rowOffset, + }); + + portal.sentRowDescription = true; + if (response.suspended) { + portal.rowOffset = response.nextRowOffset; + } else { + portal.envelope = undefined; + portal.rowOffset = 0; + portal.sentRowDescription = false; + } + + return response.message; +} + +function handleCloseMessage( + data: Uint8Array, + preparedStatements: Map, + portals: Map, +): Uint8Array { + const close = readCloseMessage(data); + if (close.targetType === "S") { + preparedStatements.delete(close.name); + } else { + portals.delete(close.name); + } + return buildCloseComplete(); +} + +function createPgProtocolError(message: string, pgCode = "XX000"): PgProtocolError { + const err = new Error(message) as PgProtocolError; + err.pgCode = pgCode; + return err; +} + +function resolveParameterFormatCode(formatCodes: number[], index: number): number { + if (formatCodes.length === 0) return 0; + if (formatCodes.length === 1) return formatCodes[0] ?? 0; + return formatCodes[index] ?? 0; +} + +function sqlExpressionForBoundParameter(value: Uint8Array | null): string { + if (value === null) { + return "NULL"; + } + return sqlStringLiteral(textDecoder.decode(value)); +} + +function sqlStringLiteral(value: string): string { + return `'${value.replaceAll("'", "''")}'`; +} + +function buildPortalQuery(portal: PortalState): string { + if (portal.parameterSqlExpressions.length === 0) { + return portal.query; + } + + const trimmedQuery = portal.query.replace(/;+\s*$/, ""); + const typeList = buildPreparedStatementTypeList(portal.parameterTypeOids); + return [ + `PREPARE "__wmill_proxy_stmt"${typeList} AS ${trimmedQuery};`, + `EXECUTE "__wmill_proxy_stmt"(${portal.parameterSqlExpressions.join(", ")});`, + ].join("\n"); +} + +function buildPreparedStatementTypeList(parameterTypeOids: number[]): string { + if (parameterTypeOids.length === 0) { + return ""; + } + + const typeNames = parameterTypeOids.map(postgresTypeNameFromOid); + if (typeNames.some((name) => name === undefined)) { + return ""; + } + + return `(${typeNames.join(", ")})`; +} + +function postgresTypeNameFromOid(oid: number): string | undefined { + switch (oid) { + case 16: + return "pg_catalog.bool"; + case 17: + return "pg_catalog.bytea"; + case 19: + return "pg_catalog.name"; + case 20: + return "pg_catalog.int8"; + case 21: + return "pg_catalog.int2"; + case 23: + return "pg_catalog.int4"; + case 25: + return "pg_catalog.text"; + case 26: + return "pg_catalog.oid"; + case 700: + return "pg_catalog.float4"; + case 701: + return "pg_catalog.float8"; + case 1042: + return "pg_catalog.bpchar"; + case 1043: + return "pg_catalog.varchar"; + case 1082: + return "pg_catalog.date"; + case 1083: + return "pg_catalog.time"; + case 1114: + return "pg_catalog.timestamp"; + case 1184: + return "pg_catalog.timestamptz"; + case 1266: + return "pg_catalog.timetz"; + case 1700: + return "pg_catalog.numeric"; + case 2950: + return "pg_catalog.uuid"; + case 114: + return "pg_catalog.json"; + case 3802: + return "pg_catalog.jsonb"; + default: + return undefined; + } +} + +function isQueryLikelySafeToDescribe(query: string): boolean { + let normalized = query.trimStart(); + while (true) { + if (normalized.startsWith("--")) { + const newlineIndex = normalized.indexOf("\n"); + normalized = + newlineIndex === -1 + ? "" + : normalized.slice(newlineIndex + 1).trimStart(); + continue; + } + if (normalized.startsWith("/*")) { + const commentEnd = normalized.indexOf("*/"); + normalized = + commentEnd === -1 + ? "" + : normalized.slice(commentEnd + 2).trimStart(); + continue; + } + break; + } + normalized = normalized.toLowerCase(); + + return ( + normalized.startsWith("select") || + normalized.startsWith("with") || + normalized.startsWith("show") || + normalized.startsWith("values") || + normalized.startsWith("table") + ); +} + +function buildPgErrorResponse( + err: unknown, + includeReadyForQuery: boolean, +): Uint8Array { + const error = err instanceof Error ? err : new Error(String(err ?? "unknown error")); + const parts = [ + BackendError.create({ + severity: "ERROR", + code: (error as PgProtocolError).pgCode ?? "XX000", + message: error.message, + }).flush(), + ]; + if (includeReadyForQuery) { + parts.push(buildReadyForQuery("I")); + } + return concat(parts); +} + +function formatProtocolName(name: string, kind: "portal" | "statement"): string { + if (name.length === 0) { + return kind === "portal" ? "" : ""; + } + return `"${name}"`; +} + +async function runQuery( + workspaceId: string, + datatableName: string, + query: string, +): Promise { + const trimmed = query.trim(); + if (trimmed.length === 0) { + return concat([buildEmptyQueryResponse(), buildReadyForQuery("I")]); + } + + const envelope = await runQueryEnvelope(workspaceId, datatableName, query); + return buildQueryResponse(envelope); +} + +async function runQueryEnvelope( + workspaceId: string, + datatableName: string, + query: string, +): Promise { + const trimmed = query.trim(); + if (trimmed.length === 0) { + return { columns: [], rows: [] }; + } + + // The annotation parser stops at the first non-blank line that isn't a + // `--` comment, so prepending guarantees it's honored even if the client + // sent leading comments of its own. + const content = `-- raw_output\n${await prepareQueryForDatatableProxy( + workspaceId, + datatableName, + query, + )}`; + + const jobId = await wmill.runScriptPreview({ + workspace: workspaceId, + requestBody: { + content, + language: "postgresql", + args: { database: `datatable://${datatableName}` }, + }, + }); + + const { result, success } = await pollJobWithQueueLogging(workspaceId, jobId); + + if (!success) { + throw createPgProtocolError(extractErrorMessage(result)); + } + + return coerceEnvelope(result); +} + +async function prepareQueryForDatatableProxy( + workspaceId: string, + datatableName: string, + query: string, +): Promise { + if (!queryTouchesVirtualDatabaseCatalog(query)) { + return query; + } + + let databaseNames = [datatableName]; + if (query.toLowerCase().includes("pg_database")) { + try { + const items = await wmill.listDataTables({ + workspace: workspaceId, + }); + databaseNames = Array.from( + new Set([datatableName, ...items.map((item) => item.name)]), + ).sort(); + } catch { + // Fall back to the current datatable so catalog queries still stay virtual. + } + } + + return virtualizeDatabaseCatalogQuery(query, datatableName, databaseNames); +} + +function extractErrorMessage(result: unknown): string { + if (typeof result === "string") return result; + if (result && typeof result === "object") { + const obj = result as Record; + const err = obj.error; + if (err && typeof err === "object") { + const errObj = err as Record; + if (typeof errObj.message === "string") return errObj.message; + if (typeof errObj.name === "string") return errObj.name; + } + if (typeof obj.message === "string") return obj.message; + } + return JSON.stringify(result); +} + +/** + * Validate that the worker returned the {columns, rows} envelope we asked for. + * If the shape doesn't match (older backend, unexpected error shape), fall back + * to an empty result — better than throwing inside the wire-protocol layer. + */ +function coerceEnvelope(result: unknown): RawOutputEnvelope { + if ( + result && + typeof result === "object" && + "columns" in result && + "rows" in result && + Array.isArray((result as any).columns) && + Array.isArray((result as any).rows) + ) { + return result as RawOutputEnvelope; + } + return { columns: [], rows: [] }; +} diff --git a/cli/src/commands/datatable/virtual_catalog.ts b/cli/src/commands/datatable/virtual_catalog.ts new file mode 100644 index 0000000000..2c8a997eb1 --- /dev/null +++ b/cli/src/commands/datatable/virtual_catalog.ts @@ -0,0 +1,75 @@ +const CURRENT_DATABASE_QUERY_RE = /\b(?:pg_catalog\.)?current_database\s*\(\s*\)/i; +const CURRENT_DATABASE_RE = /\b(?:pg_catalog\.)?current_database\s*\(\s*\)/gi; +const PG_DATABASE_RE = /\b(?:pg_catalog\.)?pg_database\b/i; +const PG_DATABASE_FROM_OR_JOIN_RE = + /(\b(?:from|join)\s+(?:only\s+)?)(?:pg_catalog\.)?pg_database\b/gi; +const PG_DATABASE_COMMA_JOIN_RE = /(,\s*)(?:pg_catalog\.)?pg_database\b/gi; +const PG_SIZE_PRETTY_DATABASE_SIZE_RE = + /\b(?:pg_catalog\.)?pg_size_pretty\s*\(\s*(?:pg_catalog\.)?pg_database_size\s*\((?:[^()]|\([^()]*\))*\)\s*\)/gi; +const PG_DATABASE_SIZE_RE = + /\b(?:pg_catalog\.)?pg_database_size\s*\((?:[^()]|\([^()]*\))*\)/gi; +const HAS_DATABASE_PRIVILEGE_RE = + /\b(?:pg_catalog\.)?has_database_privilege\s*\((?:[^()]|\([^()]*\))*\)/gi; +const PG_DATABASE_COLLATION_VERSION_RE = + /\b(?:pg_catalog\.)?pg_database_collation_actual_version\s*\((?:[^()]|\([^()]*\))*\)/gi; + +export function queryTouchesVirtualDatabaseCatalog(query: string): boolean { + return CURRENT_DATABASE_QUERY_RE.test(query) || PG_DATABASE_RE.test(query); +} + +export function virtualizeDatabaseCatalogQuery( + query: string, + currentDatabaseName: string, + databaseNames: string[], +): string { + let rewritten = query.replace( + CURRENT_DATABASE_RE, + sqlStringLiteral(currentDatabaseName), + ); + + if (!PG_DATABASE_RE.test(query)) { + return rewritten; + } + + const relation = buildVirtualPgDatabaseRelation(databaseNames); + rewritten = rewritten + .replace(PG_SIZE_PRETTY_DATABASE_SIZE_RE, "NULL::text") + .replace(PG_DATABASE_SIZE_RE, "NULL::bigint") + .replace(HAS_DATABASE_PRIVILEGE_RE, "true") + .replace(PG_DATABASE_COLLATION_VERSION_RE, "NULL::text") + .replace(PG_DATABASE_FROM_OR_JOIN_RE, `$1${relation}`) + .replace(PG_DATABASE_COMMA_JOIN_RE, `$1${relation}`); + + return rewritten; +} + +function buildVirtualPgDatabaseRelation(databaseNames: string[]): string { + const names = Array.from(new Set(databaseNames)); + const arrayLiteral = names.map(sqlStringLiteral).join(", "); + + return `( + SELECT + (20000 + src.ordinality)::oid AS oid, + src.datname, + COALESCE((SELECT oid FROM pg_catalog.pg_roles WHERE rolname = CURRENT_USER LIMIT 1), 10::oid) AS datdba, + pg_char_to_encoding(current_setting('server_encoding')) AS encoding, + 'c'::"char" AS datlocprovider, + false AS datistemplate, + true AS datallowconn, + -1 AS datconnlimit, + '0'::xid AS datfrozenxid, + '1'::xid AS datminmxid, + COALESCE((SELECT oid FROM pg_catalog.pg_tablespace WHERE spcname = 'pg_default' LIMIT 1), 1663::oid) AS dattablespace, + current_setting('lc_collate') AS datcollate, + current_setting('lc_ctype') AS datctype, + NULL::text AS daticulocale, + NULL::text AS daticurules, + NULL::text AS datcollversion, + NULL::aclitem[] AS datacl + FROM unnest(ARRAY[${arrayLiteral}]::text[]) WITH ORDINALITY AS src(datname, ordinality) +)`; +} + +function sqlStringLiteral(value: string): string { + return `'${value.replaceAll("'", "''")}'`; +} diff --git a/cli/src/guidance/skills.gen.ts b/cli/src/guidance/skills.gen.ts index 3ef4f913c5..8e8e447aed 100644 --- a/cli/src/guidance/skills.gen.ts +++ b/cli/src/guidance/skills.gen.ts @@ -6691,6 +6691,15 @@ datatable related commands - \`datatable run \` - run a SQL query on a datatable - \`-n --name \` - Datatable name (default: main) - \`-s --silent\` - Output only the final result as JSON. Useful for scripting. +- \`datatable serve\` - Serve all datatables as a Postgres-wire endpoint (psql, DBeaver, pgAdmin); the client picks the datatable via the database name in its connection string + - \`--port \` - Port to listen on (default: first free port in 5433-5500) + - \`--host \` - Bind address (default: 127.0.0.1) + - \`--password \` - Password for Postgres clients (default: generate a random password at startup) +- \`datatable psql\` - Start a serve listener and launch psql connected to it + - \`-n --name \` - Datatable to connect psql to (default: main) + - \`--port \` - Port the proxy listens on (default: first free port in 5433-5500) + - \`--host \` - Bind address for the proxy (default: 127.0.0.1) + - \`--password \` - Password for the temporary Postgres proxy (default: generate a random password at startup) ### dependencies diff --git a/cli/test/datatable_pg_wire_unit.test.ts b/cli/test/datatable_pg_wire_unit.test.ts new file mode 100644 index 0000000000..ab16842899 --- /dev/null +++ b/cli/test/datatable_pg_wire_unit.test.ts @@ -0,0 +1,193 @@ +import { expect, test } from "bun:test"; +import { + buildBindComplete, + buildCloseComplete, + buildExecuteResponse, + buildCommandComplete, + buildDataRow, + buildEmptyQueryResponse, + buildNoData, + buildParameterDescription, + buildParseComplete, + buildPortalSuspended, + buildQueryResponse, + buildReadyForQuery, + buildRowDescription, + type RawOutputEnvelope, +} from "../src/commands/datatable/pg_wire.ts"; + +// Postgres wire-protocol byte-level expectations. Every backend message is: +// 1 byte type | int32 length-including-itself | payload + +test("buildReadyForQuery: 'Z' + len=5 + 'I'", () => { + const buf = buildReadyForQuery("I"); + expect(Array.from(buf)).toEqual([ + 0x5a, // 'Z' + 0, 0, 0, 5, // length = 4 (length field) + 1 (status) + 0x49, // 'I' + ]); +}); + +test("buildCommandComplete: 'C' + tag + null terminator", () => { + const buf = buildCommandComplete("SELECT 0"); + // 'C' + len(4) + 'SELECT 0'(8) + '\0'(1) → 13 bytes total, length field = 13 + expect(buf[0]).toBe(0x43); + expect(buf.byteLength).toBe(14); + expect(buf[buf.byteLength - 1]).toBe(0); +}); + +test("buildEmptyQueryResponse: 'I' + len=4", () => { + const buf = buildEmptyQueryResponse(); + expect(Array.from(buf)).toEqual([0x49, 0, 0, 0, 4]); +}); + +test("buildParseComplete / BindComplete / CloseComplete / NoData / PortalSuspended: empty payload messages", () => { + expect(Array.from(buildParseComplete())).toEqual([0x31, 0, 0, 0, 4]); + expect(Array.from(buildBindComplete())).toEqual([0x32, 0, 0, 0, 4]); + expect(Array.from(buildCloseComplete())).toEqual([0x33, 0, 0, 0, 4]); + expect(Array.from(buildNoData())).toEqual([0x6e, 0, 0, 0, 4]); + expect(Array.from(buildPortalSuspended())).toEqual([0x73, 0, 0, 0, 4]); +}); + +test("buildParameterDescription: count + OIDs", () => { + const buf = buildParameterDescription([23, 25]); + expect(buf[0]).toBe(0x74); // 't' + expect((buf[5] << 8) | buf[6]).toBe(2); + const firstOid = + (buf[7] << 24) | + (buf[8] << 16) | + (buf[9] << 8) | + buf[10]; + const secondOid = + (buf[11] << 24) | + (buf[12] << 16) | + (buf[13] << 8) | + buf[14]; + expect(firstOid).toBe(23); + expect(secondOid).toBe(25); +}); + +test("buildRowDescription: encodes one column with given OID", () => { + const buf = buildRowDescription([{ name: "id", oid: 23, type_name: "int4" }]); + // 'T'(1) + len(4) + field_count(2) + name(2 + null=3) + tableOid(4) + attnum(2) + // + typeOid(4) + typeSize(2) + atttypmod(4) + format(2) = 1 + 4 + 2 + 3 + 4 + 2 + 4 + 2 + 4 + 2 = 28 + expect(buf[0]).toBe(0x54); // 'T' + expect(buf.byteLength).toBe(28); + // field count INT16 starts at offset 5 (after type byte + length field) + expect((buf[5] << 8) | buf[6]).toBe(1); + // typeOid INT32 is at the end of the column header — at offset 28 - 2 - 4 - 2 - 4 = 16 + const typeOidOffset = 28 - 2 /* format */ - 4 /* typmod */ - 2 /* typesize */ - 4; + const typeOid = + (buf[typeOidOffset] << 24) | + (buf[typeOidOffset + 1] << 16) | + (buf[typeOidOffset + 2] << 8) | + buf[typeOidOffset + 3]; + expect(typeOid).toBe(23); +}); + +test("buildDataRow: NULL cell is encoded as int32 -1, length 0 payload", () => { + const buf = buildDataRow([null]); + // 'D'(1) + len(4) + col_count(2) + col_length(4 = -1) = 11 bytes + expect(Array.from(buf)).toEqual([ + 0x44, // 'D' + 0, 0, 0, 10, // length field = 4 + 2 + 4 + 0, 1, // col_count = 1 + 0xff, 0xff, 0xff, 0xff, // -1 (NULL) + ]); +}); + +test("buildDataRow: text cell carries UTF-8 bytes prefixed by length", () => { + const buf = buildDataRow(["hi"]); + expect(Array.from(buf)).toEqual([ + 0x44, + 0, 0, 0, 12, // length = 4 + 2 + 4 + 2 + 0, 1, // col count + 0, 0, 0, 2, // string length + 0x68, 0x69, // "hi" + ]); +}); + +test("buildDataRow: mixed text + NULL in column order", () => { + const buf = buildDataRow(["a", null, "bb"]); + // 'D'(1) + len(4) + col_count(2) + (4 + 1) + 4 + (4 + 2) = 22 + expect(buf.byteLength).toBe(22); + // col count + expect((buf[5] << 8) | buf[6]).toBe(3); + // First cell length=1, byte 'a' at offset 11 + expect(buf[11]).toBe(0x61); + // Second cell is NULL: bytes 12-15 should be 0xff + expect(buf[12]).toBe(0xff); + expect(buf[15]).toBe(0xff); + // Third cell length=2 at offset 16, bytes 'b','b' at 20-21 + expect(buf[20]).toBe(0x62); + expect(buf[21]).toBe(0x62); +}); + +test("buildQueryResponse: empty envelope emits only CommandComplete + ReadyForQuery", () => { + const envelope: RawOutputEnvelope = { columns: [], rows: [] }; + const buf = buildQueryResponse(envelope); + // CommandComplete "SELECT 0" = 14 bytes, ReadyForQuery = 6 bytes + expect(buf.byteLength).toBe(14 + 6); + expect(buf[0]).toBe(0x43); // 'C' + expect(buf[14]).toBe(0x5a); // 'Z' +}); + +test("buildQueryResponse: one column + one row emits T, D, C, Z in that order", () => { + const envelope: RawOutputEnvelope = { + columns: [{ name: "n", oid: 23, type_name: "int4" }], + rows: [["42"]], + }; + const buf = buildQueryResponse(envelope); + const tagBytes: number[] = []; + let off = 0; + while (off < buf.byteLength) { + tagBytes.push(buf[off]); + const len = (buf[off + 1] << 24) | (buf[off + 2] << 16) | (buf[off + 3] << 8) | buf[off + 4]; + off += 1 + len; + } + expect(tagBytes).toEqual([0x54 /* T */, 0x44 /* D */, 0x43 /* C */, 0x5a /* Z */]); +}); + +test("buildExecuteResponse: emits T, D, C without ReadyForQuery", () => { + const envelope: RawOutputEnvelope = { + columns: [{ name: "n", oid: 23, type_name: "int4" }], + rows: [["42"]], + }; + const response = buildExecuteResponse(envelope); + const tagBytes: number[] = []; + let off = 0; + while (off < response.message.byteLength) { + tagBytes.push(response.message[off]); + const len = + (response.message[off + 1] << 24) | + (response.message[off + 2] << 16) | + (response.message[off + 3] << 8) | + response.message[off + 4]; + off += 1 + len; + } + expect(response.suspended).toBe(false); + expect(response.nextRowOffset).toBe(1); + expect(tagBytes).toEqual([0x54 /* T */, 0x44 /* D */, 0x43 /* C */]); +}); + +test("buildExecuteResponse: maxRows suspends portal and omits CommandComplete", () => { + const envelope: RawOutputEnvelope = { + columns: [{ name: "n", oid: 23, type_name: "int4" }], + rows: [["1"], ["2"]], + }; + const response = buildExecuteResponse(envelope, { maxRows: 1 }); + const tagBytes: number[] = []; + let off = 0; + while (off < response.message.byteLength) { + tagBytes.push(response.message[off]); + const len = + (response.message[off + 1] << 24) | + (response.message[off + 2] << 16) | + (response.message[off + 3] << 8) | + response.message[off + 4]; + off += 1 + len; + } + expect(response.suspended).toBe(true); + expect(response.nextRowOffset).toBe(1); + expect(tagBytes).toEqual([0x54 /* T */, 0x44 /* D */, 0x73 /* s */]); +}); diff --git a/cli/test/datatable_virtual_catalog_unit.test.ts b/cli/test/datatable_virtual_catalog_unit.test.ts new file mode 100644 index 0000000000..6621fe68be --- /dev/null +++ b/cli/test/datatable_virtual_catalog_unit.test.ts @@ -0,0 +1,46 @@ +import { expect, test } from "bun:test"; + +import { + queryTouchesVirtualDatabaseCatalog, + virtualizeDatabaseCatalogQuery, +} from "../src/commands/datatable/virtual_catalog.ts"; + +test("queryTouchesVirtualDatabaseCatalog: detects current_database and pg_database", () => { + expect(queryTouchesVirtualDatabaseCatalog("select current_database()")).toBe(true); + expect(queryTouchesVirtualDatabaseCatalog("select * from pg_catalog.pg_database")).toBe(true); + expect(queryTouchesVirtualDatabaseCatalog("select * from public.users")).toBe(false); +}); + +test("virtualizeDatabaseCatalogQuery: rewrites current_database to the virtual datatable name", () => { + expect( + virtualizeDatabaseCatalogQuery( + "select current_database() as db", + "sales-main", + ["sales-main"], + ), + ).toBe("select 'sales-main' as db"); +}); + +test("virtualizeDatabaseCatalogQuery: rewrites pg_database reads to the virtual datatable catalog", () => { + const rewritten = virtualizeDatabaseCatalogQuery( + "select datname from pg_catalog.pg_database where datname = current_database() order by datname", + "sales-main", + ["sales-main", "warehouse"], + ); + + expect(rewritten).toContain("FROM unnest(ARRAY['sales-main', 'warehouse']::text[]) WITH ORDINALITY"); + expect(rewritten).toContain("where datname = 'sales-main'"); + expect(rewritten).not.toContain("pg_catalog.pg_database"); +}); + +test("virtualizeDatabaseCatalogQuery: neutralizes pg_database metadata helpers that need a real upstream database", () => { + const rewritten = virtualizeDatabaseCatalogQuery( + "select pg_size_pretty(pg_database_size(datname)) as size, has_database_privilege(datname, 'CONNECT') from pg_database", + "sales-main", + ["sales-main"], + ); + + expect(rewritten).toContain("NULL::text as size"); + expect(rewritten).toContain("true"); + expect(rewritten).not.toContain("pg_database_size"); +}); diff --git a/system_prompts/auto-generated/cli/cli-commands.md b/system_prompts/auto-generated/cli/cli-commands.md index c775aa2d28..9917afd696 100644 --- a/system_prompts/auto-generated/cli/cli-commands.md +++ b/system_prompts/auto-generated/cli/cli-commands.md @@ -77,6 +77,15 @@ datatable related commands - `datatable run ` - run a SQL query on a datatable - `-n --name ` - Datatable name (default: main) - `-s --silent` - Output only the final result as JSON. Useful for scripting. +- `datatable serve` - Serve all datatables as a Postgres-wire endpoint (psql, DBeaver, pgAdmin); the client picks the datatable via the database name in its connection string + - `--port ` - Port to listen on (default: first free port in 5433-5500) + - `--host ` - Bind address (default: 127.0.0.1) + - `--password ` - Password for Postgres clients (default: generate a random password at startup) +- `datatable psql` - Start a serve listener and launch psql connected to it + - `-n --name ` - Datatable to connect psql to (default: main) + - `--port ` - Port the proxy listens on (default: first free port in 5433-5500) + - `--host ` - Bind address for the proxy (default: 127.0.0.1) + - `--password ` - Password for the temporary Postgres proxy (default: generate a random password at startup) ### dependencies diff --git a/system_prompts/auto-generated/prompts.ts b/system_prompts/auto-generated/prompts.ts index e7f4c1f02c..35637f3a6c 100644 --- a/system_prompts/auto-generated/prompts.ts +++ b/system_prompts/auto-generated/prompts.ts @@ -2608,6 +2608,15 @@ datatable related commands - \`datatable run \` - run a SQL query on a datatable - \`-n --name \` - Datatable name (default: main) - \`-s --silent\` - Output only the final result as JSON. Useful for scripting. +- \`datatable serve\` - Serve all datatables as a Postgres-wire endpoint (psql, DBeaver, pgAdmin); the client picks the datatable via the database name in its connection string + - \`--port \` - Port to listen on (default: first free port in 5433-5500) + - \`--host \` - Bind address (default: 127.0.0.1) + - \`--password \` - Password for Postgres clients (default: generate a random password at startup) +- \`datatable psql\` - Start a serve listener and launch psql connected to it + - \`-n --name \` - Datatable to connect psql to (default: main) + - \`--port \` - Port the proxy listens on (default: first free port in 5433-5500) + - \`--host \` - Bind address for the proxy (default: 127.0.0.1) + - \`--password \` - Password for the temporary Postgres proxy (default: generate a random password at startup) ### dependencies diff --git a/system_prompts/auto-generated/skills/cli-commands/SKILL.md b/system_prompts/auto-generated/skills/cli-commands/SKILL.md index 5f7a2e9991..0e4ac513c4 100644 --- a/system_prompts/auto-generated/skills/cli-commands/SKILL.md +++ b/system_prompts/auto-generated/skills/cli-commands/SKILL.md @@ -82,6 +82,15 @@ datatable related commands - `datatable run ` - run a SQL query on a datatable - `-n --name ` - Datatable name (default: main) - `-s --silent` - Output only the final result as JSON. Useful for scripting. +- `datatable serve` - Serve all datatables as a Postgres-wire endpoint (psql, DBeaver, pgAdmin); the client picks the datatable via the database name in its connection string + - `--port ` - Port to listen on (default: first free port in 5433-5500) + - `--host ` - Bind address (default: 127.0.0.1) + - `--password ` - Password for Postgres clients (default: generate a random password at startup) +- `datatable psql` - Start a serve listener and launch psql connected to it + - `-n --name ` - Datatable to connect psql to (default: main) + - `--port ` - Port the proxy listens on (default: first free port in 5433-5500) + - `--host ` - Bind address for the proxy (default: 127.0.0.1) + - `--password ` - Password for the temporary Postgres proxy (default: generate a random password at startup) ### dependencies diff --git a/system_prompts/generate.py b/system_prompts/generate.py index 885ca5d458..12792da41d 100644 --- a/system_prompts/generate.py +++ b/system_prompts/generate.py @@ -283,11 +283,12 @@ def extract_py_classes(content: str) -> list[dict]: # ============================================================================= -# Reusable option pattern for CLI parsing +# Reusable option pattern for CLI parsing. Matches both `.option(...)` and +# `.globalOption(...)` so subcommand-level global options surface in the docs. OPTION_PATTERN = re.compile( - r'\.option\(\s*"([^"]+)"\s*,\s*"([^"]+)"' # double-quoted + r'\.(?:option|globalOption)\(\s*"([^"]+)"\s*,\s*"([^"]+)"' # double-quoted r'|' - r"\.option\(\s*'([^']+)'\s*,\s*'([^']+)'", # single-quoted + r"\.(?:option|globalOption)\(\s*'([^']+)'\s*,\s*'([^']+)'", # single-quoted re.MULTILINE | re.DOTALL ) @@ -352,7 +353,24 @@ def parse_command_block(content: str, file_path: Path | None = None) -> dict: subcommand_sections = re.split(r'(?=\.command\()', block) for section in subcommand_sections: - cmd_match = re.match(r'\.command\(\s*["\']([^"\']+)["\']\s*(?:,\s*("(?:[^"\\]|\\.)*"|\'(?:[^\'\\]|\\.)*\'|[^)]+))?\s*\)', section) + # Second arg is either a quoted description, a bare identifier (imported + # command, e.g. `.command("app", app)`), or a more complex expression + # like `someWrapper(new Command()...)` — the `[^)]+` fallback covers + # the last case by matching up to the next `)`. + # + # Two subtleties: + # - The quoted-string alts are tried first so a description containing + # `(` like "(psql, DBeaver)" isn't truncated by the `[^)]+` fallback. + # - The trailing `,?` accommodates the prettier-style `\n )` close + # paren that follows a comma. Without it, the quoted alt would + # succeed but the outer `\s*\)` would fail, forcing a backtrack to + # `[^)]+` and producing a truncated description with a trailing `",`. + cmd_match = re.match( + r'\.command\(\s*["\']([^"\']+)["\']\s*' + r'(?:,\s*("(?:[^"\\]|\\.)*"|\'(?:[^\'\\]|\\.)*\'|[^)]+))?' + r'\s*,?\s*\)', + section, + ) if not cmd_match: continue