diff --git a/products/warehouse_sources/backend/models/external_data_schema.py b/products/warehouse_sources/backend/models/external_data_schema.py index 1b49950150bc..a23ce12f77bf 100644 --- a/products/warehouse_sources/backend/models/external_data_schema.py +++ b/products/warehouse_sources/backend/models/external_data_schema.py @@ -497,9 +497,18 @@ def repartition_claim(self) -> dict[str, Any] | None: return None def _save_sync_type_config(self) -> None: + # temporalio at module scope would put the Temporal client on the django.setup() path — + # this is a models module (see external_data_source.reload_schemas for the same pattern). + from posthog.temporal.common.utils import retry_on_db_connection_drop # noqa: PLC0415 + # Internal bookkeeping write — skip the activity-log SELECT (see save()) since these run # inside the sync/repartition activity where a dropped pooler connection would fail the run. - self.save(update_fields=["sync_type_config", "updated_at"], skip_activity_log=True) + # These fire once per batch across every schema sync, so a transient pooler wait_timeout + # (the pool momentarily out of free backend connections) is worth one retry rather than + # losing the write silently. + retry_on_db_connection_drop( + lambda: self.save(update_fields=["sync_type_config", "updated_at"], skip_activity_log=True) + ) def record_partition_measurement(self, max_partition_bytes: int) -> None: self.sync_type_config["max_partition_bytes"] = max_partition_bytes diff --git a/products/warehouse_sources/backend/tests/test_models.py b/products/warehouse_sources/backend/tests/test_models.py index 376226596d1c..afc23517e401 100644 --- a/products/warehouse_sources/backend/tests/test_models.py +++ b/products/warehouse_sources/backend/tests/test_models.py @@ -1,11 +1,12 @@ import uuid from datetime import date, datetime, timedelta +from typing import Any import pytest from posthog.test.base import BaseTest from unittest.mock import patch -from django.db import DatabaseError, transaction +from django.db import DatabaseError, OperationalError, transaction from django.db.models import Model from django.test import SimpleTestCase from django.utils import timezone @@ -226,6 +227,50 @@ def test_bookkeeping_save_raises_instead_of_resurrecting_deleted_row(self) -> No assert not ExternalDataSchema.objects.filter(pk=schema_id).exists() +class TestSaveSyncTypeConfigRetriesOnConnectionDrop(BaseTest): + """A pgbouncer `query_wait_timeout` (surfaced as `OperationalError`) on this per-batch + bookkeeping write used to propagate straight past `record_partition_measurement` and get + swallowed by the caller's broad except, silently losing the write for that run.""" + + def setUp(self) -> None: + super().setUp() + self.source = ExternalDataSource.objects.create( + team_id=self.team.pk, + source_id=str(uuid.uuid4()), + connection_id=str(uuid.uuid4()), + status="Completed", + source_type="Postgres", + ) + + def _create(self, **kwargs) -> ExternalDataSchema: + return ExternalDataSchema.objects.create(team_id=self.team.pk, source=self.source, name="users", **kwargs) + + def test_retries_once_on_operational_error_then_succeeds(self) -> None: + schema = self._create() + real_save = ExternalDataSchema.save + calls = 0 + + def flaky_save(self: ExternalDataSchema, *args: Any, **kwargs: Any) -> None: + nonlocal calls + calls += 1 + if calls == 1: + raise OperationalError("query_wait_timeout") + real_save(self, *args, **kwargs) + + with patch.object(ExternalDataSchema, "save", flaky_save): + schema.record_partition_measurement(123) + + assert calls == 2 + schema.refresh_from_db() + assert schema.sync_type_config["max_partition_bytes"] == 123 + + def test_second_consecutive_operational_error_propagates(self) -> None: + schema = self._create() + with patch.object(ExternalDataSchema, "save", side_effect=OperationalError("query_wait_timeout")): + with self.assertRaises(OperationalError): + schema.record_partition_measurement(123) + + class TestExternalDataSchemaOOMEvent(BaseTest): def _source(self) -> ExternalDataSource: return ExternalDataSource.objects.create(