From 70701b1f8e759ca254b085a24ef9ab236d93b5d1 Mon Sep 17 00:00:00 2001
From: Pablo Collins <pcollins@splunk.com>
Date: Fri, 1 Nov 2024 11:13:58 -0400
Subject: [PATCH] Add support for splunk access token env var

---
 src/splunk_otel/distro.py | 11 ++++++++++-
 tests/ott_lib.py          | 15 +++++++++++++++
 tests/ott_sf_token.py     | 29 +++++++++++++++++++++++++++++
 tests/ott_trace_loop.py   | 21 ++++-----------------
 tests/test_distro.py      |  8 ++++++++
 5 files changed, 66 insertions(+), 18 deletions(-)
 create mode 100644 tests/ott_lib.py
 create mode 100644 tests/ott_sf_token.py

diff --git a/src/splunk_otel/distro.py b/src/splunk_otel/distro.py
index 64c073a..580d4d2 100644
--- a/src/splunk_otel/distro.py
+++ b/src/splunk_otel/distro.py
@@ -16,8 +16,11 @@
 
 from opentelemetry.instrumentation.distro import BaseDistro
 from opentelemetry.instrumentation.system_metrics import SystemMetricsInstrumentor
+from opentelemetry.sdk.environment_variables import OTEL_EXPORTER_OTLP_HEADERS
 
-from splunk_otel.env import DEFAULTS, OTEL_METRICS_ENABLED, Env
+from splunk_otel.env import DEFAULTS, OTEL_METRICS_ENABLED, SPLUNK_ACCESS_TOKEN, Env
+
+X_SF_TOKEN = "x-sf-token"  # noqa S105
 
 
 class SplunkDistro(BaseDistro):
@@ -32,11 +35,17 @@ def __init__(self):
 
     def _configure(self, **kwargs):  # noqa: ARG002
         self.set_env_defaults()
+        self.configure_headers()
 
     def set_env_defaults(self):
         for key, value in DEFAULTS.items():
             self.env.setdefault(key, value)
 
+    def configure_headers(self):
+        tok = self.env.getval(SPLUNK_ACCESS_TOKEN)
+        if tok:
+            self.env.list_append(OTEL_EXPORTER_OTLP_HEADERS, f"{X_SF_TOKEN}={tok}")
+
     def load_instrumentor(self, entry_point, **kwargs):
         #  This method is called in a loop by opentelemetry-instrumentation
         if is_system_metrics_instrumentor(entry_point) and not self.env.is_true(OTEL_METRICS_ENABLED):
diff --git a/tests/ott_lib.py b/tests/ott_lib.py
new file mode 100644
index 0000000..e732e5c
--- /dev/null
+++ b/tests/ott_lib.py
@@ -0,0 +1,15 @@
+import time
+from pathlib import Path
+
+from opentelemetry import trace
+
+
+def project_path():
+    return str(Path(__file__).parent.parent)
+
+
+def trace_loop(loops):
+    tracer = trace.get_tracer("my-tracer")
+    for _ in range(loops):
+        with tracer.start_as_current_span("my-span"):
+            time.sleep(0.5)
diff --git a/tests/ott_sf_token.py b/tests/ott_sf_token.py
new file mode 100644
index 0000000..8d94209
--- /dev/null
+++ b/tests/ott_sf_token.py
@@ -0,0 +1,29 @@
+from oteltest import Telemetry
+from ott_lib import project_path, trace_loop
+
+if __name__ == "__main__":
+    trace_loop(12)
+
+
+class AccessTokenOtelTest:
+    def requirements(self):
+        return project_path(), "oteltest"
+
+    def environment_variables(self):
+        return {
+            "OTEL_SERVICE_NAME": "my-svc",
+            "SPLUNK_ACCESS_TOKEN": "s3cr3t",
+        }
+
+    def wrapper_command(self):
+        return "opentelemetry-instrument"
+
+    def on_start(self):
+        return None
+
+    def on_stop(self, telemetry: Telemetry, stdout: str, stderr: str, returncode: int) -> None:  # noqa: ARG002
+        for request in telemetry.get_trace_requests():
+            assert request.headers.get("x-sf-token") == "s3cr3t"
+
+    def is_http(self):
+        return False
diff --git a/tests/ott_trace_loop.py b/tests/ott_trace_loop.py
index f4c5822..0b8434e 100644
--- a/tests/ott_trace_loop.py
+++ b/tests/ott_trace_loop.py
@@ -1,27 +1,15 @@
-import time
-from pathlib import Path
-
-from opentelemetry import trace
 from oteltest.telemetry import num_spans
+from ott_lib import project_path, trace_loop
 
 NUM_SPANS = 12
 
-
-def trace_loop(loops):
-    tracer = trace.get_tracer("my-tracer")
-    for _ in range(loops):
-        with tracer.start_as_current_span("my-span"):
-            time.sleep(0.5)
-
-
 if __name__ == "__main__":
     trace_loop(NUM_SPANS)
 
 
-class MyOtelTest:
+class NumSpansOtelTest:
     def requirements(self):
-        parent_dir = str(Path(__file__).parent.parent)
-        return parent_dir, "oteltest"
+        return project_path(), "oteltest"
 
     def environment_variables(self):
         return {
@@ -35,8 +23,7 @@ def on_start(self):
         return None
 
     def on_stop(self, telemetry, stdout: str, stderr: str, returncode: int) -> None:  # noqa: ARG002
-        span_count = num_spans(telemetry)
-        assert span_count == NUM_SPANS
+        assert num_spans(telemetry) == NUM_SPANS
 
     def is_http(self):
         return False
diff --git a/tests/test_distro.py b/tests/test_distro.py
index 897f0bc..85d8f33 100644
--- a/tests/test_distro.py
+++ b/tests/test_distro.py
@@ -26,3 +26,11 @@ def test_distro_env():
     # spot check default env vars
     assert env_store["OTEL_TRACES_EXPORTER"] == "otlp"
     assert len(env_store) == 11
+
+
+def test_access_token():
+    env_store = {"SPLUNK_ACCESS_TOKEN": "abc123"}
+    sd = SplunkDistro()
+    sd.env = Env(env_store)
+    sd.configure()
+    assert env_store["OTEL_EXPORTER_OTLP_HEADERS"] == "x-sf-token=abc123"