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

[Gilbert09]

Problem

Error tracking surfaced an OperationalError: consuming input failed: server closed the connection unexpectedly from is_auto_repartition_enabled in products/warehouse_sources/backend/temporal/data_imports/pipelines/core/repartition_controller.py, during a Postgres source sync.

The failing line is a plain Django ORM read (Team.objects.only(...).get(id=schema.team_id)) against PostHog's own app database, not the customer's synced source. It runs via asyncio.to_thread inside a long-lived Temporal worker thread, so a pooler-recycled or dropped connection surfaces as OperationalError the first time it's reused — a known, transient failure mode in this codebase, not a bug in the sync logic itself.

The call is already wrapped in a broad try/except in maybe_flag_for_repartition that reports it via capture_exception and returns without breaking the sync (the function is documented to "never raise" — this is best-effort repartition detection). So the transient blip doesn't break anything, but it does add noise to error tracking on every occurrence, since the dead connection is neither evicted nor retried.

Changes

• Wrap the Team.objects.get() read in is_auto_repartition_enabled with retry_on_db_connection_drop (posthog/temporal/common/utils.py), the existing helper built for exactly this failure class: it evicts the stale pooled connection and retries once on OperationalError/InterfaceError, matching the pattern already used elsewhere in this Temporal worker tree.
• Added a regression test that raises OperationalError on the first Team lookup and asserts the function retries and resolves the flag instead of the error propagating uncaught.

I checked for an existing fix first: #73386 classifies a similar DB-connection blip as transient, but in a different function (run_pre_write_defensive_compact's job.folder_path() resolution) — it doesn't touch is_auto_repartition_enabled or this call site. No open PR covers this specific one.

How did you test this code?

Added TestIsAutoRepartitionEnabled::test_retries_once_on_transient_db_connection_drop, which mocks the Team.objects.only(...).get() call to raise OperationalError once then succeed, and asserts the function returns the resolved flag value instead of the exception propagating. Without the fix this test fails, since the error isn't caught anywhere inside is_auto_repartition_enabled — it currently escapes to the caller's outer catch-all.

Ran:

• uv run pytest products/warehouse_sources/backend/temporal/data_imports/pipelines/core/test_repartition_controller.py — 30 passed
• uv run mypy --cache-fine-grained . (repo-wide, as CI runs it) — clean on the changed files
• ruff check --fix / ruff format — clean
• hogli ci:preflight --fix — clean

Automatic notifications

• Publish to changelog?
• Alert Sales and Marketing teams?

Docs update

N/A — internal error-handling change, no user-facing behavior change.

🤖 Agent context

Autonomy: Fully autonomous

Triaged a webhook-delivered error tracking issue using the PostHog MCP error-tracking tools (query-error-tracking-issue, query-error-tracking-issue-events) to pull the stack trace and confirm the exact failing line and call path. Used a research subagent to check for an existing connection-hygiene pattern in this Temporal worker tree, which surfaced retry_on_db_connection_drop/aretry_on_db_connection_drop in posthog/temporal/common/utils.py as the established fix for this exact failure class. Checked for duplicate open PRs by exception type, message phrase, module path, and the maintainer's own open PRs before writing the fix — found a related-but-non-overlapping PR (noted above) rather than a duplicate. Invoked /writing-tests before adding the regression test.

---

Created with PostHog Code

[trunk-io]

😎 Merged manually by @gantoine - details.

[github-actions]

Hey @Gilbert09! 👋

It looks like your git author email on this PR isn't your @posthog.com address (owerstom@gmail.com). Since you're on the PostHog team, it's worth pointing your local git author email at your @posthog.com address. Why it matters:

• Consistent work identity in git history — internal tooling that attributes commits to team members keys off your @posthog.com address.
• Keeps team contributions easy to tell apart from external community ones when scanning history.

You can fix it for this repo with:

git config user.email "you@posthog.com"

Or set it globally with git config --global user.email "you@posthog.com". No need to redo this PR — just a nudge for next time. 🙂

[stamphog]

Small, contained fix wrapping one DB read with an existing, documented retry helper for a known transient failure mode; not risky territory, author owns this code with strong familiarity, and a regression test backs the fix.

• Author wrote 100% of the modified lines and has 8 merged PRs in these paths (familiarity STRONG).
• 👍 on the PR from hex-security-app[bot].

Gate mechanics and policy version

Gate		Result
prerequisites	✓	all clear
deny-list	✓	no deny categories matched
size	✓	3L, 1F substantive, 24L/2F incl. docs/generated/snapshots — within ceiling
tier	✓	T1-agent / T1b-small (24L, 2F, single-area, fix)
stamphog 2.0.0b3		.stamphog/policy.yml @ d24a844 · reviewed head b711fd5

[fuziontech]

Reviewed the complete diff and surrounding retry helper usage. The change correctly retries the Team lookup once on Django connection-drop errors while preserving existing missing-team behavior, and the focused regression test covers the intended failure path. Ruff passes. The focused local pytest could not start because the local Postgres host db is unavailable; the relevant warehouse-sources CI shards passed. No blocking issues found.

This review was generated by an automated review agent on behalf of @fuziontech.

[talyn-app]

/trunk merge

[deployment-status-posthog]

Deploy status

Environment	Status	Deployed At	Workflow
dev	✅ Deployed	2026-07-30 18:16 UTC	Run
prod-us	✅ Deployed	2026-07-30 18:39 UTC	Run
prod-eu	✅ Deployed	2026-07-30 18:35 UTC	Run