Make it possible to unset tag

edgedb · Nov 27, 2024 · dd8edc1 · dd8edc1
1 parent ff04440
commit dd8edc1
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 19 deletions.
diff --git 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
@@ -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
@@ -142,7 +142,11 @@ export class Session {
   readonly globals: Record<string, any>;
 
   /** @internal */
-  annotations: Record<string, string> = {};
+  annotations = new Map<string, string>();
+
+  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;
   }