fix(data-imports): retry transient DB connection drops on two import reads

products/warehouse_sources/backend/temporal/data_imports/pipelines/core/cdp_producer.py

@@ -4,7 +4,10 @@
 import hashlib
 
 from django.conf import settings
-from django.db.utils import OperationalError as DjangoOperationalError
+from django.db.utils import (
+    InterfaceError as DjangoInterfaceError,
+    OperationalError as DjangoOperationalError,
+)
 
 import orjson
 import pyarrow as pa
@@ -19,6 +22,7 @@
 from posthog.kafka_client.routing import KafkaClusterProfile, async_producer_scope
 from posthog.kafka_client.topics import KAFKA_DWH_CDP_RAW_TABLE
 from posthog.sync import database_sync_to_async_pool
+from posthog.temporal.common.utils import aretry_on_db_connection_drop
 
 from products.cdp.backend.models.hog_functions import HogFunction
 from products.data_warehouse.backend.facade.api import aget_s3_client, ensure_bucket_exists
@@ -97,7 +101,15 @@ def _resolve() -> str:
             raw_table_name = build_table_name(schema.source, schema.name)
             return get_data_warehouse_table_name(schema.source, raw_table_name)
 
-        self._table_name_cache = await _resolve()
+        try:
+            self._table_name_cache = await aretry_on_db_connection_drop(_resolve)
+        except (DjangoOperationalError, DjangoInterfaceError) as e:
+            # Same reasoning as `should_produce_table` below: this reads PostHog's own
+            # database, and a raw connection error stringifies like a customer's
+            # misconfigured source host, which `get_non_retryable_errors` would
+            # misclassify and permanently stop a healthy sync.
+            raise PostHogInternalDatabaseError("Failed to resolve the table name from PostHog's database") from e
+
         return self._table_name_cache
 
     def _build_event_id(self, row: object) -> str:

products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_cdp_producer.py

@@ -7,7 +7,10 @@
 from unittest import mock
 from unittest.mock import MagicMock, patch
 
-from django.db.utils import OperationalError as DjangoOperationalError
+from django.db.utils import (
+    InterfaceError as DjangoInterfaceError,
+    OperationalError as DjangoOperationalError,
+)
 
 import pyarrow as pa
 import pyarrow.parquet as pq
@@ -912,3 +915,58 @@ def test_build_event_id_changes_when_row_data_changes(row_a, row_b):
 def test_build_event_id_changes_with_job_id():
     row = {"id": 1, "name": "Alice"}
     assert _make_producer("job_1")._build_event_id(row) != _make_producer("job_2")._build_event_id(row)
+
+
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.asyncio
+async def test_get_dot_notated_table_name_retries_once_on_connection_drop(team):
+    source = await sync_to_async(ExternalDataSource.objects.create)(
+        team=team, source_type=ExternalDataSourceType.POSTGRES
+    )
+    table = await sync_to_async(DataWarehouseTable.objects.create)(
+        team=team, name="postgres_table_1", external_data_source=source
+    )
+    schema = await sync_to_async(ExternalDataSchema.objects.create)(
+        team=team, name="table_1", source=source, table=table
+    )
+
+    producer = CDPProducer(team_id=team.id, schema_id=str(schema.id), job_id="", logger=mock.AsyncMock())
+
+    real_get = ExternalDataSchema.objects.get
+    calls = []
+
+    def _flaky_get(*args, **kwargs):
+        calls.append(1)
+        if len(calls) == 1:
+            raise DjangoOperationalError("server closed the connection unexpectedly")
+        return real_get(*args, **kwargs)
+
+    with patch.object(ExternalDataSchema.objects, "get", side_effect=_flaky_get):
+        assert await producer.get_dot_notated_table_name() == "postgres.table_1"
+
+    assert len(calls) == 2
+
+
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.asyncio
+@pytest.mark.parametrize("error", [DjangoOperationalError, DjangoInterfaceError])
+async def test_get_dot_notated_table_name_database_failure_stays_retryable(team, error):
+    source = await sync_to_async(ExternalDataSource.objects.create)(
+        team=team, source_type=ExternalDataSourceType.POSTGRES
+    )
+    table = await sync_to_async(DataWarehouseTable.objects.create)(
+        team=team, name="postgres_table_1", external_data_source=source
+    )
+    schema = await sync_to_async(ExternalDataSchema.objects.create)(
+        team=team, name="table_1", source=source, table=table
+    )
+
+    producer = CDPProducer(team_id=team.id, schema_id=str(schema.id), job_id="", logger=mock.AsyncMock())
+
+    with patch.object(ExternalDataSchema.objects, "get", side_effect=error("[Errno -2] Name or service not known")):
+        with pytest.raises(PostHogInternalDatabaseError) as exc_info:
+            await producer.get_dot_notated_table_name()
+
+    non_retryable = PostgresSource().get_non_retryable_errors()
+    error_msg = str(exc_info.value)
+    assert not any(pattern in error_msg for pattern in non_retryable.keys()), error_msg

products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_sync.py

@@ -20,6 +20,7 @@
 from posthog.exceptions_capture import capture_exception
 from posthog.sync import database_sync_to_async_pool
 from posthog.temporal.common.logger import get_logger
+from posthog.temporal.common.utils import aretry_on_db_connection_drop
 
 from products.warehouse_sources.backend.models.external_data_job import ExternalDataJob
 from products.warehouse_sources.backend.models.external_data_schema import (
@@ -119,7 +120,7 @@ def _update():
         schema.last_synced_at = job.created_at
         schema.save()
 
-    await _update()
+    await aretry_on_db_connection_drop(_update)
 
 
 async def set_initial_sync_complete(schema_id: str, team_id: int) -> None: