From 19aa7ac9446c7143e0f08882c5398d22b735f065 Mon Sep 17 00:00:00 2001 From: Yury Selivanov Date: Mon, 27 Jan 2025 21:52:56 -0800 Subject: [PATCH] Introduce `withSQLRowMode(mode: 'array' | 'object'): Client` Use it to switch the return type of client.querySQL() from JS array (indexable by position) to JS object (indexable by field name). "Hide" `client.withCodecs({sql_row: ...})` by renaming "sql_row" to "_private_sql_row". If there's ever a use-case to make it possible for users to overload how SQL rows are decoded beyond choosing between 'array' and 'object' we can rename it back to `sql_row`. But I do want to keep the door open to us potentially implementing a different codec API for "collection" types -- an API that would have better performance but look more low level. See [1] for details. [1] https://github.com/edgedb/edgedb-js/pull/1187#discussion_r1931526979 --- packages/driver/src/baseClient.ts | 4 ++ packages/driver/src/codecs/codecs.ts | 4 +- packages/driver/src/codecs/context.ts | 1 + packages/driver/src/codecs/record.ts | 78 +++++++++++++++++++++------ packages/driver/src/options.ts | 31 +++++++++-- packages/driver/test/client.test.ts | 54 +++++++++++-------- 6 files changed, 127 insertions(+), 45 deletions(-) diff --git a/packages/driver/src/baseClient.ts b/packages/driver/src/baseClient.ts index 9d7c1326b..2ac04fb08 100644 --- a/packages/driver/src/baseClient.ts +++ b/packages/driver/src/baseClient.ts @@ -593,6 +593,10 @@ export class Client implements Executor { return new Client(this.pool, this.options.withCodecs(codecs)); } + withSQLRowMode(mode: "array" | "object"): Client { + return new Client(this.pool, this.options.withSQLRowMode(mode)); + } + withGlobals(globals: Record): Client { return new Client(this.pool, this.options.withGlobals(globals)); } diff --git a/packages/driver/src/codecs/codecs.ts b/packages/driver/src/codecs/codecs.ts index eb3ad85fc..09b8c79fc 100644 --- a/packages/driver/src/codecs/codecs.ts +++ b/packages/driver/src/codecs/codecs.ts @@ -160,11 +160,11 @@ export namespace Codecs { export type SQLRowCodec = { fromDatabase: (data: any[], desc: { names: string[] }) => any; - toDatabase: (data: never) => never; + toDatabase: (data: any, ...extras: any[]) => any; }; export type ContainerCodecs = { - sql_row: SQLRowCodec; + _private_sql_row: SQLRowCodec; }; export type KnownCodecs = ScalarCodecs & ContainerCodecs; diff --git a/packages/driver/src/codecs/context.ts b/packages/driver/src/codecs/context.ts index 0f7b213c7..cd38e0e9f 100644 --- a/packages/driver/src/codecs/context.ts +++ b/packages/driver/src/codecs/context.ts @@ -17,6 +17,7 @@ const NOOP: Codecs.AnyCodec = { }, }; +export type MutableCodecMap = Map; export type ReadonlyCodecMap = ReadonlyMap; type ContainerNames = keyof Codecs.ContainerCodecs; diff --git a/packages/driver/src/codecs/record.ts b/packages/driver/src/codecs/record.ts index 555934e18..fe22f09e2 100644 --- a/packages/driver/src/codecs/record.ts +++ b/packages/driver/src/codecs/record.ts @@ -18,11 +18,38 @@ import type { ICodec, uuid, CodecKind } from "./ifaces"; import { Codec } from "./ifaces"; +import type { Codecs } from "./codecs"; import type { WriteBuffer } from "../primitives/buffer"; import { ReadBuffer } from "../primitives/buffer"; import { InvalidArgumentError, ProtocolError } from "../errors"; import type { CodecContext } from "./context"; +const SQLRowArrayCodec: Codecs.SQLRowCodec = { + fromDatabase(values: any[], _desc: { names: string[] }): any[] { + return values; + }, + toDatabase() { + throw "cannot encode SQL record as a query argument"; + }, +}; + +const SQLRowObjectCodec: Codecs.SQLRowCodec = { + fromDatabase(values: any[], { names }: { names: string[] }): any { + return Object.fromEntries(names.map((key, index) => [key, values[index]])); + }, + toDatabase() { + throw "cannot encode SQL record as a query argument"; + }, +}; + +export const SQLRowModeArray = { + _private_sql_row: SQLRowArrayCodec, +}; + +export const SQLRowModeObject = { + _private_sql_row: SQLRowObjectCodec, +}; + export class RecordCodec extends Codec implements ICodec { private subCodecs: ICodec[]; private names: string[]; @@ -48,25 +75,44 @@ export class RecordCodec extends Codec implements ICodec { } const elemBuf = ReadBuffer.alloc(); - const result: any[] = new Array(els); - for (let i = 0; i < els; i++) { - buf.discard(4); // reserved - const elemLen = buf.readInt32(); - let val = null; - if (elemLen !== -1) { - buf.sliceInto(elemBuf, elemLen); - val = subCodecs[i].decode(elemBuf, ctx); - elemBuf.finish(); + const overload = ctx.getContainerOverload("_private_sql_row"); + + if (overload !== SQLRowObjectCodec) { + const result: any[] = new Array(els); + for (let i = 0; i < els; i++) { + buf.discard(4); // reserved + const elemLen = buf.readInt32(); + let val = null; + if (elemLen !== -1) { + buf.sliceInto(elemBuf, elemLen); + val = subCodecs[i].decode(elemBuf, ctx); + elemBuf.finish(); + } + result[i] = val; } - result[i] = val; - } - const overload = ctx.getContainerOverload("sql_row"); - if (overload != null) { - return overload.fromDatabase(result, { names: this.names }); - } + if (overload != null) { + return overload.fromDatabase(result, { names: this.names }); + } - return result; + return result; + } else { + // Parsing code is basically duplicated to make things fast. + const names = this.names; + const result: Record = {}; + for (let i = 0; i < els; i++) { + buf.discard(4); // reserved + const elemLen = buf.readInt32(); + let val = null; + if (elemLen !== -1) { + buf.sliceInto(elemBuf, elemLen); + val = subCodecs[i].decode(elemBuf, ctx); + elemBuf.finish(); + } + result[names[i]] = val; + } + return result; + } } getSubcodecs(): ICodec[] { diff --git a/packages/driver/src/options.ts b/packages/driver/src/options.ts index eb9340dda..bf1226de4 100644 --- a/packages/driver/src/options.ts +++ b/packages/driver/src/options.ts @@ -2,7 +2,8 @@ import * as errors from "./errors/index"; import { utf8Encoder } from "./primitives/buffer"; import type { Mutable } from "./typeutil"; import type { Codecs } from "./codecs/codecs"; -import type { ReadonlyCodecMap } from "./codecs/context"; +import { SQLRowModeObject } from "./codecs/record"; +import type { ReadonlyCodecMap, MutableCodecMap } from "./codecs/context"; import { CodecContext, NOOP_CODEC_CONTEXT } from "./codecs/context"; export type BackoffFunction = (n: number) => number; @@ -145,7 +146,7 @@ export interface CodecSpec { decode: (data: any) => any; } -export interface OptionsList { +export type OptionsList = { module?: string; moduleAliases?: Record; config?: Record; @@ -154,7 +155,11 @@ export interface OptionsList { transactionOptions?: TransactionOptions; warningHandler?: WarningHandler; codecs?: Codecs.CodecSpec; -} +}; + +type MergeOptions = OptionsList & { + _dropSQLRowCodec?: boolean; +}; export class Options { // This type is immutable. @@ -227,7 +232,7 @@ export class Options { return ctx; } - private _cloneWith(mergeOptions: OptionsList) { + private _cloneWith(mergeOptions: MergeOptions) { const clone: Mutable = Object.create(Options.prototype); clone.annotations = this.annotations; @@ -273,6 +278,14 @@ export class Options { clone.codecs = this.codecs; } + if (mergeOptions._dropSQLRowCodec && clone.codecs.has("sql_row")) { + // This is an optimization -- if "sql_row" is the only codec defined + // and it's set to "array mode", the we want the codec mapping to be + // empty instead. Why? Empty codec mapping short circuits a lot of + // custom codec code, and array is the default behavior anyway. + (clone.codecs as MutableCodecMap).delete("sql_row"); + } + clone.module = mergeOptions.module ?? this.module; return clone as Options; @@ -320,6 +333,16 @@ export class Options { return this._cloneWith({ codecs }); } + withSQLRowMode(mode: "array" | "object"): Options { + if (mode === "array") { + return this._cloneWith({ _dropSQLRowCodec: true }); + } else if (mode === "object") { + return this._cloneWith({ codecs: SQLRowModeObject }); + } else { + throw new errors.InterfaceError(`invalid mode=${mode}`); + } + } + withGlobals(globals: Record): Options { return this._cloneWith({ globals: { ...this.globals, ...globals }, diff --git a/packages/driver/test/client.test.ts b/packages/driver/test/client.test.ts index 4099a398e..37f40d3d9 100644 --- a/packages/driver/test/client.test.ts +++ b/packages/driver/test/client.test.ts @@ -2609,30 +2609,38 @@ if (getEdgeDBVersion().major >= 6) { let res = await client.querySQL("select 1"); expect(JSON.stringify(res)).toEqual("[[1]]"); - res = await client - .withCodecs({ - sql_row: { - fromDatabase(data, desc) { - return Object.fromEntries( - desc.names.map((key, index) => [key, data[index]]), - ); - }, - toDatabase() { - // toDatabase is required to be specified (limitation - // of TypeScript type system). This isn't super elegant - // and most likely we'll need a higher-level nicer API - // in top of this. - // - // Maybe `client.withCodecs(gel.SQLRowAsObject)`? - throw "cannot encode SQL record as a query argument"; + for (let i = 0; i < 2; i++) { + res = await client + .withCodecs({ + _private_sql_row: { + fromDatabase(data, desc) { + const ret = Object.fromEntries( + desc.names.map((key, index) => [key, data[index]]), + ); + ret.added = true; + return ret; + }, + toDatabase() { + throw "cannot encode SQL record as a query argument"; + }, }, - }, - }) - .querySQL("select 1 AS foo, 2 AS bar"); - expect(JSON.stringify(res)).toEqual('[{"foo":1,"bar":2}]'); - - res = await client.querySQL("select 1 + $1::int8", [41]); - expect(JSON.stringify(res)).toEqual("[[42]]"); + }) + .querySQL("select 1 AS foo, 2 AS bar"); + expect(JSON.stringify(res)).toEqual('[{"foo":1,"bar":2,"added":true}]'); + + res = await client + .withSQLRowMode("array") + .querySQL("select 1 AS foo, 2 AS bar"); + expect(JSON.stringify(res)).toEqual("[[1,2]]"); + + res = await client + .withSQLRowMode("object") + .querySQL("select 1 AS foo, 2 AS bar"); + expect(JSON.stringify(res)).toEqual('[{"foo":1,"bar":2}]'); + + res = await client.querySQL("select 1 + $1::int8", [41]); + expect(JSON.stringify(res)).toEqual("[[42]]"); + } } finally { await client.close(); }