fix: parity gaps from 80 production projects (4.0.2)

CHANGELOG.md

@@ -4,6 +4,19 @@ Notable changes to `@tinbase/pg-mem`, the tinbase fork of pg-mem.
 
 Released from `main`, which carries the scoped package name. Upstream is tracked through the `upstream` remote (`oguimbal/pg-mem`) rather than a branch; the leftover `master` is vestigial.
 
+## 4.0.2
+
+Gaps found by running 80 production projects' migrations and seeds through the agent's validator on pg-mem and on PGlite; all 80 now match. Requires `@tinbase/pgsql-ast-parser` 12.2.1.
+
+- **Wrong results, fixed:** an aggregate inside a select-list subquery aggregated the *outer* query (`select id, (select count(*) from c where c.p = p.id) from p` returned one row); `agg(x ORDER BY y)` ignored its ORDER BY; `json_agg` / `jsonb_agg` dropped NULL inputs (postgres keeps them as json null).
+- `f(unnest(arr))` calls `f` per element (`lower(unnest(emails))` resolved to the range overload); `RETURNS SETOF <scalar>` SQL functions.
+- A set-returning call in a subquery's FROM reads the outer row (`from jsonb_array_elements(outer.col)`).
+- `UPDATE t alias` / `DELETE FROM t alias`: the alias names the table, in the statement and its subqueries.
+- plpgsql: a `FOR rec IN <query>` record's fields are usable in SQL statements inside the loop (`update … set x = rec.col`), including nested loops.
+- `json -> key` / `->>` with a column, variable or expression key.
+- `ALTER COLUMN … TYPE` and `RENAME COLUMN` keep the column's position.
+- `INSERT … VALUES` checks each row against the target columns first ("INSERT has more expressions than target columns").
+
 ## 4.0.1
 
 - Comparisons across type categories fail with Postgres' own error, `operator does not exist: uuid = text` (was `cannot cast type text to uuid`), types spelled as Postgres spells them. Also for `x IN (select col ...)`. The SQL that is rejected does not change.

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tinbase/pg-mem",
-  "version": "4.0.1",
+  "version": "4.0.2",
   "description": "Fork of pg-mem with extended Postgres conformance (PL/pgSQL, triggers, RLS, correlated subqueries, MERGE, ranges, full-text, partitioning, ...). Tracks oguimbal/pg-mem; pending upstream PR #476.",
   "main": "index.js",
   "publishConfig": {
@@ -49,7 +49,7 @@
     "immutable": "^4.3.4",
     "lru-cache": "^6.0.0",
     "object-hash": "^2.0.3",
-    "pgsql-ast-parser": "npm:@tinbase/pgsql-ast-parser@^12.2.0"
+    "pgsql-ast-parser": "npm:@tinbase/pgsql-ast-parser@^12.2.1"
   },
   "devDependencies": {
     "@babel/cli": "^7.10.5",

src/column.ts

@@ -221,9 +221,8 @@ export class ColRef implements _Column {
         const nn = newId;
         this.expression = columnEvaluator(this.table, newId, newType);
 
-        // replace in table
-        this.table.columnMgr.delete(on);
-        this.table.columnMgr.set(nn, this);
+        // replace in table, keeping the column's position (postgres never moves a column)
+        this.table.columnMgr.replace(on, nn, this);
     }
 
     drop(t: _Transaction): void {

src/execution/plpgsql.ts

@@ -541,6 +541,10 @@ interface Helpers {
     prepareSql(sql: string): (ctx: RunCtx) => any;
     /** run a dynamic (runtime-built) SQL string */
     runDynamic(ctx: RunCtx, sqlText: string): any;
+    /** the same toolbox with extra variables in scope for embedded SQL */
+    withVars(extra: { name: string; value: IValue }[]): Helpers;
+    /** rewrites embedded SQL for the enclosing FOR-loop records (rec.col -> variable) */
+    rewriteSql?: (sql: string) => string;
 }
 
 /** does this block (recursively) use RETURN NEXT / RETURN QUERY? => set-returning */
@@ -610,6 +614,7 @@ function makeHelpers(schema: _ISchema, params: Parameter[], returns: _IType | ni
     };
     const outColumns = ((returns as any)?.of?.columns as { name: string }[] | undefined) ?? null;
     return {
+        withVars: (extra) => makeHelpers(schema, [...params, ...extra.map((x, i) => ({ index: params.length + i, value: x.value }))], returns, setof),
         returns,
         setof,
         compileExpr,
@@ -864,6 +869,16 @@ function expandRowTypes(toks: string[], schema: _ISchema): string[] {
     return out;
 }
 
+/** `rec.col` -> fieldVar(col) in SQL text, outside string literals and quoted identifiers, for known columns */
+function rewriteRecordFields(sql: string, rec: string, known: Set<string>, fieldVar: (col: string) => string): string {
+    const re = new RegExp(`(^|[^\\w$."])${rec.replace(/[^\w]/g, '')}\\s*\\.\\s*([a-zA-Z_][\\w$]*)`, 'gi');
+    // split on '…' and "…" so literals and quoted identifiers are left alone
+    return sql.split(/('(?:[^']|'')*'|"(?:[^"]|"")*")/).map((part, i) => i % 2
+        ? part
+        : part.replace(re, (m, pre, col) => known.has(col.toLowerCase()) ? pre + fieldVar(col.toLowerCase()) : m)
+    ).join('');
+}
+
 function parseTypeDef(typeSrc: string): any {
     const e = parseExpr(`null::${typeSrc}`);
     if (e.type !== 'cast') {
@@ -1007,7 +1022,27 @@ function compileBody(stmts: GStmt[], h: Helpers): GCompiled[] {
                 // (aliased) selection; the current row is fed to the body's expressions
                 const sel = h.querySelection(s.sql).setAlias(s.varName);
                 const holder: { row: any } = { row: DUMMY_ROW };
-                const body = compileBody(s.body, { ...h, compileExpr: h.mkCompiler(sel, () => holder.row) });
+                // embedded SQL in the body (`update … set x = rec.col`) sees `rec.col` as a table-qualified
+                // column, not a variable: expose each field as a variable and rewrite the references to it
+                const rec = s.varName.toLowerCase();
+                const fields = sel.columns.filter(c => !!c.id);
+                const fieldVar = (col: string) => `__rec_${rec}__${col}`;
+                const withFields = h.withVars(fields.map(c => ({
+                    name: fieldVar(c.id!),
+                    value: new Evaluator(c.type, fieldVar(c.id!), `plpgsql_rec_${rec}_${c.id}`, [],
+                        (_raw, t) => c.get(holder.row, t), { forceNotConstant: true }),
+                })));
+                const known = new Set(fields.map(c => c.id!.toLowerCase()));
+                // enclosing loops' records first, so a nested loop still sees them
+                const rewrite = (sql: string) => rewriteRecordFields(h.rewriteSql ? h.rewriteSql(sql) : sql, rec, known, fieldVar);
+                const body = compileBody(s.body, {
+                    ...withFields,
+                    compileExpr: h.mkCompiler(sel, () => holder.row),
+                    mkCompiler: withFields.mkCompiler,
+                    prepareSql: (sql) => withFields.prepareSql(rewrite(sql)),
+                    querySelection: (sql) => withFields.querySelection(rewrite(sql)),
+                    rewriteSql: rewrite,
+                });
                 return {
                     run(ctx) {
                         for (const row of sel.enumerate(ctx.t)) {

src/execution/records-mutations/deletion.ts

@@ -12,7 +12,7 @@ export class Deletion extends MutationDataSourceBase {
         const { schema } = buildCtx();
         const table = asTable(schema.getObject(ast.from));
         // row-level security: DELETE can only affect rows visible via DELETE policies
-        const mutatedSel = applyReadRls(table, table.selection, 'delete', statementReadsColumns(ast as any))
+        const mutatedSel = applyReadRls(table, ast.from.alias ? table.selection.setAlias(ast.from.alias) : table.selection, 'delete', statementReadsColumns(ast as any))
             .filter(ast.where);
 
         super(table, mutatedSel, ast);

src/execution/records-mutations/insert.ts

@@ -30,6 +30,21 @@ export class Insert extends MutationDataSourceBase {
         // init super
         super(table, selection, ast);
 
+        // postgres checks each VALUES row against the target columns before the rows against
+        // each other, so a seed row with one value too many says so (not "VALUES lists must
+        // all be the same length")
+        if (ast.insert.type === 'values') {
+            const target = ast.columns?.length ?? table.selection.columns.length;
+            for (const row of ast.insert.values) {
+                if (row.length > target) {
+                    throw new QueryError('INSERT has more expressions than target columns', '42601');
+                }
+                if (ast.columns && row.length < target) {
+                    throw new QueryError('INSERT has more target columns than expressions', '42601');
+                }
+            }
+        }
+
         // get data to insert
         this.valueRawSource = ast.insert.type === 'values'
             ? buildValues(ast.insert, true)

src/execution/records-mutations/update.ts

@@ -63,7 +63,9 @@ export class Update extends MutationDataSourceBase {
 
             //  => REGULAR UPDATE
             // row-level security: UPDATE can only affect rows visible via UPDATE policies
-            mutatedSel = applyReadRls(into, into.selection, 'update', statementReadsColumns(ast as any))
+            // `update t alias set ... where alias.col ...`: the alias names the table, as in FROM
+            const base = ast.table.alias ? into.selection.setAlias(ast.table.alias) : into.selection;
+            mutatedSel = applyReadRls(into, base, 'update', statementReadsColumns(ast as any))
                 .filter(ast.where);
         }
 

src/execution/schema-amends/create-function.ts

@@ -94,6 +94,7 @@ export class CreateFunction extends ExecHelper implements _IStatementExecutor {
             args,
             code: fn.code,
             returns,
+            setof: !!fn.setof,
             functionName: fn.name.name,
             schema: schema,
         });

src/execution/select.ts

@@ -160,7 +160,10 @@ function buildRawSelectSubject(p: SelectFromStatement): _ISelection | nil {
                 const fnName = from.alias?.name ?? from.function?.name;
                 // built against the left FROM items: function calls referencing them
                 // are implicitly lateral in postgres
-                const fromValue: IValue = sel ? withSelection(sel, () => buildValue(from)) : buildValue(from);
+                // a first FROM item reads no row of its own: build it against the empty selection,
+                // so `from f(outer.col)` in a subquery resolves outer.col as a correlated reference
+                // rather than as a column of whatever selection happens to be current
+                const fromValue: IValue = withSelection(sel ?? buildCtx().schema.dualTable.selection, () => buildValue(from));
                 if ((from.lateral || !fromValue.isConstant) && sel) {
                     if (!ArrayType.matches(fromValue.type) || RecordType.matches(fromValue.type.of)) {
                         throw new NotSupported('lateral function calls returning records');

src/interfaces.ts

@@ -195,6 +195,8 @@ export interface ToCompile {
     args: ArgDefDetails[];
     /** Expected return type (if any) */
     returns?: IType | nil;
+    /** RETURNS SETOF <type>: `returns` is then that type's array, one element per row */
+    setof?: boolean;
 }
 
 export class AdvancedResult {

src/parser/function-call.ts

@@ -7,10 +7,17 @@ import hash from 'object-hash';
 import { parseArrayLiteral, QName } from 'pgsql-ast-parser';
 import { asSingleQName, nullIsh, qnameToStr } from '../utils';
 import { buildCtx } from './context';
-import { markSetReturning } from '../transforms/expand-srf';
+import { markSetReturning, isSetReturning } from '../transforms/expand-srf';
 
 
 export function buildCall(name: string | QName, args: IValue[]): IValue {
+    // f(unnest(arr)): postgres calls f once per element and the result is itself set-returning
+    // (`select lower(unnest(emails))`). Applying f to the whole array instead resolved
+    // lower(text[]) to the range overload.
+    const srfAt = args.findIndex(a => isSetReturning(a) && a.type instanceof ArrayType);
+    if (srfAt >= 0) {
+        return buildCallPerElement(name, args, srfAt);
+    }
     let type: _IType | nil = null;
     let get: (...args: any[]) => any;
 
@@ -312,6 +319,30 @@ export function buildCall(name: string | QName, args: IValue[]): IValue {
 }
 
 
+function buildCallPerElement(name: string | QName, args: IValue[], srfAt: number): IValue {
+    const srf = args[srfAt];
+    let current: any;
+    // stands for "the element being processed"; impure so it is never folded into a constant
+    const element = new Evaluator((srf.type as ArrayType).of, null, hash({ srfElement: srf.hash }), srf, () => current, { unpure: true });
+    const perElement = buildCall(name, args.map((a, i) => i === srfAt ? element : a));
+    const ret = new Evaluator(
+        perElement.type.asArray()
+        , null
+        , hash({ call: name, perElement: perElement.hash })
+        , args
+        , (raw, t) => {
+            const arr = srf.get(raw, t);
+            if (nullIsh(arr)) {
+                return null;
+            }
+            return (arr as any[]).map(x => {
+                current = x;
+                return perElement.get(raw, t);
+            });
+        }, { unpure: true });
+    return markSetReturning(ret);
+}
+
 function expectArgs(name: string | QName, args: IValue[], count: number | [number, number]) {
     const [min, max] = typeof count === 'number' ? [count, count] : count;
     if (args.length < min || args.length > max) {

src/schema/function-call-table.ts

@@ -25,7 +25,12 @@ export class FunctionCallTable extends DataSourceBase {
     }
 
     enumerate(t: _Transaction): Iterable<any> {
-        const results = this.evaluator.get(null, t);
+        // nothing precedes this FROM item, so the arguments read no row: they are constants, or
+        // outer-query references in a correlated subquery (`from jsonb_array_elements(outer.col)`),
+        // which are bound per outer row before this runs - evaluated now, not as constants
+        const results = this.evaluator.isConstant
+            ? this.evaluator.get()
+            : this.evaluator.get({}, t);
         for (const result of results ?? []) {
             result[this.symbol] = true;
         }
@@ -50,10 +55,14 @@ export class FunctionCallTable extends DataSourceBase {
 
 
     explain(e: _Explainer): _SelectExplanation {
-        throw new Error('Method not implemented.');
+        return {
+            id: e.idFor(this),
+            _: 'table',
+            table: 'function call',
+        } as any;
     }
 
     stats(t: _Transaction): Stats | null {
-        throw new Error('Method not implemented.');
+        return null;
     }
 }

src/schema/pg-catalog/binary-operators.ts

@@ -258,6 +258,33 @@ function jsonbHasKey(doc: any, key: string): boolean {
 }
 
 function registerJsonOperators(schema: _ISchema) {
+    // ======= "json -> key" / "json ->> key" with a non-literal key (a column, a variable,
+    // an expression); a literal key is compiled as a member access instead, same semantics
+    for (const j of [Types.json, Types.jsonb]) {
+        for (const [key, pick] of [
+            [Types.text(), (doc: any, k: string) => doc && typeof doc === 'object' && !Array.isArray(doc) ? doc[k] : undefined],
+            [Types.integer, (doc: any, i: number) => Array.isArray(doc) ? doc[i < 0 ? doc.length + i : i] : undefined],
+        ] as const) {
+            schema.registerOperator({
+                operator: '->',
+                left: j,
+                right: key,
+                returns: j,
+                implementation: (doc: any, k: any) => (pick as any)(doc, k) ?? null,
+            });
+            schema.registerOperator({
+                operator: '->>',
+                left: j,
+                right: key,
+                returns: Types.text(),
+                implementation: (doc: any, k: any) => {
+                    const v = (pick as any)(doc, k);
+                    return v === undefined || v === null ? null : typeof v === 'string' ? v : JSON.stringify(v);
+                },
+            });
+        }
+    }
+
     // ======= "jsonb ? text" / "?|" (any) / "?&" (all) key existence
     schema.registerOperator({
         operator: '?',

src/schema/pg-catalog/sql-function-language.ts

@@ -13,7 +13,7 @@ import { RecordType } from '../../datatypes/t-record';
 let execId = 0;
 
 export function registerSqlFunctionLanguage(db: _IDb) {
-    db.registerLanguage('sql', ({ code, schema: _schema, args, returns: _returns }) => {
+    db.registerLanguage('sql', ({ code, schema: _schema, args, returns: _returns, setof }) => {
         const schema = _schema as _ISchema;
         const returns = _returns as _IType;
         // parse SQL
@@ -69,6 +69,14 @@ export function registerSqlFunctionLanguage(db: _IDb) {
             transformResult = (v, t, eid) => v?.map(x => {
                 return transformItem(x, t, eid);
             });
+        } else if (setof && ArrayType.matches(returns)) {
+            // RETURNS SETOF <scalar>: one value per row
+            const cols = executor.selection.columns;
+            if (cols.length !== 1 || !cols[0].type.canConvertImplicit(returns.of)) {
+                throw new QueryError(`return type mismatch in function declared to return ${returns.of.name}`, '42P13');
+            }
+            const col = cols[0].cast(returns.of);
+            transformResult = (v, t) => (v ?? []).map(x => col.get(x, t));
         } else {
             // returns a single value
             const cols = executor.selection.columns;

src/table.ts

@@ -61,6 +61,20 @@ class ColumnManager {
         this.invalidateColumns();
         return this.map.delete(name);
     }
+
+    /** Puts `colDef` under `newName` where `oldName` was, keeping the column's position (rename, retype) */
+    replace(oldName: string, newName: string, colDef: ColRef) {
+        this.invalidateColumns();
+        const entries = Array.from(this.map.entries());
+        this.map.clear();
+        for (const [k, v] of entries) {
+            if (k === oldName) {
+                this.map.set(newName, colDef);
+            } else {
+                this.map.set(k, v);
+            }
+        }
+    }
 }
 
 export class MemoryTable extends DataSourceBase implements IMemoryTable<any>, _ITable {