mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-10-17 16:03:06 +00:00
feat: CLI datatable serve / psql (#9267)
* feat(cli): add datatable list and run commands * feat(cli): render datatable query results as a table * feat(cli): serve datatables as a postgres-wire endpoint * feat(cli): add 'datatable psql' to launch psql against the proxy * feat(cli): route datatable serve by client-supplied database name * override database list + password option * fix: support extended queries in datatable serve * fix: correct cloud size threshold log and parse CLI descriptions with parens/trailing comma * refactor: extract raw_output envelope encoding into pg_raw_output module --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
This commit is contained in:
Diego Imbert
committed by
Diego Imbert
Diego Imbert
co-authored by
claude[bot]
claude[bot]
parent
7ecd096457
commit
45dc3194e7
@@ -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";
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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<BoxFuture<'a, error::Result<Vec<Box<RawValue>>>>> {
|
||||
let mut query_params = vec![];
|
||||
let mut param_types: Vec<Type> = 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 {
|
||||
|
||||
@@ -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<RawOutputColumn>,
|
||||
pub rows: Vec<Vec<Option<String>>>,
|
||||
}
|
||||
|
||||
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> {
|
||||
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<Row>,
|
||||
format_state: &ResultFormatState,
|
||||
siz: &AtomicUsize,
|
||||
) -> error::Result<RawOutputEnvelope> {
|
||||
let columns: Vec<RawOutputColumn> = 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<Option<String>>> = 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<Option<String>> = 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<RawValue>` 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<Vec<Box<RawValue>>>) -> Box<RawValue> {
|
||||
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<String> {
|
||||
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\":[]}");
|
||||
}
|
||||
}
|
||||
@@ -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=="],
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:number>",
|
||||
"Port to listen on (default: first free port in 5433-5500)",
|
||||
)
|
||||
.option(
|
||||
"--host <host:string>",
|
||||
"Bind address (default: 127.0.0.1)",
|
||||
)
|
||||
.option(
|
||||
"--password <password:string>",
|
||||
"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 <name:string>",
|
||||
"Datatable to connect psql to (default: main)",
|
||||
)
|
||||
.option(
|
||||
"--port <port:number>",
|
||||
"Port the proxy listens on (default: first free port in 5433-5500)",
|
||||
)
|
||||
.option(
|
||||
"--host <host:string>",
|
||||
"Bind address for the proxy (default: 127.0.0.1)",
|
||||
)
|
||||
.option(
|
||||
"--password <password:string>",
|
||||
"Password for the temporary Postgres proxy (default: generate a random password at startup)",
|
||||
)
|
||||
.action(psql as any);
|
||||
|
||||
export default command;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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<void> {
|
||||
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<number>((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<string, number> = {
|
||||
SIGINT: 2,
|
||||
SIGQUIT: 3,
|
||||
SIGTERM: 15,
|
||||
SIGHUP: 1,
|
||||
SIGKILL: 9,
|
||||
};
|
||||
return table[name];
|
||||
}
|
||||
@@ -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://…/<dbName>` URL for a given datatable name. */
|
||||
connectionString: (datatableName: string) => string;
|
||||
close: () => Promise<void>;
|
||||
}
|
||||
|
||||
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<ServeHandle> {
|
||||
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<void>((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<void>((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<void> {
|
||||
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("<datatable_name>")}'`,
|
||||
);
|
||||
}
|
||||
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<void>(() => {});
|
||||
}
|
||||
|
||||
async function handleConnection(
|
||||
socket: Socket,
|
||||
workspaceId: string,
|
||||
preHashedPassword: string,
|
||||
): Promise<void> {
|
||||
const preparedStatements = new Map<string, PreparedStatementState>();
|
||||
const portals = new Map<string, PortalState>();
|
||||
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://…/<dbname>`). 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<string, PreparedStatementState>,
|
||||
portals: Map<string, PortalState>,
|
||||
): 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<string, PreparedStatementState>,
|
||||
portals: Map<string, PortalState>,
|
||||
): 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<string, PreparedStatementState>,
|
||||
portals: Map<string, PortalState>,
|
||||
executeQuery: (query: string) => Promise<RawOutputEnvelope>,
|
||||
): Promise<Uint8Array> {
|
||||
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<string, PortalState>,
|
||||
executeQuery: (query: string) => Promise<RawOutputEnvelope>,
|
||||
): Promise<Uint8Array> {
|
||||
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<string, PreparedStatementState>,
|
||||
portals: Map<string, PortalState>,
|
||||
): 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" ? "<unnamed portal>" : "<unnamed statement>";
|
||||
}
|
||||
return `"${name}"`;
|
||||
}
|
||||
|
||||
async function runQuery(
|
||||
workspaceId: string,
|
||||
datatableName: string,
|
||||
query: string,
|
||||
): Promise<Uint8Array> {
|
||||
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<RawOutputEnvelope> {
|
||||
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<string> {
|
||||
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<string, unknown>;
|
||||
const err = obj.error;
|
||||
if (err && typeof err === "object") {
|
||||
const errObj = err as Record<string, unknown>;
|
||||
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: [] };
|
||||
}
|
||||
@@ -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("'", "''")}'`;
|
||||
}
|
||||
@@ -6691,6 +6691,15 @@ datatable related commands
|
||||
- \`datatable run <sql:string>\` - run a SQL query on a datatable
|
||||
- \`-n --name <name:string>\` - 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:number>\` - Port to listen on (default: first free port in 5433-5500)
|
||||
- \`--host <host:string>\` - Bind address (default: 127.0.0.1)
|
||||
- \`--password <password:string>\` - 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 <name:string>\` - Datatable to connect psql to (default: main)
|
||||
- \`--port <port:number>\` - Port the proxy listens on (default: first free port in 5433-5500)
|
||||
- \`--host <host:string>\` - Bind address for the proxy (default: 127.0.0.1)
|
||||
- \`--password <password:string>\` - Password for the temporary Postgres proxy (default: generate a random password at startup)
|
||||
|
||||
### dependencies
|
||||
|
||||
|
||||
@@ -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 */]);
|
||||
});
|
||||
@@ -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");
|
||||
});
|
||||
@@ -77,6 +77,15 @@ datatable related commands
|
||||
- `datatable run <sql:string>` - run a SQL query on a datatable
|
||||
- `-n --name <name:string>` - 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:number>` - Port to listen on (default: first free port in 5433-5500)
|
||||
- `--host <host:string>` - Bind address (default: 127.0.0.1)
|
||||
- `--password <password:string>` - 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 <name:string>` - Datatable to connect psql to (default: main)
|
||||
- `--port <port:number>` - Port the proxy listens on (default: first free port in 5433-5500)
|
||||
- `--host <host:string>` - Bind address for the proxy (default: 127.0.0.1)
|
||||
- `--password <password:string>` - Password for the temporary Postgres proxy (default: generate a random password at startup)
|
||||
|
||||
### dependencies
|
||||
|
||||
|
||||
@@ -2608,6 +2608,15 @@ datatable related commands
|
||||
- \`datatable run <sql:string>\` - run a SQL query on a datatable
|
||||
- \`-n --name <name:string>\` - 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:number>\` - Port to listen on (default: first free port in 5433-5500)
|
||||
- \`--host <host:string>\` - Bind address (default: 127.0.0.1)
|
||||
- \`--password <password:string>\` - 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 <name:string>\` - Datatable to connect psql to (default: main)
|
||||
- \`--port <port:number>\` - Port the proxy listens on (default: first free port in 5433-5500)
|
||||
- \`--host <host:string>\` - Bind address for the proxy (default: 127.0.0.1)
|
||||
- \`--password <password:string>\` - Password for the temporary Postgres proxy (default: generate a random password at startup)
|
||||
|
||||
### dependencies
|
||||
|
||||
|
||||
@@ -82,6 +82,15 @@ datatable related commands
|
||||
- `datatable run <sql:string>` - run a SQL query on a datatable
|
||||
- `-n --name <name:string>` - 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:number>` - Port to listen on (default: first free port in 5433-5500)
|
||||
- `--host <host:string>` - Bind address (default: 127.0.0.1)
|
||||
- `--password <password:string>` - 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 <name:string>` - Datatable to connect psql to (default: main)
|
||||
- `--port <port:number>` - Port the proxy listens on (default: first free port in 5433-5500)
|
||||
- `--host <host:string>` - Bind address for the proxy (default: 127.0.0.1)
|
||||
- `--password <password:string>` - Password for the temporary Postgres proxy (default: generate a random password at startup)
|
||||
|
||||
### dependencies
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user