diff --git a/postgres/query-builder/__tests__/query-builder.test.ts b/postgres/query-builder/__tests__/query-builder.test.ts index 3c01f69500..ae83b3b95f 100644 --- a/postgres/query-builder/__tests__/query-builder.test.ts +++ b/postgres/query-builder/__tests__/query-builder.test.ts @@ -1,4 +1,4 @@ -import { QueryBuilder } from '../src/query-builder'; +import { col, fn, lit, param, QueryBuilder } from '../src/query-builder'; // The deparser outputs pretty-printed SQL with newlines. // Use the `s` flag on regex so `.` matches newlines, or use [\s\S]. @@ -884,4 +884,116 @@ describe('QueryBuilder', () => { expect(values).toEqual(['inactive', '2024-01-01']); }); }); + + // ========================================================================= + // Expression system (col / lit / param / fn) + computed SELECT columns + // ========================================================================= + describe('Expressions', () => { + it('should pass a column reference as a function arg (not a param)', () => { + const { text, values } = new QueryBuilder() + .schema('priv') + .call('secrets_get', { secret_name: col('s.name'), namespace_id: col('s.namespace_id') }) + .build(); + + expect(text).toMatch(/priv\.secrets_get\(/); + expect(text).toMatch(/secret_name\s*=>\s*s\.name/); + expect(text).toMatch(/namespace_id\s*=>\s*s\.namespace_id/); + // Column refs are not parameterized. + expect(values).toEqual([]); + }); + + it('should emit a literal NULL via lit(null) (not a bound param)', () => { + const { text, values } = new QueryBuilder() + .call('f', { d: lit(null) }) + .build(); + + expect(text).toMatch(/d\s*=>\s*NULL/i); + expect(values).toEqual([]); + }); + + it('should bind an explicit param() and auto-bind a plain value', () => { + const { text, values } = new QueryBuilder() + .call('f', { a: param('x'), b: 'y' }) + .build(); + + expect(text).toMatch(/a\s*=>\s*\$1/); + expect(text).toMatch(/b\s*=>\s*\$2/); + expect(values).toEqual(['x', 'y']); + }); + + it('should mix column-ref, param, and literal args in one named call', () => { + const { text, values } = new QueryBuilder() + .call('f', { + owner_id: col('s.owner_id'), + secret_name: col('s.name'), + ns: param('11111111-1111-1111-1111-111111111111'), + d: lit(null), + }) + .build(); + + expect(text).toMatch(/owner_id\s*=>\s*s\.owner_id/); + expect(text).toMatch(/secret_name\s*=>\s*s\.name/); + expect(text).toMatch(/ns\s*=>\s*\$1/); + expect(text).toMatch(/d\s*=>\s*NULL/i); + expect(values).toEqual(['11111111-1111-1111-1111-111111111111']); + }); + + it('should support fn() with schema-qualified name', () => { + const { text } = new QueryBuilder() + .selectExpr('v', fn('priv.secrets_get', { secret_name: col('s.name') })) + .table('secrets', 's') + .schema('priv') + .build(); + + expect(text).toMatch(/priv\.secrets_get\(\s*secret_name\s*=>\s*s\.name\s*\)\s+AS\s+v/); + }); + + it('should build a computed function-call column over a base table (scope-switch shape)', () => { + // Mirrors decryptScopeSecrets for an entity scope. + const nsId = '22222222-2222-2222-2222-222222222222'; + const { text, values } = new QueryBuilder() + .schema('priv') + .table('secrets', 's') + .select(['s.name']) + .selectCall('decrypted_value', 'secrets_get', { + owner_id: col('s.owner_id'), + secret_name: col('s.name'), + namespace_id: col('s.namespace_id'), + default_value: lit(null), + }, { schema: 'priv' }) + .where('s.namespace_id', '=', nsId) + .where('s.owner_id', '=', 'owner-1') + .where('s.retired_at', 'IS', null) + .orderBy('s.created_at', 'ASC') + .build(); + + expect(text).toMatch(/SELECT/); + expect(text).toMatch(/s\.name/); + expect(text).toMatch(/priv\.secrets_get\(/); + expect(text).toMatch(/owner_id\s*=>\s*s\.owner_id/); + expect(text).toMatch(/default_value\s*=>\s*NULL/i); + expect(text).toMatch(/AS\s+decrypted_value/); + expect(text).toMatch(/FROM\s+priv\.secrets\s+(AS\s+)?s/); + expect(text).toMatch(/s\.namespace_id\s*=\s*\$1/); + expect(text).toMatch(/s\.owner_id\s*=\s*\$2/); + expect(text).toMatch(/s\.retired_at\s+IS\s+NULL/); + expect(text).toMatch(/ORDER\s+BY\s+s\.created_at/); + // Column-ref getter args bind nothing; only WHERE values are params. + expect(values).toEqual([nsId, 'owner-1']); + }); + + it('should keep param numbering correct when a computed column also binds a param', () => { + const { text, values } = new QueryBuilder() + .table('t', 's') + .select(['s.id']) + .selectExpr('v', fn('f', { a: param('A'), b: col('s.b') })) + .where('s.active', '=', true) + .build(); + + // Target-list param comes before WHERE param. + expect(text).toMatch(/a\s*=>\s*\$1/); + expect(text).toMatch(/s\.active\s*=\s*\$2/); + expect(values).toEqual(['A', true]); + }); + }); }); diff --git a/postgres/query-builder/src/index.ts b/postgres/query-builder/src/index.ts index 036333812a..b84e62999e 100644 --- a/postgres/query-builder/src/index.ts +++ b/postgres/query-builder/src/index.ts @@ -1,2 +1,11 @@ -export { QueryBuilder } from './query-builder'; -export type { QueryOutput, SqlValue } from './query-builder'; +export { QueryBuilder, col, lit, param, fn } from './query-builder'; +export type { + QueryOutput, + SqlValue, + Expr, + FnArg, + FnArgs, + ParamAllocator, + SelectExpr, + SelectItem, +} from './query-builder'; diff --git a/postgres/query-builder/src/query-builder.ts b/postgres/query-builder/src/query-builder.ts index 3c60d064bf..9ce5f0d18c 100644 --- a/postgres/query-builder/src/query-builder.ts +++ b/postgres/query-builder/src/query-builder.ts @@ -17,6 +17,13 @@ type SortDir = 'ASC' | 'DESC'; type NullsOrder = 'FIRST' | 'LAST'; type ConflictAction = 'nothing' | 'update'; +// A SELECT target: a column name string, or a computed expression with alias. +export interface SelectExpr { + expr: Expr; + as?: string; +} +export type SelectItem = string | SelectExpr; + interface OrderBySpec { column: string; direction: SortDir; @@ -180,6 +187,73 @@ function namedArgNode(name: string, val: unknown) { }); } +// --------------------------------------------------------------------------- +// Expression system +// +// Function-call arguments and computed SELECT columns are built lazily via +// thunks so that parameter numbering ($1, $2, ...) stays correct regardless of +// where the expression lands in the final statement. An `Expr` is a function +// that, given the builder's parameter allocator, returns a pg-ast node. +// --------------------------------------------------------------------------- + +// Allocates a positional parameter ($n) and returns its AST node. +export type ParamAllocator = (value: SqlValue) => unknown; + +// A deferred AST node. Produced by `col`, `lit`, `param`, `fn`. +export type Expr = (alloc: ParamAllocator) => unknown; + +// A function-call argument: an explicit expression, or a plain value (which is +// auto-parameterized as $n). +export type FnArg = Expr | SqlValue; +export type FnArgs = FnArg[] | Record; + +function isExpr(x: unknown): x is Expr { + return typeof x === 'function'; +} + +function argToNode(arg: FnArg, alloc: ParamAllocator): unknown { + return isExpr(arg) ? arg(alloc) : alloc(arg); +} + +function buildFnArgNodes(args: FnArgs | undefined, alloc: ParamAllocator): unknown[] { + if (!args) return []; + if (Array.isArray(args)) { + return args.map((a) => argToNode(a, alloc)); + } + return Object.entries(args).map(([name, a]) => + namedArgNode(name, argToNode(a, alloc)) + ); +} + +// Column reference, e.g. `col('s.owner_id')` or `col('name')`. +export function col(name: string): Expr { + return () => colRef(name); +} + +// SQL literal (not a bound parameter), e.g. `lit(null)` -> NULL, `lit(3)` -> 3. +export function lit(value: SqlValue): Expr { + return () => constNode(value); +} + +// Explicitly bound parameter ($n). +export function param(value: SqlValue): Expr { + return (alloc) => alloc(value); +} + +// Function/procedure call as an expression. `name` may be schema-qualified +// ('schema.fn') or the schema passed via opts. Args may be positional (array) +// or named (object); each arg is an Expr or a plain value. +export function fn(name: string, args?: FnArgs, opts?: { schema?: string }): Expr { + let schema = opts?.schema; + let fname = name; + if (!schema && name.includes('.')) { + const idx = name.lastIndexOf('.'); + schema = name.slice(0, idx); + fname = name.slice(idx + 1); + } + return (alloc) => funcCallNode(fname, buildFnArgNodes(args, alloc), schema); +} + // --------------------------------------------------------------------------- // QueryBuilder // --------------------------------------------------------------------------- @@ -188,7 +262,7 @@ export class QueryBuilder { private _schema: string | undefined; private _table: string | undefined; private _tableAlias: string | undefined; - private _columns: string[] = []; + private _selectItems: SelectItem[] = []; private _distinct: boolean = false; private _insertData: Record[] | undefined; private _updateData: Record | undefined; @@ -203,7 +277,7 @@ export class QueryBuilder { private _returning: string[] | undefined; private _ctes: CteSpec[] = []; private _onConflict: OnConflictSpec | undefined; - private _funcCall: { name: string; args?: SqlValue[] | Record; schema?: string } | undefined; + private _funcCall: { name: string; args?: FnArgs; schema?: string } | undefined; // ------------------------------------------------------------------------- // Schema / Table @@ -224,8 +298,21 @@ export class QueryBuilder { // SELECT // ------------------------------------------------------------------------- - select(columns: string[] = ['*']): this { - this._columns = columns; + select(columns: SelectItem[] = ['*']): this { + this._selectItems = columns; + return this; + } + + // Append a computed function-call column, e.g. + // .selectCall('decrypted_value', 'priv.secrets_get', { secret_name: col('s.name') }) + selectCall(as: string, name: string, args?: FnArgs, opts?: { schema?: string }): this { + this._selectItems.push({ expr: fn(name, args, opts), as }); + return this; + } + + // Append an arbitrary computed column expression with an alias. + selectExpr(as: string, expr: Expr): this { + this._selectItems.push({ expr, as }); return this; } @@ -374,8 +461,8 @@ export class QueryBuilder { // Function / Procedure call // ------------------------------------------------------------------------- - call(name: string, args?: SqlValue[] | Record): this { - this._funcCall = { name, args, schema: this._schema }; + call(name: string, args?: FnArgs, opts?: { schema?: string }): this { + this._funcCall = { name, args, schema: opts?.schema ?? this._schema }; return this; } @@ -424,6 +511,22 @@ export class QueryBuilder { return paramNode(this._paramIndex); } + private get _alloc(): ParamAllocator { + return (value: SqlValue) => this._addParam(value); + } + + private _buildTargetList(items: SelectItem[]): unknown[] { + const alloc = this._alloc; + return items.map((item) => { + if (typeof item === 'string') { + return item === '*' + ? resTargetVal(colRef('*')) + : resTargetVal(colRef(item)); + } + return resTargetVal(item.expr(alloc), item.as); + }); + } + // ------------------------------------------------------------------------- // Internal: WHERE clause builder // ------------------------------------------------------------------------- @@ -558,11 +661,7 @@ export class QueryBuilder { private _buildSelectInternal(startParamIndex: number = 0): unknown { this._paramIndex = startParamIndex; - const targetList = this._columns.map((col) => - col === '*' - ? resTargetVal(colRef('*')) - : resTargetVal(colRef(col)) - ); + const targetList = this._buildTargetList(this._selectItems); const whereClause = this._buildWhereNode(this._whereConditions); const havingClause = this._buildWhereNode(this._havingConditions); @@ -763,22 +862,12 @@ export class QueryBuilder { this._resetParams(); const fc = this._funcCall!; - let funcArgs: unknown[] = []; - - if (Array.isArray(fc.args)) { - funcArgs = fc.args.map((v) => this._addParam(v)); - } else if (fc.args && typeof fc.args === 'object') { - funcArgs = Object.entries(fc.args).map(([name, val]) => - namedArgNode(name, this._addParam(val)) - ); - } + const funcArgs = buildFnArgNodes(fc.args, this._alloc); const funcNode = funcCallNode(fc.name, funcArgs, fc.schema); - if (this._columns.length > 0) { - const targetList = this._columns.map((col) => - resTargetVal(colRef(col)) - ); + if (this._selectItems.length > 0) { + const targetList = this._buildTargetList(this._selectItems); const rangeFunc = nodes.rangeFunction({ functions: [nodes.list({ items: [funcNode] as any[] })] as any[],