From c3fdce829833d692ea2e81e3d424659bc023a274 Mon Sep 17 00:00:00 2001
From: Fantix King <fantix.king@gmail.com>
Date: Tue, 3 Dec 2024 15:32:25 -0500
Subject: [PATCH] API fix: with_query_tag()/without_query_tag()

---
 gel/options.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/gel/options.py b/gel/options.py
index 239cb4c1..fb63b811 100644
--- a/gel/options.py
+++ b/gel/options.py
@@ -13,6 +13,7 @@
 
 
 _RetryRule = namedtuple("_RetryRule", ["attempts", "backoff"])
+TAG_NAME = "tag"
 
 
 def default_backoff(attempt):
@@ -413,17 +414,25 @@ def without_globals(self, *global_names):
         )
         return result
 
-    def with_annotation(self, name: str, value: str):
+    def with_query_tag(self, tag: str):
+        for prefix in ["edgedb/", "gel/"]:
+            if tag.startswith(prefix):
+                raise errors.InvalidArgumentError(f"reserved tag: {prefix}*")
+        if len(tag) > 128:
+            raise errors.InvalidArgumentError(
+                "tag too long (> 128 characters)"
+            )
+
         result = self._shallow_clone()
         result._options = self._options.with_annotations(
-            self._options.annotations | {name: value}
+            self._options.annotations | {TAG_NAME: tag}
         )
         return result
 
-    def without_annotation(self, name: str):
+    def without_query_tag(self):
         result = self._shallow_clone()
         annotations = self._options.annotations.copy()
-        annotations.pop(name, None)
+        annotations.pop(TAG_NAME, None)
         result._options = self._options.with_annotations(annotations)
         return result