From af8ab2fad4a0811408619f1be41f4cfc9e6c89cb Mon Sep 17 00:00:00 2001 From: Adithya Samavedhi Date: Mon, 17 Aug 2026 10:05:56 -0700 Subject: [PATCH] return-dataframe-with-succeeded-rows-even-during-TTDAPIError --- README.md | 11 +- ... Connector Data SDK Example Notebook.ipynb | 5 +- pyproject.toml | 2 +- tests/unit/test_call_api.py | 118 +++++++++++++--- tests/unit/test_exceptions.py | 6 - tests/unit/test_process_partitions.py | 129 +++++++++++++++--- tests/unit/test_push_data.py | 61 +++++++++ .../ttd_databricks/batching.py | 109 ++++++++------- .../ttd_databricks/constants.py | 18 +++ .../ttd_databricks/exceptions.py | 11 +- .../ttd_databricks/ttd_client.py | 87 +++++++----- ttd_databricks_python/ttd_databricks/utils.py | 26 ++++ 12 files changed, 439 insertions(+), 144 deletions(-) diff --git a/README.md b/README.md index c98ec8d..bf9e929 100644 --- a/README.md +++ b/README.md @@ -188,7 +188,8 @@ client = TtdDatabricksClient.from_params( api_token="", # your TTD platform API token spark=spark, # optional; auto-detected from Databricks context # server_url="https://..." # optional; see Server Selection - # retry_config=RetryConfig(...) # optional; retry transient errors (429/5xx), see Custom HTTP Client + # retry_config=RetryConfig(...) # optional; transient errors (429/5xx) are retried by + # default, pass None to disable. See Custom HTTP Client # timeout_ms=10000 # optional; per-request timeout in milliseconds ) ``` @@ -382,10 +383,11 @@ To submit email addresses or phone numbers, set the `id_type` column in your inp All SDK exceptions inherit from `TTDError`. +Both `push_data` abd `batch_process` do not raise API call failures — a batch that hits an unrecoverable error fails with its error code, and rows succeeding that were never sent The Trade Desk and fail with `error_code="ABORTED"`. Both are captured inline in the result DataFrame via the `success`, `error_code`, and `error_message` columns, so processing is never interrupted by API or row-level failures. + ```python from ttd_databricks_python.ttd_databricks.exceptions import ( TTDError, - TTDApiError, TTDConfigurationError, TTDSchemaValidationError, ) @@ -394,8 +396,6 @@ try: result_df = client.push_data(df=input_df, context=context) except TTDSchemaValidationError as e: print(f"Missing columns: {e.missing_columns}") -except TTDApiError as e: - print(f"API error on batch {e.batch_index}: {e.status_code} — {e.response_text}") except TTDConfigurationError as e: print(f"Configuration error: {e}") ``` @@ -403,11 +403,8 @@ except TTDConfigurationError as e: | Exception | Cause | |---|---| | `TTDSchemaValidationError` | DataFrame is missing required columns for the endpoint | -| `TTDApiError` | HTTP error or no response from the TTD Data API | | `TTDConfigurationError` | SparkSession not found or PySpark not installed | -For `push_data`, row-level errors are also captured inline in the result DataFrame via the `success`, `error_code`, and `error_message` columns — so processing is not interrupted by individual row failures. - --- ## Server Selection diff --git a/example_notebook/TTD Connector Data SDK Example Notebook.ipynb b/example_notebook/TTD Connector Data SDK Example Notebook.ipynb index 4f1c123..e507827 100644 --- a/example_notebook/TTD Connector Data SDK Example Notebook.ipynb +++ b/example_notebook/TTD Connector Data SDK Example Notebook.ipynb @@ -76,7 +76,6 @@ " TTDEndpoint,\n", " get_ttd_input_schema,\n", " TTDSchemaValidationError,\n", - " TTDApiError,\n", ")\n", "\n", "client = TtdDatabricksClient.from_params(api_token=API_TOKEN)\n", @@ -197,9 +196,7 @@ " )\n", " display(result_df)\n", "except TTDSchemaValidationError as e:\n", - " print(f\"Schema validation failed: {e}\")\n", - "except TTDApiError as e:\n", - " print(f\"API call failed (HTTP {e.status_code}): {e}\")" + " print(f\"Schema validation failed: {e}\")" ] }, { diff --git a/pyproject.toml b/pyproject.toml index 4a30ffb..ba08aed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "ttd-databricks" -version = "0.4.0" +version = "0.5.0" description = "Client implementation and helper functions for integrating with the TTD Databricks services." readme = "README.md" requires-python = ">=3.10" diff --git a/tests/unit/test_call_api.py b/tests/unit/test_call_api.py index f284e03..0269661 100644 --- a/tests/unit/test_call_api.py +++ b/tests/unit/test_call_api.py @@ -5,9 +5,9 @@ 2. Maps failed_lines (by item number) to per-row result dicts. Rows with an item_number get their specific error. Rows without one fall back to the unattributable error (if any). - 3. On 5xx, marks all rows as failed (transient — caller can retry). - On 4xx, raises TTDApiError — unrecoverable client error. - 4. Raises TTDApiError on unexpected non-HTTP exceptions. + 3. Marks all rows in the batch as failed for any error other than auth/permission. + 4. Raises TTDApiError on an auth/permission failure, carrying the error_code the + failing batch's rows should get. The handler module import is patched so no real API calls are made. """ @@ -22,7 +22,7 @@ import httpx from ttd_data import DataClient -from ttd_data.errors import DataError, NoResponseError +from ttd_data.errors import DataError, NoResponseError, ResponseValidationError from ttd_databricks_python.ttd_databricks.contexts import AdvertiserContext from ttd_databricks_python.ttd_databricks.exceptions import TTDApiError @@ -147,9 +147,9 @@ def test_attributable_row_gets_specific_error_others_get_unattributable_fallback results = client._call_api(_CONTEXT, rows, batch_index=0) assert results[0]["success"] is False - assert results[0]["error_code"] == "INVALID_ID" # specific error preserved + assert results[0]["error_code"] == "INVALID_ID" # specific error preserved assert results[1]["success"] is False - assert results[1]["error_code"] == "UNKNOWN" # unattributable as fallback + assert results[1]["error_code"] == "UNKNOWN" # unattributable as fallback def test_failed_line_with_null_message_and_code_fails_all_rows(): @@ -196,10 +196,13 @@ def __str__(self): results = client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=0) assert len(results) == 1 assert results[0]["success"] is False - assert results[0]["error_message"] == "No response" + assert results[0]["error_message"] == "_FakeNoResponseError: No response" + # No HTTP status to report, so the exception name is the code — never NULL, which + # would be indistinguishable from a succeeded row. + assert results[0]["error_code"] == "_FakeNoResponseError" -def test_4xx_error_raises_ttd_api_error(): +def test_400_error_fails_only_its_own_batch(): client = _make_client() rows = _make_rows(_ROW, _ROW) mock_handler = _make_mock_handler() @@ -211,25 +214,108 @@ def test_4xx_error_raises_ttd_api_error(): raw.headers = httpx.Headers({}) mock_handler.call_api.side_effect = DataError("batch error", raw) + with patch("importlib.import_module", return_value=mock_handler): + results = client._call_api(_CONTEXT, rows, batch_index=2) + + assert len(results) == 2 + assert all(r["success"] is False for r in results) + assert all(r["error_code"] == "Bad Request" for r in results) + assert all("not configured" in r["error_message"] for r in results) + + +@pytest.mark.parametrize(("status_code", "expected_error_code"), [(401, "Unauthorized"), (403, "Forbidden")]) +def test_401_and_403_raise_ttd_api_error(status_code: int, expected_error_code: str): + client = _make_client() + rows = _make_rows(_ROW, _ROW) + mock_handler = _make_mock_handler() + mock_handler.build_items.return_value = [MagicMock(), MagicMock()] + + raw = MagicMock(spec=httpx.Response) + raw.status_code = status_code + raw.text = "not authorized" + raw.headers = httpx.Headers({}) + mock_handler.call_api.side_effect = DataError("auth error", raw) + with patch("importlib.import_module", return_value=mock_handler): with pytest.raises(TTDApiError) as exc_info: client._call_api(_CONTEXT, rows, batch_index=2) - assert exc_info.value.status_code == 400 + assert exc_info.value.error_code == expected_error_code assert exc_info.value.batch_index == 2 - assert "not configured" in exc_info.value.response_text + assert "not authorized" in exc_info.value.response_text + + +def test_response_validation_failure_fails_only_its_own_batch(): + # Schema drift: the server returns 200 but the body doesn't match the SDK's model. + client = _make_client() + mock_handler = _make_mock_handler() + mock_handler.build_items.return_value = [MagicMock()] + + raw = MagicMock(spec=httpx.Response) + raw.status_code = 200 + raw.text = '{"FailedLines": "not-a-list"}' + raw.headers = httpx.Headers({}) + mock_handler.call_api.side_effect = ResponseValidationError( + "Response validation failed", raw, ValueError("type mismatch"), body=raw.text + ) + + with patch("importlib.import_module", return_value=mock_handler): + results = client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=3) + assert len(results) == 1 + assert results[0]["success"] is False + assert results[0]["error_code"] == "ResponseValidationError" -def test_unexpected_exception_from_handler_raises_ttd_api_error_with_message(): +def test_unexpected_exception_is_reported_named_after_the_failure(): client = _make_client() mock_handler = _make_mock_handler() mock_handler.build_items.return_value = [MagicMock()] - mock_handler.call_api.side_effect = ValueError("unexpected error") + mock_handler.call_api.side_effect = ValueError("bug in the SDK") with patch("importlib.import_module", return_value=mock_handler): - with pytest.raises(TTDApiError) as exc_info: - client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=5) + results = client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=5) + + assert len(results) == 1 + assert results[0]["success"] is False + assert results[0]["error_code"] == "ValueError" + assert results[0]["error_message"] == "ValueError: bug in the SDK" + + +def test_build_items_failure_fails_only_its_own_batch(): + # A malformed row makes build_items raise. Nothing was sent, and the failure is specific + # to this batch's own rows, so it must fail just this batch rather than raising and + # aborting every later batch too. + client = _make_client() + mock_handler = _make_mock_handler() + mock_handler.build_items.side_effect = ValueError("id_type 'Banana' is not supported") + + with patch("importlib.import_module", return_value=mock_handler): + results = client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=1) + + mock_handler.call_api.assert_not_called() + assert len(results) == 1 + assert results[0]["success"] is False + assert results[0]["error_code"] == "ValueError" + assert "not supported" in results[0]["error_message"] - assert exc_info.value.batch_index == 5 - assert "unexpected error" in exc_info.value.response_text + +def test_non_standard_status_code_does_not_raise_out_of_the_handler(): + # Load balancers and proxies (AWS ALB 460/463/464, nginx 499) return codes HTTPStatus() + # rejects. The phrase lookup must not blow up and escape as an unhandled ValueError. + client = _make_client() + mock_handler = _make_mock_handler() + mock_handler.build_items.return_value = [MagicMock()] + + raw = MagicMock(spec=httpx.Response) + raw.status_code = 520 + raw.text = "web server returned an unknown error" + raw.headers = httpx.Headers({}) + mock_handler.call_api.side_effect = DataError("unknown error", raw, body=raw.text) + + with patch("importlib.import_module", return_value=mock_handler): + results = client._call_api(_CONTEXT, _make_rows(_ROW), batch_index=0) + + # Not a 401/403, so this batch fails and later batches still run. + assert results[0]["success"] is False + assert results[0]["error_code"] == "520" diff --git a/tests/unit/test_exceptions.py b/tests/unit/test_exceptions.py index 120c5ac..bad4a9a 100644 --- a/tests/unit/test_exceptions.py +++ b/tests/unit/test_exceptions.py @@ -13,9 +13,3 @@ def test_all_sdk_exceptions_inherit_ttd_error(): assert issubclass(TTDApiError, TTDError) assert issubclass(TTDConfigurationError, TTDError) assert issubclass(TTDSchemaValidationError, TTDError) - - -def test_api_error_none_status_code_shows_no_response_in_message(): - # status_code=None means no HTTP response was received - err = TTDApiError(status_code=None, response_text="timeout", batch_index=0) - assert "no response" in str(err) diff --git a/tests/unit/test_process_partitions.py b/tests/unit/test_process_partitions.py index 3e10c8c..3805292 100644 --- a/tests/unit/test_process_partitions.py +++ b/tests/unit/test_process_partitions.py @@ -1,25 +1,27 @@ -"""mapInPandas wiring test for process_partitions. +"""mapInPandas wiring and abort-path tests for process_partitions. Proves that Spark invokes our partition function in a worker, ships rows through -Arrow, and reassembles the output DataFrame with the declared schema. Scope is -strictly the .mapInPandas(...) plumbing — error-handling branches, handler-specific -behaviour, and SDK error mapping are out of scope. +Arrow, and reassembles the output DataFrame with the declared schema, and that a +401/403 aborts the partition instead of failing the job. A local HTTP server is used as the only cross-process-safe way to stand in for the TTD API: mocks in the driver process do not propagate to Spark Python workers. -The server returns 500 to every request so the run terminates deterministically +It returns a configured status to every request so runs terminate deterministically without needing real credentials or wire-format responses. """ from __future__ import annotations import threading +from collections import Counter from collections.abc import Iterator +from http import HTTPStatus from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import pytest from pyspark.sql import SparkSession from pyspark.sql.types import StringType, StructField, StructType, TimestampType +from ttd_data import ClientConfig from ttd_databricks_python.ttd_databricks.batching import process_partitions from ttd_databricks_python.ttd_databricks.contexts import AdvertiserContext @@ -27,20 +29,39 @@ pytestmark = pytest.mark.spark +# Pinned so request counts stay exact: retry_config=None leaves the SDK's retry wrapper +# off, so each batch makes exactly one call even when the stub returns a retryable 5xx. +# Shared by both tests — workers cache one DataClient per process, so a differing config +# in a second test would be silently ignored. +_NO_RETRY_CLIENT_CONFIG = ClientConfig( + server_url=None, + retry_config=None, + timeout_ms=10_000, + uid2_config=None, +) -class _FailingHandler(BaseHTTPRequestHandler): - """Responds 500 to every request. Tracks request count to prove the server was hit.""" +class _StubHandler(BaseHTTPRequestHandler): + """Responds with the configured status to every request. Tracks request count to prove the server was hit.""" + + status_code = 500 request_count = 0 + # 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. + counter_lock = threading.Lock() @classmethod - def reset_count(cls) -> None: - cls.request_count = 0 + def configure(cls, status_code: int) -> None: + with cls.counter_lock: + cls.status_code = status_code + cls.request_count = 0 def do_POST(self) -> None: # noqa: N802 — required by stdlib BaseHTTPRequestHandler - type(self).request_count += 1 - body = b'{"Message":"forced server error for test"}' - self.send_response(500) + with type(self).counter_lock: + type(self).request_count += 1 + body = b'{"Message":"forced error for test"}' + self.send_response(type(self).status_code) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) self.end_headers() @@ -51,9 +72,9 @@ def log_message(self, format: str, *args: object) -> None: # noqa: A002 — bas @pytest.fixture(scope="module") -def failing_server() -> Iterator[str]: - """Run a localhost HTTP server that 500s on every request. Yields base URL.""" - server = ThreadingHTTPServer(("127.0.0.1", 0), _FailingHandler) +def stub_server() -> Iterator[str]: + """Run a localhost HTTP server standing in for the TTD API. Yields base URL.""" + server = ThreadingHTTPServer(("127.0.0.1", 0), _StubHandler) thread = threading.Thread(target=server.serve_forever, daemon=True) thread.start() try: @@ -77,16 +98,16 @@ def _advertiser_input_schema() -> StructType: ) -def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, failing_server: str) -> None: +def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, stub_server: str) -> None: """mapInPandas invokes our partition function in a worker, preserves row count, preserves input values through Arrow, and returns the declared schema.""" input_ids = [f"id-{i}" for i in range(7)] rows = [("TDID", id_value, "seg-a", None, None) for id_value in input_ids] input_df = spark.createDataFrame(rows, _advertiser_input_schema()) output_schema = get_output_schema(input_df.schema) - context = AdvertiserContext(advertiser_id="adv-test", base_url_override=failing_server) + context = AdvertiserContext(advertiser_id="adv-test", base_url_override=stub_server) - _FailingHandler.reset_count() + _StubHandler.configure(500) result_df = process_partitions( df=input_df, batch_size=3, @@ -94,14 +115,84 @@ def test_mapinpandas_wires_up_and_round_trips(spark: SparkSession, failing_serve api_token="not-a-real-token", context=context, parallelism=2, + client_config=_NO_RETRY_CLIENT_CONFIG, ) result_rows = result_df.collect() # 1. Our partition function executed in a Spark worker as expected: 7 rows / batch_size=3 → 3 batches. - assert _FailingHandler.request_count == 3 + assert _StubHandler.request_count == 3 # 2. Row count round-trips through the Arrow pipeline. assert len(result_rows) == 7 # 3. Declared output schema columns are present (loose check — avoid nullable/metadata flakes). 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) + + +@pytest.mark.parametrize( + "status_code", + [ + 401, # missing / expired TTD auth token + 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: + """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)] + input_df = spark.createDataFrame(rows, _advertiser_input_schema()) + output_schema = get_output_schema(input_df.schema) + context = AdvertiserContext(advertiser_id="adv-test", base_url_override=stub_server) + + _StubHandler.configure(status_code) + result_df = process_partitions( + 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, + ) + result_rows = result_df.collect() + + # Single partition: only the first batch is sent; the remaining two never leave the worker. + assert _StubHandler.request_count == 1 + # No rows discarded — every input row is accounted for. + assert len(result_rows) == 7 + assert all(row["success"] is False for row in result_rows) + + by_code = Counter(row["error_code"] for row in result_rows) + # The 3 rows that were sent carry the server's own status, not ABORTED. + assert by_code[HTTPStatus(status_code).phrase] == 3 + # The 4 rows after them were never submitted, so they are safe to re-run. + assert by_code["ABORTED"] == 4 + + +def test_other_4xx_fails_only_its_own_batch(spark: SparkSession, stub_server: str) -> None: + """A 4xx other than 401/403 fails only its own batch; the partition keeps calling the API.""" + status_code = 400 # malformed request body + rows = [("TDID", f"id-{i}", "seg-a", None, None) for i in range(7)] + input_df = spark.createDataFrame(rows, _advertiser_input_schema()) + output_schema = get_output_schema(input_df.schema) + context = AdvertiserContext(advertiser_id="adv-test", base_url_override=stub_server) + + _StubHandler.configure(status_code) + result_df = process_partitions( + 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, + ) + result_rows = result_df.collect() + + # All 3 batches are attempted since the failure doesn't abort the partition. + assert _StubHandler.request_count == 3 + assert len(result_rows) == 7 + assert all(row["success"] is False for row in result_rows) + + by_code = Counter(row["error_code"] for row in result_rows) + assert by_code[HTTPStatus(status_code).phrase] == 7 diff --git a/tests/unit/test_push_data.py b/tests/unit/test_push_data.py index 06ccc5a..58233bf 100644 --- a/tests/unit/test_push_data.py +++ b/tests/unit/test_push_data.py @@ -130,6 +130,67 @@ def test_partial_failure_maps_error_to_correct_row(spark: SparkSession) -> None: assert by_id["def456"]["success"] is True +def test_400_error_fails_only_its_own_batch_and_continues(spark: SparkSession) -> None: + # batch_size=1 -> 3 separate API calls: first succeeds, second hits a 400 (fails just + # that batch), third still runs and succeeds. + import httpx + from ttd_data.errors import DataError + + data = [("TDID", "ok", "seg1"), ("TDID", "bad-request", "seg2"), ("TDID", "later-succeeds", "seg3")] + df = spark.createDataFrame(data, _REQUIRED_SCHEMA) + + raw = MagicMock(spec=httpx.Response) + raw.status_code = 400 + raw.text = "invalid segment" + raw.headers = httpx.Headers({}) + client_error = DataError("bad request", raw) + + mock_handler = _make_handler() + mock_handler.call_api.side_effect = [([], {}), client_error, ([], {})] + + with patch("importlib.import_module", return_value=mock_handler): + result = _make_client(spark).push_data(df, _CONTEXT, batch_size=1) + + by_id = {r["id_value"]: r for r in result.collect()} + assert len(by_id) == 3 + assert by_id["ok"]["success"] is True + assert by_id["bad-request"]["success"] is False + assert by_id["bad-request"]["error_code"] == "Bad Request" + assert by_id["later-succeeds"]["success"] is True # later batch still ran + + +@pytest.mark.parametrize("status_code", [401, 403]) +def test_401_or_403_aborts_remaining_batches(spark: SparkSession, status_code: int) -> None: + # batch_size=1 -> 3 separate API calls: first succeeds, second hits a 401/403 (stops + # the run), third is never attempted. The successful row survives, the rejected row + # keeps the server's status, and only the untried row is ABORTED. + import httpx + from ttd_data.errors import DataError + + data = [("TDID", "ok", "seg1"), ("TDID", "unauthorized", "seg2"), ("TDID", "never-attempted", "seg3")] + df = spark.createDataFrame(data, _REQUIRED_SCHEMA) + + raw = MagicMock(spec=httpx.Response) + raw.status_code = status_code + raw.text = "not authorized" + raw.headers = httpx.Headers({}) + client_error = DataError("auth error", raw) + + mock_handler = _make_handler() + mock_handler.call_api.side_effect = [([], {}), client_error] + + with patch("importlib.import_module", return_value=mock_handler): + result = _make_client(spark).push_data(df, _CONTEXT, batch_size=1) + + by_id = {r["id_value"]: r for r in result.collect()} + assert len(by_id) == 3 + assert by_id["ok"]["success"] is True # earlier batch's result is not discarded + assert by_id["unauthorized"]["success"] is False + assert "not authorized" in by_id["unauthorized"]["error_message"] # sent and rejected + assert by_id["never-attempted"]["success"] is False + assert by_id["never-attempted"]["error_code"] == "ABORTED" # never sent + + def test_missing_required_column_raises_schema_validation_error(spark: SparkSession) -> None: schema = StructType( [ diff --git a/ttd_databricks_python/ttd_databricks/batching.py b/ttd_databricks_python/ttd_databricks/batching.py index 7eb909c..dfdedc0 100644 --- a/ttd_databricks_python/ttd_databricks/batching.py +++ b/ttd_databricks_python/ttd_databricks/batching.py @@ -60,6 +60,9 @@ def process_partitions( client_config is a snapshot of the driver DataClient's settings (server_url, retry_config, timeout_ms, uid2_config), used to rebuild an equivalent DataClient per worker. + + Does not raise. An auth or permission failure aborts its own partition — later rows there + get error_code="ABORTED", never submitted and safe to re-run. Other partitions carry on. """ if parallelism is None: try: @@ -72,17 +75,16 @@ def process_partitions( handler_module = context.endpoint.handler_module def partition_to_results(pandas_df_iter: Iterable[pd.DataFrame]) -> Iterator[pd.DataFrame]: - import http import importlib from datetime import datetime, timezone - import httpx import pandas as pd from ttd_data import DataClient - from ttd_data.errors import DataError, NoResponseError + from ttd_databricks_python.ttd_databricks.constants import ABORTED_ERROR_CODE, DEFAULT_RETRY_CONFIG from ttd_databricks_python.ttd_databricks.utils import ( attach_resolutions, + classify_failure, empty_resolution_value, parse_failed_lines, ) @@ -92,62 +94,71 @@ def partition_to_results(pandas_df_iter: Iterable[pd.DataFrame]) -> Iterator[pd. # Workers rebuild the client from the picklable client_config snapshot; # DataClient itself can't be cloudpickled. if client_config is None: - _worker_client = DataClient(timeout_ms=10_000) + _worker_client = DataClient(timeout_ms=10_000, retry_config=DEFAULT_RETRY_CONFIG) else: _worker_client = DataClient.from_config(client_config) client = _worker_client handler = importlib.import_module(handler_module) + # Why this partition stopped. Once set, later batches are never sent to The Trade Desk. + abort_reason: Optional[str] = None + + def build_result_df( + batch_rows: list[dict[str, Any]], + timestamp: datetime, + row_results: list[dict[str, Any]], + ) -> pd.DataFrame: + merged = [ + {**row_dict, **row_result, "processed_timestamp": timestamp} + for row_dict, row_result in zip(batch_rows, row_results, strict=True) + ] + return pd.DataFrame(merged, columns=output_field_names) + + def failed_batch(batch_rows: list[dict[str, Any]], error_code: str, error_message: str) -> pd.DataFrame: + """Mark every row in the batch with the same failure.""" + row_results = [ + { + "success": False, + "error_code": error_code, + "error_message": error_message, + **empty_resolution_value(), + } + for _ in batch_rows + ] + return build_result_df(batch_rows, datetime.now(timezone.utc), row_results) + def call_batch(batch_rows: list[dict[str, Any]]) -> pd.DataFrame: timestamp = datetime.now(timezone.utc) - items = handler.build_items(batch_rows) - raw_pii_ids_per_row = handler.collect_raw_pii_ids_per_row(batch_rows) - - def fail_batch(error_code: str | None, error_message: str) -> pd.DataFrame: - results = [ - { - **row_dict, - "success": False, - "error_code": error_code, - "error_message": error_message, - "processed_timestamp": timestamp, - **empty_resolution_value(), - } - for row_dict in batch_rows - ] - return pd.DataFrame(results, columns=output_field_names) - - failed_lines: list[Any] = [] - identity_resolutions: dict[str, Any] = {} + + def abort(error_code: str, error_message: str) -> pd.DataFrame: + """Record this batch's own outcome, then stop sending the rest of the partition.""" + nonlocal abort_reason + abort_reason = error_message + return failed_batch(batch_rows, error_code, error_message) + try: + items = handler.build_items(batch_rows) + raw_pii_ids_per_row = handler.collect_raw_pii_ids_per_row(batch_rows) failed_lines, identity_resolutions = handler.call_api( client, context, items, api_token, data_load_trace_id ) - except ( - httpx.TimeoutException, - httpx.RemoteProtocolError, - NoResponseError, - ) as exc: - # Transient: timeout, stale pooled connection, or no response. - # Mark batch as failed and continue. - return fail_batch(None, str(exc)) - except DataError as exc: - error_code = http.HTTPStatus(exc.status_code).phrase - if exc.status_code >= 500: - # Transient server error, mark batch as failed and continue. - return fail_batch(error_code, exc.body) - # 4xx errors (auth, bad request) — fail the job. - raise RuntimeError(f"TTD API unrecoverable error: {exc}") from exc + row_results = parse_failed_lines(failed_lines, len(batch_rows)) + attach_resolutions(row_results, raw_pii_ids_per_row, identity_resolutions) except Exception as exc: - raise RuntimeError(f"Unexpected error during API call: {exc}") from exc - - row_results = parse_failed_lines(failed_lines, len(batch_rows)) - attach_resolutions(row_results, raw_pii_ids_per_row, identity_resolutions) - merged = [ - {**row_dict, **row_result, "processed_timestamp": timestamp} - for row_dict, row_result in zip(batch_rows, row_results, strict=True) - ] - return pd.DataFrame(merged, columns=output_field_names) + transient, error_code, error_message = classify_failure(exc) + if transient: + return failed_batch(batch_rows, error_code, error_message) + return abort(error_code, error_message) + + return build_result_df(batch_rows, timestamp, row_results) + + def process_batch(batch_rows: list[dict[str, Any]]) -> pd.DataFrame: + """Call the API, unless an earlier batch already aborted this partition.""" + if abort_reason is not None: + return failed_batch( + batch_rows, ABORTED_ERROR_CODE, f"Aborted batch due to unrecoverable error: {abort_reason}" + ) + return call_batch(batch_rows) batch: list[dict[str, Any]] = [] for pandas_df in pandas_df_iter: @@ -158,9 +169,9 @@ def fail_batch(error_code: str | None, error_message: str) -> pd.DataFrame: for row_dict in cast(list[dict[str, Any]], normalised_df.to_dict(orient="records")): batch.append(row_dict) if len(batch) == batch_size: - yield call_batch(batch) + yield process_batch(batch) batch = [] if batch: - yield call_batch(batch) + yield process_batch(batch) return df.select(*all_input_cols).repartition(parallelism).mapInPandas(partition_to_results, schema=output_schema) diff --git a/ttd_databricks_python/ttd_databricks/constants.py b/ttd_databricks_python/ttd_databricks/constants.py index 6879620..43bda2f 100644 --- a/ttd_databricks_python/ttd_databricks/constants.py +++ b/ttd_databricks_python/ttd_databricks/constants.py @@ -1,5 +1,23 @@ """Package-level constants for the TTD Databricks SDK.""" +from ttd_data.utils import BackoffStrategy, RetryConfig + # DataOrigin ID automatically appended to every API call to identify # data submitted via this SDK. TTD_DATABRICKS_SDK_ORIGIN_ID = "ttd_databricks_sdk" + +# Error Code for rows never sent to The Trade Desk, so re-running them is safe. Failures +# that may already have been ingested are named after their exception instead. +ABORTED_ERROR_CODE = "ABORTED" + +DEFAULT_RETRY_CONFIG = RetryConfig( + "backoff", + BackoffStrategy( + initial_interval=500, + max_interval=8_000, + exponent=2.0, + max_elapsed_time=30_000, + jitter_ms=250, + ), + retry_connection_errors=True, +) diff --git a/ttd_databricks_python/ttd_databricks/exceptions.py b/ttd_databricks_python/ttd_databricks/exceptions.py index d000968..7e39f39 100644 --- a/ttd_databricks_python/ttd_databricks/exceptions.py +++ b/ttd_databricks_python/ttd_databricks/exceptions.py @@ -2,8 +2,6 @@ from __future__ import annotations -from typing import Optional - class TTDError(Exception): """Base exception for all TTD Databricks SDK errors.""" @@ -12,14 +10,13 @@ class TTDError(Exception): class TTDApiError(TTDError): - """Raised when the TTD API returns a non-2xx status for an entire batch request.""" + """Raised when a batch hits a failure no later batch could survive, so the run must stop.""" - def __init__(self, status_code: Optional[int], response_text: str, batch_index: int) -> None: - self.status_code = status_code + def __init__(self, response_text: str, batch_index: int, error_code: str) -> None: self.response_text = response_text self.batch_index = batch_index - status_str = str(status_code) if status_code is not None else "no response" - super().__init__(f"TTD API error (HTTP {status_str}) for batch {batch_index}: {response_text}") + self.error_code = error_code + super().__init__(f"TTD API error ({error_code}) for batch {batch_index}: {response_text}") class TTDConfigurationError(TTDError): diff --git a/ttd_databricks_python/ttd_databricks/ttd_client.py b/ttd_databricks_python/ttd_databricks/ttd_client.py index ee80624..e24e086 100644 --- a/ttd_databricks_python/ttd_databricks/ttd_client.py +++ b/ttd_databricks_python/ttd_databricks/ttd_client.py @@ -13,6 +13,7 @@ from ttd_data.uid2 import UID2Config from ttd_data.utils import RetryConfig +from ttd_databricks_python.ttd_databricks.constants import ABORTED_ERROR_CODE, DEFAULT_RETRY_CONFIG from ttd_databricks_python.ttd_databricks.contexts import TTDContext from ttd_databricks_python.ttd_databricks.endpoints import TTDEndpoint @@ -66,7 +67,7 @@ def from_params( api_token: str, spark: Optional[SparkSession] = None, uid2_config: Optional[UID2Config] = None, - retry_config: OptionalNullable[RetryConfig] = None, + retry_config: OptionalNullable[RetryConfig] = DEFAULT_RETRY_CONFIG, server_url: Optional[str] = None, timeout_ms: Optional[int] = None, ) -> TtdDatabricksClient: @@ -80,6 +81,7 @@ def from_params( - uid2_config: Optional. Enables client-side resolution of raw PII identifiers (Email/Phone/HashedEmail/HashedPhone) to UID2/EUID. - retry_config: Optional. Retry behavior for transient API errors (429/5xx). + Defaults to DEFAULT_RETRY_CONFIG; pass None to disable retries. - server_url: Optional. Override the default TTD Data API server URL. - timeout_ms: Optional. Per-request timeout in milliseconds. @@ -112,6 +114,9 @@ def push_data( Returns: Original columns + success, error_code, error_message, processed_timestamp, and `uid2_resolutions` (array, empty unless `uid2_config` was provided). + Does not raise; results already collected are always returned. An auth or permission + failure stops the run — later rows get error_code="ABORTED", safe to re-run. + - df: Input Spark DataFrame. Must contain all non-nullable columns for context.endpoint. Nullable columns may be omitted — they will be filled with null automatically. Extra columns are preserved in the output but ignored during API submission. @@ -121,7 +126,9 @@ def push_data( - data_load_trace_id: Optional trace ID passed for debugging. Passed as DataLoadTraceId in the API request body. If None, omitted from the request. """ + from ttd_databricks_python.ttd_databricks.exceptions import TTDApiError from ttd_databricks_python.ttd_databricks.schemas import get_output_schema, validate_ttd_schema + from ttd_databricks_python.ttd_databricks.utils import empty_resolution_value df = self._fill_nullable_columns(df, context.endpoint) validate_ttd_schema(df, context.endpoint) @@ -130,10 +137,34 @@ def push_data( all_rows = df.collect() result_rows: list[dict[str, Any]] = [] + def record_failed(rows: list[Row], error_code: str, error_message: str, timestamp: datetime) -> None: + for row in rows: + merged = row.asDict() + merged.update( + success=False, + error_code=error_code, + error_message=error_message, + processed_timestamp=timestamp, + **empty_resolution_value(), + ) + result_rows.append(merged) + for batch_index, i in enumerate(range(0, len(all_rows), batch_size)): batch = all_rows[i : i + batch_size] timestamp = datetime.now(timezone.utc) - api_results = self._call_api(context, batch, batch_index, data_load_trace_id) + + try: + api_results = self._call_api(context, batch, batch_index, data_load_trace_id) + except TTDApiError as exc: + # Attempted, so it keeps its own error; everything after it is unsent, so ABORTED. + record_failed(batch, exc.error_code, exc.response_text, timestamp) + record_failed( + all_rows[i + len(batch) :], + ABORTED_ERROR_CODE, + f"Aborted batch due to unrecoverable error: {exc}", + timestamp, + ) + break for row, result in zip(batch, api_results, strict=True): merged = row.asDict() @@ -164,6 +195,9 @@ def batch_process( Updates metadata table if provided. + Does not raise. An auth or permission failure aborts its own partition — later rows + there get error_code="ABORTED", safe to re-run. Other partitions carry on. + - context: Typed context object (AdvertiserContext, ThirdPartyContext, etc.) Contains endpoint-specific config (data_provider_id, advertiser_id, etc.) - input_table: Delta table name to read from. Must contain all mandatory columns for @@ -392,36 +426,32 @@ def _call_api( Delegates item-building and the API call to the endpoint-specific handler module, then applies shared failed_lines parsing to produce per-row result dicts with keys: success (bool), error_code (Optional[str]), error_message (Optional[str]), - and `uid2_resolutions` (list[dict], empty on transient/4xx errors). + and `uid2_resolutions` (list[dict], empty on every failure path). - rows: Spark Rows from df.collect(); must contain the mandatory columns for context.endpoint. - batch_index: Zero-based batch number used in TTDApiError if the call fails. - Transient errors (timeouts, stale connections, no response, server 5xx) return - all-failed results so the caller can continue with remaining batches. + Any failure other than auth or permission fails just this batch, so the caller can + continue with the remaining ones. Raises: - TTDApiError: On unrecoverable errors (4xx client errors, unexpected exceptions). + TTDApiError: On an auth or permission failure, carrying the error_code the caller + should label this batch's rows with. """ - import http import importlib - import httpx - from ttd_data.errors import DataError, NoResponseError - from ttd_databricks_python.ttd_databricks.exceptions import TTDApiError from ttd_databricks_python.ttd_databricks.utils import ( attach_resolutions, + classify_failure, empty_resolution_value, parse_failed_lines, ) handler = importlib.import_module(context.endpoint.handler_module) rows_data = [row.asDict() for row in rows] - items = handler.build_items(rows_data) - raw_pii_ids_per_row = handler.collect_raw_pii_ids_per_row(rows_data) - def fail_all(error_code: str | None, error_message: str) -> list[dict[str, Any]]: + def fail_all(error_code: str, error_message: str) -> list[dict[str, Any]]: return [ { "success": False, @@ -432,37 +462,24 @@ def fail_all(error_code: str | None, error_message: str) -> list[dict[str, Any]] for _ in rows ] - failed_lines: list[Any] = [] - identity_resolutions: dict[str, Any] = {} try: + items = handler.build_items(rows_data) + raw_pii_ids_per_row = handler.collect_raw_pii_ids_per_row(rows_data) failed_lines, identity_resolutions = handler.call_api( self._data_api_client, context, items, self._api_token, data_load_trace_id ) - except ( - httpx.TimeoutException, - httpx.RemoteProtocolError, - NoResponseError, - ) as exc: - return fail_all(None, str(exc)) - except DataError as exc: - error_code = http.HTTPStatus(exc.status_code).phrase - if exc.status_code >= 500: - # Transient server error, mark batch as failed and continue. - return fail_all(error_code, exc.body) - raise TTDApiError( - status_code=exc.status_code, - response_text=exc.body, - batch_index=batch_index, - ) from exc + results = parse_failed_lines(failed_lines, len(rows)) + attach_resolutions(results, raw_pii_ids_per_row, identity_resolutions) except Exception as exc: + transient, error_code, error_message = classify_failure(exc) + if transient: + return fail_all(error_code, error_message) raise TTDApiError( - status_code=None, - response_text=str(exc), + response_text=error_message, batch_index=batch_index, + error_code=error_code, ) from exc - results = parse_failed_lines(failed_lines, len(rows)) - attach_resolutions(results, raw_pii_ids_per_row, identity_resolutions) return results @staticmethod diff --git a/ttd_databricks_python/ttd_databricks/utils.py b/ttd_databricks_python/ttd_databricks/utils.py index ee19e86..a5c03e2 100644 --- a/ttd_databricks_python/ttd_databricks/utils.py +++ b/ttd_databricks_python/ttd_databricks/utils.py @@ -2,6 +2,7 @@ from __future__ import annotations +import http from typing import Any, Optional from ttd_databricks_python.ttd_databricks.schemas import UID2_RESOLUTIONS_COLUMN @@ -12,6 +13,31 @@ def empty_resolution_value() -> dict[str, list[Any]]: return {UID2_RESOLUTIONS_COLUMN: []} +def classify_failure(exc: Exception) -> tuple[bool, str, str]: + """Classify a failed API call as (transient, error_code, error_message). + + Transient failures fail only their own batch; a later batch can still succeed. Auth and + permission failures are not transient — they recur on every batch, so the run stops. + """ + from ttd_data.errors import DataError, ResponseValidationError + + if isinstance(exc, DataError) and not isinstance(exc, ResponseValidationError): + if exc.status_code in (http.HTTPStatus.UNAUTHORIZED, http.HTTPStatus.FORBIDDEN): + return False, _status_label(exc.status_code), exc.body + return True, _status_label(exc.status_code), exc.body + # No HTTP status to report, so the exception name is the code: ReadTimeout, ValueError, etc. + # Never NULL, which would be indistinguishable from a succeeded row. + return True, type(exc).__name__, f"{type(exc).__name__}: {getattr(exc, 'body', exc)}" + + +def _status_label(status_code: int) -> str: + """Label for an HTTP status: its reason phrase, or the bare code if not IANA-registered.""" + try: + return http.HTTPStatus(status_code).phrase + except ValueError: + return str(status_code) + + def parse_failed_lines(failed_lines: list[Any], row_count: int) -> list[dict[str, Any]]: """Map API failed_lines to per-row result dicts with success, error_code, error_message.