Merge pull request #72 from aflorithmic/feat/tracing-headers

feat: add customer trace id header

Author: harry97uk

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to `audiostack` will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [2.10.0] - 2025-02-13
+
+### Added
+
+- Customer trace id header can now be set per production using the `use_trace` context manager.
+- Can pass through custom headers into `send_request`.
+
 ## [2.9.0] - 2025-01-23
 
 ### Fixed

audiostack/helpers/request_interface.py

@@ -1,12 +1,18 @@
+import contextlib
 import json
 import shutil
-from typing import Any, Callable, Dict, Optional, Union
+from contextvars import ContextVar
+from typing import Any, Callable, Dict, Generator, Optional, Union
 
 import requests
 
 import audiostack
 from audiostack.helpers.request_types import RequestTypes
 
+_current_trace_id: ContextVar[Optional[str]] = ContextVar(
+    "current_trace_id", default=None
+)
+
 
 def remove_empty(data: Any) -> Any:
     if not (isinstance(data, dict) or isinstance(data, list)):
@@ -32,14 +38,19 @@ def __init__(self, family: str) -> None:
         self.family = family
 
     @staticmethod
-    def make_header() -> dict:
-        header = {
+    def make_header(headers: Optional[dict] = None) -> dict:
+        new_headers = {
             "x-api-key": audiostack.api_key,
             "x-python-sdk-version": audiostack.sdk_version,
         }
+        current_trace_id = _current_trace_id.get()
+        if current_trace_id is not None:
+            new_headers["x-customer-trace-id"] = current_trace_id
         if audiostack.assume_org_id:
-            header["x-assume-org"] = audiostack.assume_org_id
-        return header
+            new_headers["x-assume-org"] = audiostack.assume_org_id
+        if headers:
+            new_headers.update(headers)
+        return new_headers
 
     def resolve_response(self, r: Any) -> dict:
         if self.DEBUG_PRINT:
@@ -82,6 +93,7 @@ def send_request(
         path_parameters: Optional[Union[dict, str]] = None,
         query_parameters: Optional[Union[dict, str]] = None,
         overwrite_base_url: Optional[str] = None,
+        headers: Optional[dict] = None,
     ) -> Any:
         if overwrite_base_url:
             url = overwrite_base_url
@@ -111,15 +123,15 @@ def send_request(
             }
 
             return self.resolve_response(
-                FUNC_MAP[rtype](url=url, json=json, headers=self.make_header())
+                FUNC_MAP[rtype](url=url, json=json, headers=self.make_header(headers))
             )
         elif rtype == RequestTypes.GET:
             if path_parameters:
                 url = f"{url}/{path_parameters}"
 
             return self.resolve_response(
                 requests.get(
-                    url=url, params=query_parameters, headers=self.make_header()
+                    url=url, params=query_parameters, headers=self.make_header(headers)
                 )
             )
         elif rtype == RequestTypes.DELETE:
@@ -128,7 +140,7 @@ def send_request(
 
             return self.resolve_response(
                 requests.delete(
-                    url=url, params=query_parameters, headers=self.make_header()
+                    url=url, params=query_parameters, headers=self.make_header(headers)
                 )
             )
 
@@ -142,3 +154,12 @@ def download_url(cls, url: str, name: str, destination: str) -> None:
         local_filename = f"{destination}/{name}"
         with open(local_filename, "wb") as f:
             shutil.copyfileobj(r.raw, f)
+
+
+@contextlib.contextmanager
+def use_trace(trace_id: str) -> Generator[None, None, None]:
+    token = _current_trace_id.set(trace_id)
+    try:
+        yield
+    finally:
+        _current_trace_id.reset(token)

audiostack/tests/helpers/test_request_interface.py

@@ -3,7 +3,8 @@
 import pytest
 
 import audiostack
-from audiostack.helpers.request_interface import RequestInterface
+from audiostack.helpers.request_interface import RequestInterface, use_trace
+from audiostack.helpers.request_types import RequestTypes
 
 
 @patch("audiostack.helpers.request_interface.open")
@@ -59,3 +60,49 @@ def test_RequestInterface_download_url_4XX(
     mock_requests.get.return_value.status_code = 400
     with pytest.raises(Exception):
         RequestInterface.download_url(url="foo", name="bar", destination="baz")
+
+
+@patch("audiostack.helpers.request_interface.requests")
+def test_RequestInterface_with_trace_id(mock_requests: Mock) -> None:
+    body = {
+        "scriptText": "scriptText",
+        "projectName": "projectName",
+        "moduleName": "moduleName",
+        "scriptName": "scriptName",
+        "metadata": "metadata",
+    }
+    with use_trace("trace_id"):
+        mock_requests.post.return_value.status_code = 200
+        interface = RequestInterface(family="content")
+        interface.send_request(rtype=RequestTypes.POST, route="script", json=body)
+        mock_requests.post.assert_called_once_with(
+            url=f"{audiostack.api_base}/content/script",
+            headers={
+                "x-api-key": audiostack.api_key,
+                "x-python-sdk-version": audiostack.sdk_version,
+                "x-customer-trace-id": "trace_id",
+            },
+            json=body,
+        )
+
+
+@patch("audiostack.helpers.request_interface.requests")
+def test_RequestInterface_with_no_trace_id(mock_requests: Mock) -> None:
+    body = {
+        "scriptText": "scriptText",
+        "projectName": "projectName",
+        "moduleName": "moduleName",
+        "scriptName": "scriptName",
+        "metadata": "metadata",
+    }
+    mock_requests.post.return_value.status_code = 200
+    interface = RequestInterface(family="content")
+    interface.send_request(rtype=RequestTypes.POST, route="script", json=body)
+    mock_requests.post.assert_called_once_with(
+        url=f"{audiostack.api_base}/content/script",
+        headers={
+            "x-api-key": audiostack.api_key,
+            "x-python-sdk-version": audiostack.sdk_version,
+        },
+        json=body,
+    )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "audiostack"
-version = "2.9.1"
+version = "2.10.0"
 description = "Python SDK for Audiostack API"
 authors = ["Aflorithmic <[email protected]>"]
 repository = "https://github.com/aflorithmic/audiostack-python"