From b711fd55ff1964b386df113465676884c6912328 Mon Sep 17 00:00:00 2001 From: Tom Owers Date: Thu, 30 Jul 2026 14:32:30 +0100 Subject: [PATCH] fix(data-imports): retry transient DB connection drop in repartition flag check Evict the stale connection and retry once when the Team lookup in `is_auto_repartition_enabled` hits a transient Postgres connection drop, using the existing `retry_on_db_connection_drop` helper. Adds a regression test covering the retry. Generated-By: PostHog Code Task-Id: 6a029cab-6a4a-4099-a438-75ce9612dc6c --- .../pipelines/core/repartition_controller.py | 3 ++- .../core/test_repartition_controller.py | 21 ++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/repartition_controller.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/repartition_controller.py index ea2f3bfec95c..f36867bd8626 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/repartition_controller.py +++ b/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: diff --git a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test_repartition_controller.py b/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test_repartition_controller.py index 2e9aad0eb6dd..3efceab0d936 100644 --- a/products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test_repartition_controller.py +++ b/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)