Retry transient failures in Xylem Data Lake adapter

Vendor-side instability this past week killed standard DAG runs in two
places that previously had no retry:
  - 5xx from /api/v1/sqllab/execute/ during chunk queries (fatal at
    raise_for_status)
  - 401 from /api/v1/me/ during initial Superset login (fatal in
    _create_superset_session)

Adds retry-with-backoff (10s/20s/40s, 3 attempts) at both points, plus
the same retry path for ConnectionError and Timeout exceptions on the
query POST. Re-auths triggered mid-query also flow through the new
auth-retry wrapper.

Author: dmarulli

amiadapters/adapters/xylem_datalake.py

@@ -91,6 +91,8 @@ class DatalakeWaterRegister:
 DELAY_SECONDS = 2
 MAX_AUTH_ATTEMPTS = 3
 AUTH_BACKOFF_BASE = 10
+MAX_SERVER_ATTEMPTS = 3
+SERVER_BACKOFF_BASE = 10
 CSRF_REFRESH_INTERVAL = 10
 
 
@@ -204,6 +206,39 @@ def _create_superset_session(
     return session
 
 
+def _create_superset_session_with_retry(
+    datalake_url: str,
+    superset_url: str,
+    client_id: str,
+    username: str,
+    password: str,
+) -> requests.Session:
+    """Wrap _create_superset_session with retry-with-backoff on transient failures.
+
+    Vendor-side flakes have produced 401 from /api/v1/me/ even after a successful
+    Keycloak handshake. A short backoff and retry gets through most of these.
+    """
+    last_exc = None
+    for attempt in range(MAX_AUTH_ATTEMPTS + 1):
+        if attempt > 0:
+            delay = AUTH_BACKOFF_BASE * (2 ** (attempt - 1))
+            logger.info(
+                f"Backing off {delay}s before auth retry "
+                f"(attempt {attempt + 1}/{MAX_AUTH_ATTEMPTS + 1})"
+            )
+            time.sleep(delay)
+        try:
+            return _create_superset_session(
+                datalake_url, superset_url, client_id, username, password
+            )
+        except (RuntimeError, requests.exceptions.RequestException) as e:
+            logger.warning(
+                f"Auth attempt {attempt + 1}/{MAX_AUTH_ATTEMPTS + 1} failed: {e}"
+            )
+            last_exc = e
+    raise last_exc
+
+
 def _refresh_csrf(session: requests.Session, superset_url: str) -> None:
     """Refresh the CSRF token on an existing session."""
     csrf_resp = session.get(
@@ -291,7 +326,7 @@ def _extract(
         extract_range_start: datetime,
         extract_range_end: datetime,
     ) -> ExtractOutput:
-        self._session = _create_superset_session(
+        self._session = _create_superset_session_with_retry(
             self.datalake_url,
             self.superset_url,
             self.client_id,
@@ -365,6 +400,7 @@ def _query_chunked(
     def _query(self, sql: str) -> List[dict]:
         """Execute a SQL query via the Superset SQL Lab API with session management."""
         auth_failures = 0
+        server_failures = 0
 
         while True:
             # Refresh CSRF periodically
@@ -393,6 +429,36 @@ def _query(self, sql: str) -> List[dict]:
                 self._session = self._reauth(auth_failures - 1)
                 self._requests_since_csrf = 0
                 continue
+            except (
+                requests.exceptions.ConnectionError,
+                requests.exceptions.Timeout,
+            ) as e:
+                server_failures += 1
+                logger.warning(f"Connection error ({type(e).__name__}): {e}")
+                if server_failures > MAX_SERVER_ATTEMPTS:
+                    raise
+                delay = SERVER_BACKOFF_BASE * (2 ** (server_failures - 1))
+                logger.info(
+                    f"Backing off {delay}s before retry "
+                    f"(attempt {server_failures}/{MAX_SERVER_ATTEMPTS})"
+                )
+                time.sleep(delay)
+                continue
+
+            if 500 <= resp.status_code < 600:
+                server_failures += 1
+                logger.warning(
+                    f"Server error HTTP {resp.status_code}: {resp.text[:300]}"
+                )
+                if server_failures > MAX_SERVER_ATTEMPTS:
+                    resp.raise_for_status()
+                delay = SERVER_BACKOFF_BASE * (2 ** (server_failures - 1))
+                logger.info(
+                    f"Backing off {delay}s before retry "
+                    f"(attempt {server_failures}/{MAX_SERVER_ATTEMPTS})"
+                )
+                time.sleep(delay)
+                continue
 
             if _session_expired(resp):
                 auth_failures += 1
@@ -413,6 +479,7 @@ def _query(self, sql: str) -> List[dict]:
             resp.raise_for_status()
             result = resp.json()
             auth_failures = 0
+            server_failures = 0
 
             if result.get("status") == "error":
                 msg = result.get("error") or result.get("message") or "unknown"
@@ -433,7 +500,7 @@ def _reauth(self, consecutive_failures: int) -> requests.Session:
             time.sleep(delay)
         else:
             logger.info("Re-authenticating")
-        return _create_superset_session(
+        return _create_superset_session_with_retry(
             self.datalake_url,
             self.superset_url,
             self.client_id,

test/amiadapters/test_xylem_datalake.py

@@ -1,8 +1,14 @@
 import json
+from unittest.mock import MagicMock, patch
+
+import requests
 
 from amiadapters.outputs.base import ExtractOutput
 from amiadapters.models import DataclassJSONEncoder, GeneralMeter, GeneralMeterRead
-from amiadapters.adapters.xylem_datalake import XylemDatalakeAdapter
+from amiadapters.adapters.xylem_datalake import (
+    XylemDatalakeAdapter,
+    _create_superset_session_with_retry,
+)
 from test.base_test_case import BaseTestCase
 
 
@@ -235,3 +241,114 @@ def test_parse_flowtime_with_space(self):
 
         result = self.adapter._parse_flowtime("2026-04-14 08:00:00")
         self.assertEqual(result, datetime(2026, 4, 14, 8, 0, 0))
+
+
+def _query_response(
+    status_code, json_data=None, text="", content_type="application/json"
+):
+    resp = MagicMock(spec=requests.Response)
+    resp.status_code = status_code
+    resp.headers = {"Content-Type": content_type}
+    resp.text = text
+    resp.json.return_value = json_data if json_data is not None else {}
+    if 400 <= status_code < 600:
+        resp.raise_for_status.side_effect = requests.exceptions.HTTPError(
+            f"{status_code} Error", response=resp
+        )
+    else:
+        resp.raise_for_status.return_value = None
+    return resp
+
+
+class TestXylemDatalakeQueryRetry(BaseTestCase):
+    """Tests for retry behavior in _query and the auth retry wrapper."""
+
+    def setUp(self):
+        self.adapter = XylemDatalakeAdapter(
+            org_id="test_org",
+            org_timezone="America/Los_Angeles",
+            pipeline_configuration=None,
+            configured_task_output_controller=self.TEST_TASK_OUTPUT_CONTROLLER_CONFIGURATION,
+            configured_metrics=self.TEST_METRICS_CONFIGURATION,
+            agency_code="hlsbo",
+            database_id=1,
+            client_id="test-client-id",
+            username="test",
+            password="test",
+            configured_sinks=[],
+        )
+        self.adapter._session = MagicMock()
+        self.adapter._requests_since_csrf = 0
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    def test_query_retries_on_5xx_then_succeeds(self, _mock_sleep):
+        success = _query_response(
+            200, json_data={"status": "success", "data": [{"x": 1}]}
+        )
+        self.adapter._session.post.side_effect = [
+            _query_response(500, text="boom"),
+            _query_response(502, text="bad gateway"),
+            success,
+        ]
+        rows = self.adapter._query("SELECT 1")
+        self.assertEqual(rows, [{"x": 1}])
+        self.assertEqual(self.adapter._session.post.call_count, 3)
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    def test_query_raises_after_max_5xx_attempts(self, _mock_sleep):
+        self.adapter._session.post.return_value = _query_response(500, text="boom")
+        with self.assertRaises(requests.exceptions.HTTPError):
+            self.adapter._query("SELECT 1")
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    def test_query_retries_on_connection_error_then_succeeds(self, _mock_sleep):
+        success = _query_response(
+            200, json_data={"status": "success", "data": [{"x": 1}]}
+        )
+        self.adapter._session.post.side_effect = [
+            requests.exceptions.ConnectionError("dns failure"),
+            requests.exceptions.Timeout("read timeout"),
+            success,
+        ]
+        rows = self.adapter._query("SELECT 1")
+        self.assertEqual(rows, [{"x": 1}])
+        self.assertEqual(self.adapter._session.post.call_count, 3)
+
+
+class TestCreateSupersetSessionWithRetry(BaseTestCase):
+    """Tests for the auth retry wrapper."""
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    @patch("amiadapters.adapters.xylem_datalake._create_superset_session")
+    def test_returns_session_on_first_success(self, mock_create, _mock_sleep):
+        sentinel = MagicMock(name="session")
+        mock_create.return_value = sentinel
+        result = _create_superset_session_with_retry(
+            "https://dl", "https://sup", "client", "user", "pw"
+        )
+        self.assertIs(result, sentinel)
+        self.assertEqual(mock_create.call_count, 1)
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    @patch("amiadapters.adapters.xylem_datalake._create_superset_session")
+    def test_retries_then_succeeds(self, mock_create, _mock_sleep):
+        sentinel = MagicMock(name="session")
+        mock_create.side_effect = [
+            RuntimeError("Superset login returned unexpected shape (status 401)"),
+            RuntimeError("Superset login returned unexpected shape (status 401)"),
+            sentinel,
+        ]
+        result = _create_superset_session_with_retry(
+            "https://dl", "https://sup", "client", "user", "pw"
+        )
+        self.assertIs(result, sentinel)
+        self.assertEqual(mock_create.call_count, 3)
+
+    @patch("amiadapters.adapters.xylem_datalake.time.sleep")
+    @patch("amiadapters.adapters.xylem_datalake._create_superset_session")
+    def test_raises_after_max_attempts(self, mock_create, _mock_sleep):
+        mock_create.side_effect = RuntimeError("persistent failure")
+        with self.assertRaises(RuntimeError):
+            _create_superset_session_with_retry(
+                "https://dl", "https://sup", "client", "user", "pw"
+            )