From ff04440593f742f89852cf45d2efacfc3c91e7ad Mon Sep 17 00:00:00 2001 From: James Clarke Date: Tue, 26 Nov 2024 14:30:33 +0000 Subject: [PATCH 1/6] Add support for query tagging --- packages/driver/src/baseClient.ts | 7 ++++++ packages/driver/src/baseConn.ts | 8 +++++-- packages/driver/src/options.ts | 39 ++++++++++++++++++++++++++----- 3 files changed, 46 insertions(+), 8 deletions(-) diff --git a/packages/driver/src/baseClient.ts b/packages/driver/src/baseClient.ts index b25ae7e91..d4d9217c9 100644 --- a/packages/driver/src/baseClient.ts +++ b/packages/driver/src/baseClient.ts @@ -604,6 +604,13 @@ export class Client implements Executor { ); } + withTag(tag: string): Client { + return new Client( + this.pool, + this.options.withSession(this.options.session.withTag(tag)), + ); + } + withWarningHandler(handler: WarningHandler): Client { return new Client(this.pool, this.options.withWarningHandler(handler)); } diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts index 730c91e8c..261ac112c 100644 --- a/packages/driver/src/baseConn.ts +++ b/packages/driver/src/baseConn.ts @@ -899,6 +899,12 @@ export class BaseRawConnection { options: QueryOptions | undefined, language: Language, ) { + const annotations = Object.entries(state.annotations); + wb.writeUInt16(annotations.length); + for (const [name, value] of annotations) { + wb.writeString(name); + wb.writeString(value); + } wb.writeFlags(0xffff_ffff, capabilitiesFlags); wb.writeFlags( 0, @@ -953,7 +959,6 @@ export class BaseRawConnection { ): Promise { const wb = new WriteMessageBuffer(); wb.beginMessage(chars.$P); - wb.writeUInt16(0); // no headers this._encodeParseParams( wb, @@ -1080,7 +1085,6 @@ export class BaseRawConnection { ): Promise { const wb = new WriteMessageBuffer(); wb.beginMessage(chars.$O); - wb.writeUInt16(0); // no headers this._encodeParseParams( wb, diff --git a/packages/driver/src/options.ts b/packages/driver/src/options.ts index 08b0a3201..979e88f30 100644 --- a/packages/driver/src/options.ts +++ b/packages/driver/src/options.ts @@ -1,4 +1,5 @@ import * as errors from "./errors"; +import { utf8Encoder } from "./primitives/buffer"; export type BackoffFunction = (n: number) => number; @@ -118,6 +119,8 @@ export class TransactionOptions { } } +const TAG_ANNOTATION_KEY = "tag"; + export interface SessionOptions { module?: string; moduleAliases?: Record; @@ -138,6 +141,9 @@ export class Session { readonly config: Record; readonly globals: Record; + /** @internal */ + annotations: Record = {}; + constructor({ module = "default", moduleAliases = {}, @@ -150,33 +156,54 @@ export class Session { this.globals = globals; } + private _clone(mergeOptions: SessionOptions) { + const session = new Session({ ...this, ...mergeOptions }); + session.annotations = this.annotations; + return session; + } + withModuleAliases({ module, ...aliases }: { [name: string]: string; }): Session { - return new Session({ - ...this, + return this._clone({ module: module ?? this.module, moduleAliases: { ...this.moduleAliases, ...aliases }, }); } withConfig(config: { [name: string]: any }): Session { - return new Session({ - ...this, + return this._clone({ config: { ...this.config, ...config }, }); } withGlobals(globals: { [name: string]: any }): Session { - return new Session({ - ...this, + return this._clone({ globals: { ...this.globals, ...globals }, }); } + withTag(tag: string): Session { + if (tag.startsWith("edgedb/")) { + throw new errors.InterfaceError("reserved tag: edgedb/*"); + } + if (tag.startsWith("gel/")) { + throw new errors.InterfaceError("reserved tag: gel/*"); + } + if (utf8Encoder.encode(tag).length > 128) { + throw new errors.InterfaceError("tag too long (> 128 bytes)"); + } + const session = new Session({ ...this }); + session.annotations = { + ...this.annotations, + [TAG_ANNOTATION_KEY]: tag, + }; + return session; + } + /** @internal */ _serialise() { const state: SerializedSessionState = {}; From dd8edc19ac3b5ae7d905ea1001737b11e530f5fc Mon Sep 17 00:00:00 2001 From: James Clarke Date: Wed, 27 Nov 2024 22:33:57 +0000 Subject: [PATCH 2/6] Make it possible to unset tag --- packages/driver/src/baseClient.ts | 2 +- packages/driver/src/baseConn.ts | 5 ++--- packages/driver/src/options.ts | 36 ++++++++++++++++++------------- 3 files changed, 24 insertions(+), 19 deletions(-) diff --git a/packages/driver/src/baseClient.ts b/packages/driver/src/baseClient.ts index d4d9217c9..c1ad4c00d 100644 --- a/packages/driver/src/baseClient.ts +++ b/packages/driver/src/baseClient.ts @@ -604,7 +604,7 @@ export class Client implements Executor { ); } - withTag(tag: string): Client { + withTag(tag: string | null): Client { return new Client( this.pool, this.options.withSession(this.options.session.withTag(tag)), diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts index 261ac112c..fdd303cdc 100644 --- a/packages/driver/src/baseConn.ts +++ b/packages/driver/src/baseConn.ts @@ -899,9 +899,8 @@ export class BaseRawConnection { options: QueryOptions | undefined, language: Language, ) { - const annotations = Object.entries(state.annotations); - wb.writeUInt16(annotations.length); - for (const [name, value] of annotations) { + wb.writeUInt16(state.annotations.size); + for (const [name, value] of state.annotations) { wb.writeString(name); wb.writeString(value); } diff --git a/packages/driver/src/options.ts b/packages/driver/src/options.ts index 979e88f30..9b9e55b35 100644 --- a/packages/driver/src/options.ts +++ b/packages/driver/src/options.ts @@ -142,7 +142,11 @@ export class Session { readonly globals: Record; /** @internal */ - annotations: Record = {}; + annotations = new Map(); + + get tag(): string | null { + return this.annotations.get(TAG_ANNOTATION_KEY) ?? null; + } constructor({ module = "default", @@ -186,21 +190,23 @@ export class Session { }); } - withTag(tag: string): Session { - if (tag.startsWith("edgedb/")) { - throw new errors.InterfaceError("reserved tag: edgedb/*"); - } - if (tag.startsWith("gel/")) { - throw new errors.InterfaceError("reserved tag: gel/*"); - } - if (utf8Encoder.encode(tag).length > 128) { - throw new errors.InterfaceError("tag too long (> 128 bytes)"); - } + withTag(tag: string | null): Session { const session = new Session({ ...this }); - session.annotations = { - ...this.annotations, - [TAG_ANNOTATION_KEY]: tag, - }; + session.annotations = new Map(this.annotations); + if (tag != null) { + if (tag.startsWith("edgedb/")) { + throw new errors.InterfaceError("reserved tag: edgedb/*"); + } + if (tag.startsWith("gel/")) { + throw new errors.InterfaceError("reserved tag: gel/*"); + } + if (utf8Encoder.encode(tag).length > 128) { + throw new errors.InterfaceError("tag too long (> 128 bytes)"); + } + session.annotations.set(TAG_ANNOTATION_KEY, tag); + } else { + session.annotations.delete(TAG_ANNOTATION_KEY); + } return session; } From 421b064fa2aa47a5d536a73af676670510a6271f Mon Sep 17 00:00:00 2001 From: James Clarke Date: Wed, 27 Nov 2024 22:38:32 +0000 Subject: [PATCH 3/6] Check annotation count --- packages/driver/src/baseConn.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts index fdd303cdc..582d9fcc4 100644 --- a/packages/driver/src/baseConn.ts +++ b/packages/driver/src/baseConn.ts @@ -899,6 +899,9 @@ export class BaseRawConnection { options: QueryOptions | undefined, language: Language, ) { + if (state.annotations.size >= 1 << 16) { + throw new errors.InternalClientError("too many annotations"); + } wb.writeUInt16(state.annotations.size); for (const [name, value] of state.annotations) { wb.writeString(name); From 47b899aaa36956ea62530a5d47deb9c290104799 Mon Sep 17 00:00:00 2001 From: James Clarke Date: Wed, 27 Nov 2024 22:46:47 +0000 Subject: [PATCH 4/6] Only send annotations for protocol version >= 3 --- packages/driver/src/baseConn.ts | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/packages/driver/src/baseConn.ts b/packages/driver/src/baseConn.ts index 582d9fcc4..bd2581fc5 100644 --- a/packages/driver/src/baseConn.ts +++ b/packages/driver/src/baseConn.ts @@ -899,13 +899,17 @@ export class BaseRawConnection { options: QueryOptions | undefined, language: Language, ) { - if (state.annotations.size >= 1 << 16) { - throw new errors.InternalClientError("too many annotations"); - } - wb.writeUInt16(state.annotations.size); - for (const [name, value] of state.annotations) { - wb.writeString(name); - wb.writeString(value); + if (versionGreaterThanOrEqual(this.protocolVersion, [3, 0])) { + if (state.annotations.size >= 1 << 16) { + throw new errors.InternalClientError("too many annotations"); + } + wb.writeUInt16(state.annotations.size); + for (const [name, value] of state.annotations) { + wb.writeString(name); + wb.writeString(value); + } + } else { + wb.writeUInt16(0); } wb.writeFlags(0xffff_ffff, capabilitiesFlags); wb.writeFlags( From c0adda60ae9383a7d48902d315ce27f7aeb3feee Mon Sep 17 00:00:00 2001 From: James Clarke Date: Thu, 5 Dec 2024 21:41:42 +0000 Subject: [PATCH 5/6] Rename method to `withQueryTag` --- packages/driver/src/baseClient.ts | 4 ++-- packages/driver/src/options.ts | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/driver/src/baseClient.ts b/packages/driver/src/baseClient.ts index c1ad4c00d..a06dfff43 100644 --- a/packages/driver/src/baseClient.ts +++ b/packages/driver/src/baseClient.ts @@ -604,10 +604,10 @@ export class Client implements Executor { ); } - withTag(tag: string | null): Client { + withQueryTag(tag: string | null): Client { return new Client( this.pool, - this.options.withSession(this.options.session.withTag(tag)), + this.options.withSession(this.options.session.withQueryTag(tag)), ); } diff --git a/packages/driver/src/options.ts b/packages/driver/src/options.ts index 9b9e55b35..a2263c224 100644 --- a/packages/driver/src/options.ts +++ b/packages/driver/src/options.ts @@ -190,7 +190,7 @@ export class Session { }); } - withTag(tag: string | null): Session { + withQueryTag(tag: string | null): Session { const session = new Session({ ...this }); session.annotations = new Map(this.annotations); if (tag != null) { From 7be49e52f467031afa5b47db360824c517f4e836 Mon Sep 17 00:00:00 2001 From: James Clarke Date: Thu, 5 Dec 2024 22:24:27 +0000 Subject: [PATCH 6/6] Fix browser tests --- packages/driver/test/browser.test.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/driver/test/browser.test.ts b/packages/driver/test/browser.test.ts index e9b110c57..ad39edcb3 100644 --- a/packages/driver/test/browser.test.ts +++ b/packages/driver/test/browser.test.ts @@ -73,9 +73,11 @@ import { const brokenConnectOpts = JSON.parse( process.env._JEST_EDGEDB_CONNECT_CONFIG || "", ); +const edgedbVersion = JSON.parse(process.env._JEST_EDGEDB_VERSION!); const connectOpts = { ...brokenConnectOpts, + user: edgedbVersion.major >= 6 ? "admin" : "edgedb", tlsCAFile: undefined, tlsSecurity: "insecure", };