use-ttd-auth-from-dataclient-instead-of-passing-it-per-call

README.md

@@ -198,20 +198,21 @@ client = TtdDatabricksClient.from_params(
 
 Provide your own [`DataClient`](https://github.com/thetradedesk/ttd-data-python/blob/main/src/ttd_data/sdk.py) instance to control the underlying HTTP transport directly.
 Use this when you need to configure options not exposed by `from_params()`, or to inject a mock in tests.
+The `DataClient` you pass in must carry your API token as `ttd_auth`; every request the SDK makes authenticates with it.
 
 ```python
 from ttd_data import DataClient
 from ttd_databricks_python.ttd_databricks import TtdDatabricksClient
 
-# Configure DataClient with custom HTTP settings.
+# Configure DataClient with your API token and custom HTTP settings.
 data_client = DataClient(
+    ttd_auth="<ttd-auth-token>",                      # your TTD platform API token
     server_url="https://custom-server.example.com",  # override default server URL
     timeout_ms=10000,                                 # request timeout in milliseconds
 )
 
 client = TtdDatabricksClient(
     data_api_client=data_client,
-    api_token="<ttd-auth-token>",
     spark=spark,  # optional; spark variable available from the Databricks notebook runtime
 )
 ```
@@ -456,15 +457,13 @@ from ttd_data.utils.retries import BackoffStrategy, RetryConfig
 from ttd_databricks_python.ttd_databricks import TtdDatabricksClient
 
 data_client = DataClient(
+    ttd_auth="<ttd-auth-token>",                      # your TTD platform API token
     server_url="https://custom-server.example.com",  # override default server URL
     timeout_ms=10000,                                 # request timeout in milliseconds
     retry_config=RetryConfig("backoff", BackoffStrategy(1000, 60000, 1.5, 3600000), True),  # custom retry config
 )
 
-client = TtdDatabricksClient(
-    data_api_client=data_client,
-    api_token="<ttd-auth-token>",
-)
+client = TtdDatabricksClient(data_api_client=data_client)
 ```
 
 In batch processing mode, a `DataClient` singleton is maintained per Spark worker process to enable HTTP connection reuse across batches, reducing overhead during distributed execution.

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ttd-databricks"
-version = "0.5.0"
+version = "0.6.0"
 description = "Client implementation and helper functions for integrating with the TTD Databricks services."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -15,7 +15,7 @@ authors = [
 ]
 
 dependencies = [
-    "ttd-data>=0.2.6,<0.3.0",
+    "ttd-data>=0.3.1,<0.4.0",
     "pandas>=1.0.5",
     "pyarrow>=4.0.0",
     "setuptools>=63.4.1",

tests/test_placeholder.py

@@ -13,4 +13,3 @@ def test_placeholder(self) -> None:
 
 if __name__ == "__main__":
     unittest.main()
-

tests/unit/test_batch_process_early_exit.py

@@ -34,7 +34,6 @@
 def _make_client(spark: SparkSession) -> TtdDatabricksClient:
     return TtdDatabricksClient(
         data_api_client=MagicMock(spec=DataClient),
-        api_token="test-token",
         spark=spark,
     )
 

tests/unit/test_call_api.py

@@ -17,10 +17,8 @@
 from typing import Any
 from unittest.mock import MagicMock, patch
 
-import pytest
-
 import httpx
-
+import pytest
 from ttd_data import DataClient
 from ttd_data.errors import DataError, NoResponseError, ResponseValidationError
 
@@ -30,7 +28,7 @@
 
 
 def _make_client() -> TtdDatabricksClient:
-    return TtdDatabricksClient(data_api_client=MagicMock(spec=DataClient), api_token="test-token")
+    return TtdDatabricksClient(data_api_client=MagicMock(spec=DataClient))
 
 
 def _make_rows(*dicts: dict[str, Any]) -> list[MagicMock]:

tests/unit/test_client_helpers.py

@@ -19,7 +19,6 @@
 def _make_client(**kwargs) -> TtdDatabricksClient:  # type: ignore[no-untyped-def]
     return TtdDatabricksClient(
         data_api_client=MagicMock(spec=DataClient),
-        api_token="test-token",
         **kwargs,
     )
 

tests/unit/test_contexts.py

@@ -14,7 +14,6 @@
 )
 from ttd_databricks_python.ttd_databricks.endpoints import TTDEndpoint
 
-
 _REQUEST_TYPE = PartnerDsrRequestType.OPT_OUT
 _DATA_ORIGINS = [DataOrigin(id="test-origin", type=DataOriginType.DATA_PROVIDER)]
 
@@ -72,9 +71,7 @@ def test_deletion_optout_advertiser_context_verify_context_pickling():
 
 
 def test_deletion_optout_thirdparty_context_verify_context_pickling():
-    ctx = DeletionOptOutThirdPartyContext(
-        data_provider_id="prov123", request_type=_REQUEST_TYPE, brand_id="brand99"
-    )
+    ctx = DeletionOptOutThirdPartyContext(data_provider_id="prov123", request_type=_REQUEST_TYPE, brand_id="brand99")
     restored = pickle.loads(pickle.dumps(ctx))
     assert restored.data_provider_id == "prov123"
     assert restored.request_type == _REQUEST_TYPE

tests/unit/test_handlers_build_items.py

@@ -4,20 +4,19 @@
 """
 
 from datetime import datetime, timezone
-from typing import Union
 
 import numpy as np
 import pytest
+from ttd_data.models import AdvertiserDataItem, OfflineConversionDataItem, PartnerDsrDataItem, ThirdPartyDataItem
+from ttd_data.types import UNSET
 
 import ttd_databricks_python.ttd_databricks.handlers.advertiser as adv_handler
-from ttd_databricks_python.ttd_databricks.id_types import normalize_id_type
 import ttd_databricks_python.ttd_databricks.handlers.deletion_optout_advertiser as del_adv_handler
 import ttd_databricks_python.ttd_databricks.handlers.deletion_optout_merchant as del_merch_handler
 import ttd_databricks_python.ttd_databricks.handlers.deletion_optout_thirdparty as del_tp_handler
 import ttd_databricks_python.ttd_databricks.handlers.offline_conversion as oc_handler
 import ttd_databricks_python.ttd_databricks.handlers.third_party as tp_handler
-from ttd_data.models import AdvertiserDataItem, OfflineConversionDataItem, PartnerDsrDataItem, ThirdPartyDataItem
-from ttd_data.types import UNSET
+from ttd_databricks_python.ttd_databricks.id_types import normalize_id_type
 
 # UNSET is not a singleton — the SDK creates fresh Unset() instances per field.
 # Use isinstance check rather than identity (is).
@@ -27,7 +26,7 @@
 # An array<struct> column reaches build_items as a list via the adhoc path
 # (collect + asDict) and as a numpy array via the batch path (mapInPandas).
 # build_items must handle both, so array-column tests run against each shape.
-def _build_array_column(array_type: type, items: list[dict]) -> Union[list, np.ndarray]:
+def _build_array_column(array_type: type, items: list[dict]) -> list | np.ndarray:
     return items if array_type is list else np.array(items, dtype=object)
 
 
@@ -43,7 +42,7 @@ def test_builds_advertiser_data_item_with_correct_fields(self):
         # Handler maps id_type → AdvertiserDataItem field dynamically: {d["id_type"]: d["id_value"]}
         item = adv_handler.build_items([self._MINIMAL])[0]
         assert isinstance(item, AdvertiserDataItem)
-        assert getattr(item, "tdid") == "test-tdid-value"
+        assert item.tdid == "test-tdid-value"
         assert item.data[0].name == "test-segment-name"
 
     def test_none_optional_fields_are_not_sent_to_api(self):
@@ -76,7 +75,7 @@ class TestThirdPartyBuildItems:
     def test_builds_third_party_data_item_with_correct_fields(self):
         item = tp_handler.build_items([self._MINIMAL])[0]
         assert isinstance(item, ThirdPartyDataItem)
-        assert getattr(item, "tdid") == "test-tdid-value"
+        assert item.tdid == "test-tdid-value"
         assert item.data[0].name == "test-segment-name"
 
     def test_none_optional_fields_are_not_sent_to_api(self):
@@ -99,19 +98,19 @@ def test_optional_fields_are_passed_through_when_provided(self):
 def test_deletion_optout_advertiser_returns_partner_dsr_item_with_correct_id():
     item = del_adv_handler.build_items([{"id_type": "TDID", "id_value": "test-advertiser-tdid"}])[0]
     assert isinstance(item, PartnerDsrDataItem)
-    assert getattr(item, "tdid") == "test-advertiser-tdid"
+    assert item.tdid == "test-advertiser-tdid"
 
 
 def test_deletion_optout_thirdparty_returns_partner_dsr_item_with_correct_id():
     item = del_tp_handler.build_items([{"id_type": "UID2", "id_value": "test-thirdparty-uid2"}])[0]
     assert isinstance(item, PartnerDsrDataItem)
-    assert getattr(item, "uid2") == "test-thirdparty-uid2"
+    assert item.uid2 == "test-thirdparty-uid2"
 
 
 def test_deletion_optout_merchant_returns_partner_dsr_item_with_correct_id():
     item = del_merch_handler.build_items([{"id_type": "TDID", "id_value": "test-merchant-tdid"}])[0]
     assert isinstance(item, PartnerDsrDataItem)
-    assert getattr(item, "tdid") == "test-merchant-tdid"
+    assert item.tdid == "test-merchant-tdid"
 
 
 # --------------------------------------------------------------------------- #
@@ -143,8 +142,13 @@ def test_user_ids_converted_to_user_id_array_with_type_codes(self, array_type):
 
     def test_all_user_id_types_map_to_correct_codes(self):
         type_map = {
-            "TDID": "0", "DAID": "1", "UID2": "2", "UID2Token": "3",
-            "EUID": "4", "EUIDToken": "5", "RampID": "6",
+            "TDID": "0",
+            "DAID": "1",
+            "UID2": "2",
+            "UID2Token": "3",
+            "EUID": "4",
+            "EUIDToken": "5",
+            "RampID": "6",
         }
         for id_type, expected_code in type_map.items():
             row = {**self._MINIMAL, "user_ids": [{"type": id_type, "id": f"test-{id_type}-value"}]}
@@ -209,4 +213,4 @@ def test_collect_raw_pii_ids_keeps_only_pii_types(self, array_type):
         assert oc_handler.collect_raw_pii_ids_per_row(rows) == [["a@example.com"]]
 
     def test_collect_raw_pii_ids_handles_missing_user_ids(self):
-        assert oc_handler.collect_raw_pii_ids_per_row([self._MINIMAL]) == [[]]
+        assert oc_handler.collect_raw_pii_ids_per_row([self._MINIMAL]) == [[]]

tests/unit/test_process_partitions.py

@@ -37,7 +37,9 @@
     server_url=None,
     retry_config=None,
     timeout_ms=10_000,
+    ttd_auth="not-a-real-token",
     uid2_config=None,
+    graphql_server_url=None,
 )
 
 
@@ -46,6 +48,7 @@ class _StubHandler(BaseHTTPRequestHandler):
 
     status_code = 500
     request_count = 0
+    auth_headers: list[str | None] = []
     # ThreadingHTTPServer handles each request on its own thread; `+= 1` is a
     # non-atomic read-modify-write, so guard it rather than relying on the spark
     # fixture staying single-threaded.
@@ -56,10 +59,12 @@ def configure(cls, status_code: int) -> None:
         with cls.counter_lock:
             cls.status_code = status_code
             cls.request_count = 0
+            cls.auth_headers = []
 
     def do_POST(self) -> None:  # noqa: N802 — required by stdlib BaseHTTPRequestHandler
         with type(self).counter_lock:
             type(self).request_count += 1
+            type(self).auth_headers.append(self.headers.get("TTD-Auth"))
         body = b'{"Message":"forced error for test"}'
         self.send_response(type(self).status_code)
         self.send_header("Content-Type", "application/json")
@@ -112,7 +117,6 @@ def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, stub_server:
         df=input_df,
         batch_size=3,
         output_schema=output_schema,
-        api_token="not-a-real-token",
         context=context,
         parallelism=2,
         client_config=_NO_RETRY_CLIENT_CONFIG,
@@ -127,6 +131,9 @@ def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, stub_server:
     assert result_df.schema.fieldNames() == output_schema.fieldNames()
     # 4. Input column values survive Arrow → pandas → dict → pandas → Arrow round-trip.
     assert {row["id_value"] for row in result_rows} == set(input_ids)
+    # 5. The worker's rebuilt DataClient authenticates: ttd_auth travels in the client_config
+    #    snapshot, not as a separate per-call argument.
+    assert set(_StubHandler.auth_headers) == {"not-a-real-token"}
 
 
 @pytest.mark.parametrize(
@@ -136,7 +143,9 @@ def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, stub_server:
         403,  # token valid but not entitled to this advertiser or data provider
     ],
 )
-def test_401_and_403_stop_partition_without_failing_job(spark: SparkSession, stub_server: str, status_code: int) -> None:
+def test_401_and_403_stop_partition_without_failing_job(
+    spark: SparkSession, stub_server: str, status_code: int
+) -> None:
     """401/403 stop the partition without raising. The batch that was sent keeps
     the server's own status; every row after it is ABORTED, meaning it was never submitted."""
     rows = [("TDID", f"id-{i}", "seg-a", None, None) for i in range(7)]
@@ -149,7 +158,6 @@ def test_401_and_403_stop_partition_without_failing_job(spark: SparkSession, stu
         df=input_df,
         batch_size=3,
         output_schema=output_schema,
-        api_token="not-a-real-token",
         context=context,
         parallelism=1,
         client_config=_NO_RETRY_CLIENT_CONFIG,
@@ -182,7 +190,6 @@ def test_other_4xx_fails_only_its_own_batch(spark: SparkSession, stub_server: st
         df=input_df,
         batch_size=3,
         output_schema=output_schema,
-        api_token="not-a-real-token",
         context=context,
         parallelism=1,
         client_config=_NO_RETRY_CLIENT_CONFIG,

tests/unit/test_push_data.py

@@ -32,7 +32,6 @@
 def _make_client(spark: SparkSession) -> TtdDatabricksClient:
     return TtdDatabricksClient(
         data_api_client=MagicMock(spec=DataClient),
-        api_token="test-token",
         spark=spark,
     )
 

tests/unit/test_uid2_resolutions.py

@@ -16,11 +16,11 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from ttd_data import DataClient
+from ttd_data.uid2 import UID2Resolution
 
 import ttd_databricks_python.ttd_databricks.handlers.advertiser as adv_handler
 import ttd_databricks_python.ttd_databricks.handlers.offline_conversion as oc_handler
-from ttd_data import DataClient
-from ttd_data.uid2 import UID2Resolution
 from ttd_databricks_python.ttd_databricks.contexts import AdvertiserContext, OfflineConversionContext
 from ttd_databricks_python.ttd_databricks.endpoints import TTDEndpoint
 from ttd_databricks_python.ttd_databricks.id_types import is_raw_pii_id_type
@@ -31,7 +31,6 @@
 from ttd_databricks_python.ttd_databricks.ttd_client import TtdDatabricksClient
 from ttd_databricks_python.ttd_databricks.utils import attach_resolutions
 
-
 # --------------------------------------------------------------------------- #
 # id_types normalization                                                        #
 # --------------------------------------------------------------------------- #
@@ -306,7 +305,7 @@ def test_raises_with_alter_table_hint_when_column_missing(self) -> None:
 
 
 def _make_client() -> TtdDatabricksClient:
-    return TtdDatabricksClient(data_api_client=MagicMock(spec=DataClient), api_token="test-token")
+    return TtdDatabricksClient(data_api_client=MagicMock(spec=DataClient))
 
 
 def _make_rows(*dicts: dict) -> list[MagicMock]:
@@ -355,9 +354,7 @@ def test_call_api_attaches_uid2_resolutions_array_for_offline_conversion() -> No
     )
 
     with patch("importlib.import_module", return_value=mock_handler):
-        results = client._call_api(
-            OfflineConversionContext(data_provider_id="dp"), rows, batch_index=0
-        )
+        results = client._call_api(OfflineConversionContext(data_provider_id="dp"), rows, batch_index=0)
 
     assert len(results[0][UID2_RESOLUTIONS_COLUMN]) == 1
     assert results[0][UID2_RESOLUTIONS_COLUMN][0]["current_uid2"] == "uid2-y"
@@ -420,4 +417,4 @@ def test_batch_process_config_is_derived_from_data_api_client() -> None:
 
     assert client._data_api_client.config.uid2_config is uid2_cfg
     assert client._data_api_client.config.retry_config is retry_cfg
-
+    assert client._data_api_client.config.ttd_auth == "tok"