From 9dc0b23a3f9f918cc92929df0c139ceae4639252 Mon Sep 17 00:00:00 2001
From: Mose Mueller <mosemueller@gmail.com>
Date: Mon, 5 May 2025 11:34:53 +0200
Subject: [PATCH] feat: adds support for headers to InfluxDBv1Session and
 InfluxDBSession

---
 pydase_service_base/database/config.py             | 2 ++
 pydase_service_base/database/influxdb_session.py   | 5 +++++
 pydase_service_base/database/influxdbv1_session.py | 2 ++
 3 files changed, 9 insertions(+)

diff --git a/pydase_service_base/database/config.py b/pydase_service_base/database/config.py
index 94a1cb2..65cb77a 100644
--- a/pydase_service_base/database/config.py
+++ b/pydase_service_base/database/config.py
@@ -29,6 +29,7 @@ class InfluxDBConfig(BaseConfig):  # type: ignore
     url: str
     org: str
     token: SecretStr
+    headers: dict[str, str] = {}  # noqa: RUF012
 
 
 class InfluxDBv1Config(BaseConfig):  # type: ignore
@@ -39,3 +40,4 @@ class InfluxDBv1Config(BaseConfig):  # type: ignore
     database: str
     ssl: bool = True
     verify_ssl: bool = True
+    headers: dict[str, str] = {}  # noqa: RUF012
diff --git a/pydase_service_base/database/influxdb_session.py b/pydase_service_base/database/influxdb_session.py
index af5177b..2c905ef 100644
--- a/pydase_service_base/database/influxdb_session.py
+++ b/pydase_service_base/database/influxdb_session.py
@@ -80,12 +80,17 @@ class InfluxDBSession:
         self.url = self._config.url
         self.token = self._config.token.get_secret_value()
         self.org = self._config.org
+        self.headers = self._config.headers
         self._client: InfluxDBClient
         self._write_api: WriteApi
         self._buckets_api: BucketsApi
 
     def __enter__(self) -> Self:
         self._client = InfluxDBClient(url=self.url, token=self.token, org=self.org)
+
+        for header_name, header_value in self.headers.items():
+            self._client.api_client.set_default_header(header_name, header_value)
+
         self._write_api = self._client.write_api(write_options=SYNCHRONOUS)  # type: ignore
         return self
 
diff --git a/pydase_service_base/database/influxdbv1_session.py b/pydase_service_base/database/influxdbv1_session.py
index d08da40..ac85856 100644
--- a/pydase_service_base/database/influxdbv1_session.py
+++ b/pydase_service_base/database/influxdbv1_session.py
@@ -65,6 +65,7 @@ class InfluxDBv1Session:
         self._database = self._config.database
         self._ssl = self._config.ssl
         self._verify_ssl = self._config.verify_ssl
+        self._headers = self._config.headers
 
     def __enter__(self) -> Self:
         self._client = influxdb.InfluxDBClient(
@@ -75,6 +76,7 @@ class InfluxDBv1Session:
             database=self._database,
             ssl=self._ssl,
             verify_ssl=self._verify_ssl,
+            headers=self._headers,
         )
         return self