fix(data-imports): retry transient DB connection drop in repartition flag check

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

@@ -20,6 +20,7 @@
 from structlog.types import FilteringBoundLogger
 
 from posthog.exceptions_capture import capture_exception
+from posthog.temporal.common.utils import retry_on_db_connection_drop
 from posthog.utils import get_machine_id
 
 from products.warehouse_sources.backend.models.external_data_job import ExternalDataJob
@@ -76,7 +77,7 @@ def is_auto_repartition_enabled(schema: ExternalDataSchema) -> bool:
     from posthog.models import Team
 
     try:
-        team = Team.objects.only("uuid", "organization_id").get(id=schema.team_id)
+        team = retry_on_db_connection_drop(lambda: Team.objects.only("uuid", "organization_id").get(id=schema.team_id))
     except Team.DoesNotExist:
         return False
     try:

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

@@ -4,7 +4,7 @@
 import tempfile
 
 import pytest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from django.db import OperationalError
 
@@ -229,6 +229,25 @@ def test_unpartitioned_over_budget_with_keys_enables_partitioning(self, team):
         assert schema.repartition_pending["partition_keys"] == ["id"]
 
 
+class TestIsAutoRepartitionEnabled:
+    def test_retries_once_on_transient_db_connection_drop(self, team):
+        # The Team lookup runs on a long-lived Temporal worker thread; a pooler-dropped connection
+        # raises OperationalError on first use. Without a retry this propagates out of
+        # is_auto_repartition_enabled uncaught (it's outside the function's Team.DoesNotExist/
+        # feature_enabled try blocks) instead of resolving the flag.
+        schema = _make_schema(team, {})
+        mock_queryset = MagicMock()
+        mock_queryset.get.side_effect = [OperationalError("server closed the connection unexpectedly"), team]
+
+        with (
+            patch("posthog.models.Team.objects.only", return_value=mock_queryset),
+            patch.object(ctrl.posthoganalytics, "feature_enabled", return_value=True),
+        ):
+            assert ctrl.is_auto_repartition_enabled(schema) is True
+
+        assert mock_queryset.get.call_count == 2
+
+
 class TestRepartitionOOMHistoryTrigger:
     def _detect(self, team, schema: ExternalDataSchema, delta: deltalake.DeltaTable) -> None:
         async_to_sync(ctrl.maybe_flag_for_repartition)(schema, schema.source, _make_job(team, schema), delta, logger)