Feature: Call tagging

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 1.1.7
+current_version = 1.1.8
 commit = True
 tag = True
 

setup.py

@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name="supergood",
-    version="1.1.7",
+    version="1.1.8",
     author="Alex Klarfeld",
     description="The Python client for Supergood",
     long_description=long_description,

src/supergood/client.py

@@ -2,11 +2,12 @@
 
 import atexit
 import os
+import threading
 import traceback
 from base64 import b64encode
+from contextlib import contextmanager
 from datetime import datetime
 from importlib.metadata import version
-from threading import Lock, Thread
 from urllib.parse import urlparse
 
 from dotenv import load_dotenv
@@ -42,6 +43,7 @@ def __init__(
         config={},
         metadata={},
     ):
+        self.thread_local = threading.local()
         # This PID is used to detect when the client is running in a forked process
         self.main_pid = os.getpid()
         self.base_url = base_url if base_url else DEFAULT_SUPERGOOD_BASE_URL
@@ -84,7 +86,7 @@ def __init__(
 
         self.remote_config = None
         if auto_config and self.base_config["useRemoteConfig"]:
-            self.remote_config_initial_pull = Thread(
+            self.remote_config_initial_pull = threading.Thread(
                 daemon=True, target=self._get_config
             )
             self.remote_config_initial_pull.start()
@@ -114,7 +116,7 @@ def __init__(
         self.flush_thread = RepeatingThread(
             self.flush_cache, self.base_config["flushInterval"] / 1000
         )
-        self.flush_lock = Lock()
+        self.flush_lock = threading.Lock()
         if auto_flush:
             self.flush_thread.start()
         else:
@@ -124,6 +126,17 @@ def __init__(
         if self.base_config["runThreads"]:
             atexit.register(self.close)
 
+    @contextmanager
+    def tagging(self, tags):
+        # tags should be a dict. wrap non-dicts
+        if not isinstance(tags, dict):
+            tags = {"tags": tags}
+        self.thread_local.current_tag = tags
+        try:
+            yield
+        finally:
+            self.thread_local.current_tag = None
+
     def _build_log_payload(self, urls=None, size=None, num_events=None):
         payload = {}
         payload["config"] = self.base_config
@@ -219,6 +232,9 @@ def _cache_request(self, request_id, url, method, body, headers):
                     "search": parsed_url.query,
                     "requestedAt": now,
                 }
+                tags = getattr(self.thread_local, "current_tag", None)
+                if tags is not None:
+                    request["metadata"]["tags"] = tags
                 self._request_cache[request_id] = request
         except Exception:
             payload = self._build_log_payload(
@@ -374,6 +390,7 @@ def flush_cache(self, force=False) -> None:
                 self.log.error(ERRORS["REDACTION"], trace, payload)
             else:  # Only post if no exceptions
                 self.log.debug(f"Flushing {len(data)} items")
+                self.log.debug(data)
                 try:
                     self.api.post_telemetry(
                         {

tests/conftest.py

@@ -31,6 +31,7 @@ def broken_client(broken_redaction, monkeysession):
         client_id="client_id",
         client_secret_id="client_secret_id",
         base_url="https://api.supergood.ai",
+        telemetry_url="https://telemetry.supergood.ai",
         config=config,
     )
     monkeysession.setenv("SG_OVERRIDE_AUTO_FLUSH", "false")
@@ -70,6 +71,7 @@ def supergood_client(request, session_mocker, monkeysession):
         client_id="client_id",
         client_secret_id="client_secret_id",
         base_url="https://api.supergood.ai",
+        telemetry_url="https://telemetry.supergood.ai",
         config=config,
     )
     client._get_config()

tests/test_core.py

@@ -114,6 +114,18 @@ def test_posting_errors(self, supergood_client, session_mocker):
         session_mocker.patch("supergood.api.Api.post_events").start()
         supergood_client.kill()
 
+    def test_tagging(self, httpserver: HTTPServer, supergood_client):
+        with supergood_client.tagging({"key": "val"}):
+            requests.get(httpserver.url_for("/200"))
+
+        supergood_client.flush_cache()
+        assert Api.post_events.call_args is not None
+        args = Api.post_events.call_args[0][0]
+        assert args[0]["request"] is not None
+        assert args[0]["response"] is not None
+        assert args[0]["metadata"]["tags"] == {"key": "val"}
+        supergood_client.kill()
+
     def test_different_http_library(self, httpserver: HTTPServer, supergood_client):
         http = urllib3.PoolManager()
         http.request("GET", f"{TEST_BED_URL}/200")