From ce90321876a35280d8d99a520d4e2da724d55f27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gu=C3=B0mundur=20Bj=C3=B6rn=20Birkisson?=
 <gudmundur.birkisson@oda.com>
Date: Tue, 9 Jan 2024 13:55:59 +0100
Subject: [PATCH] Add way to set trace exporter request headers

---
 README.md                             | 17 ++++++++++++++++-
 tests/tracing/test_writer.py          | 27 ++++++++++++++++++++++++++-
 troncos/contrib/structlog/__init__.py |  2 +-
 troncos/tracing/__init__.py           |  4 ++--
 troncos/tracing/_enums.py             | 17 ++++++++++++++++-
 troncos/tracing/_otel.py              | 12 ++++++------
 6 files changed, 67 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index b970a15d..6e8df32a 100644
--- a/README.md
+++ b/README.md
@@ -21,7 +21,6 @@
 </p>
 
 <!-- TOC -->
-
 - [Etymology](#etymology)
 - [Installation](#installation)
 - [Tracing](#tracing)
@@ -136,6 +135,22 @@ configure_tracer(
 )
 ```
 
+### Setting headers for the exporter
+
+```python
+from troncos.tracing import configure_tracer, Exporter, ExporterType
+
+TRACE_HOST = "127.0.0.1" # Usually obtained from env variables.
+TRACE_PORT = "4317"
+
+configure_tracer(
+    enabled=False, # Set to True when TRACE_HOST is configured.
+    service_name='SERVICE_NAME',
+    endpoint=f"http://{TRACE_HOST}:{TRACE_PORT}",
+    exporter=Exporter(ExporterType.GRPC, headers={"my", "header"}),
+)
+```
+
 ### Instrument your code
 
 Manual instrumentation of your code is described in the [ddtrace docs](https://ddtrace.readthedocs.io/en/stable/basic_usage.html#manual-instrumentation).
diff --git a/tests/tracing/test_writer.py b/tests/tracing/test_writer.py
index d3e4c818..f8f30a71 100644
--- a/tests/tracing/test_writer.py
+++ b/tests/tracing/test_writer.py
@@ -4,7 +4,7 @@
 from ddtrace import Tracer
 from pytest_httpserver import HTTPServer
 
-from troncos.tracing._enums import Exporter
+from troncos.tracing._enums import Exporter, ExporterType
 from troncos.tracing._writer import OTELWriter
 
 
@@ -73,3 +73,28 @@ def test_exceptions(httpserver: HTTPServer) -> None:
     data = tracer_assert(httpserver)
     assert b"service.name\x12\x10\n\x0etest_exception" in data
     assert b"exception.type\x12\x19\n\x17builtins.AssertionError" in data
+
+
+def test_headers(httpserver: HTTPServer) -> None:
+    httpserver.expect_oneshot_request("/v1/trace").respond_with_data("OK")
+
+    tracer = Tracer()
+    tracer.configure(
+        writer=OTELWriter(
+            service_name="test_headers",
+            service_attributes={},
+            endpoint=httpserver.url_for("/v1/trace"),
+            exporter=Exporter(
+                exporter_type=ExporterType.HTTP, headers={"test-header": "works"}
+            ),
+        )
+    )
+
+    with tracer.trace("test"):
+        pass
+
+    tracer.flush()  # type: ignore[no-untyped-call]
+
+    assert len(httpserver.log), "We should have gotten 1 request"
+    req, _ = httpserver.log[0]
+    assert req.headers.get("test-header") == "works"
diff --git a/troncos/contrib/structlog/__init__.py b/troncos/contrib/structlog/__init__.py
index 0eb283f3..636de118 100644
--- a/troncos/contrib/structlog/__init__.py
+++ b/troncos/contrib/structlog/__init__.py
@@ -12,7 +12,7 @@
 try:
     from structlog_sentry import SentryProcessor
 except ImportError:
-    SentryProcessor = None  # type: ignore
+    SentryProcessor = None
 
 shared_processors: list[structlog.types.Processor] = [
     # Add the name of the logger to event dict.
diff --git a/troncos/tracing/__init__.py b/troncos/tracing/__init__.py
index f80473fa..ab6bbf72 100644
--- a/troncos/tracing/__init__.py
+++ b/troncos/tracing/__init__.py
@@ -2,10 +2,10 @@
 
 import ddtrace
 
-from ._enums import Exporter
+from ._enums import Exporter, ExporterType
 from ._writer import OTELWriter
 
-__all__ = ["Exporter"]
+__all__ = ["Exporter", "ExporterType"]
 
 
 def configure_tracer(
diff --git a/troncos/tracing/_enums.py b/troncos/tracing/_enums.py
index 6a3e1ff1..468ac567 100644
--- a/troncos/tracing/_enums.py
+++ b/troncos/tracing/_enums.py
@@ -1,6 +1,21 @@
 from enum import Enum
 
 
-class Exporter(Enum):
+class ExporterType(Enum):
     HTTP = "http"
     GRPC = "grpc"
+
+
+class Exporter:
+    HTTP: "Exporter" = None  # type: ignore[assignment]
+    GRPC: "Exporter" = None  # type: ignore[assignment]
+
+    def __init__(
+        self, exporter_type: ExporterType, headers: dict[str, str] | None = None
+    ) -> None:
+        self.exporter_type = exporter_type
+        self.headers = headers
+
+
+Exporter.HTTP = Exporter(ExporterType.HTTP)
+Exporter.GRPC = Exporter(ExporterType.GRPC)
diff --git a/troncos/tracing/_otel.py b/troncos/tracing/_otel.py
index bbd47b0d..27747bb7 100644
--- a/troncos/tracing/_otel.py
+++ b/troncos/tracing/_otel.py
@@ -13,14 +13,14 @@
 )
 from structlog import get_logger
 
-from ._enums import Exporter
+from ._enums import Exporter, ExporterType
 
 try:
     from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
         OTLPSpanExporter as GRPCSpanExporter,
     )
 except ImportError:
-    GRPCSpanExporter = None  # type: ignore
+    GRPCSpanExporter = None
 
 
 logger = get_logger()
@@ -41,16 +41,16 @@ def get_otel_span_processors(
     span_exporter: SpanExporter
 
     # Exporter
-    if exporter == Exporter.HTTP:
-        span_exporter = HTTPSpanExporter(endpoint=endpoint)
-    elif exporter == Exporter.GRPC:
+    if exporter.exporter_type == ExporterType.HTTP:
+        span_exporter = HTTPSpanExporter(endpoint=endpoint, headers=exporter.headers)
+    elif exporter.exporter_type == ExporterType.GRPC:
         if GRPCSpanExporter is None:
             raise RuntimeError(
                 "opentelemetry-exporter-otlp-proto-grpc needs to be installed "
                 "to use the GRPC exporter."
             )
 
-        span_exporter = GRPCSpanExporter(endpoint=endpoint)
+        span_exporter = GRPCSpanExporter(endpoint=endpoint, headers=exporter.headers)
     else:
         raise RuntimeError("Unsupported span exporter.")