diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/cdp_producer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/cdp_producer.py index 5ce64dc20111..ed06e69bd39d 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/cdp_producer.py +++ b/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: diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_cdp_producer.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_cdp_producer.py index a7ed39ce7a8c..5b1dbd333f5e 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test/test_cdp_producer.py +++ b/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 diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_sync.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_sync.py index 119f081875c0..df478d73f835 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/pipeline_sync.py +++ b/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: