fix(flags): fix cohort version update

[matheus-vb]

Problem

A customer reported their cohort showed stale data ("no new people since a day ago") despite last_calculation being recent. Removing and re-adding a filter fixed it by forcing a version bump.

The probable root cause: calculate_people_ch updates the Postgres Cohort.version after the try/finally block. If self.save() or _safe_reset_calculating_state() in the finally block raises, the version update is silently skipped. ClickHouse has the new version's data, but Postgres still points to the old version, so all queries return stale results.

Changes

Moved the Cohort.objects.filter(...).update(version=..., count=...) call from after the try/finally block into the try block, right after recalculate_cohortpeople succeeds. This ensures the version is persisted before the finally block runs, so exceptions there can no longer silently skip it.

• Renamed update_fields to version_update_fields to avoid confusion with the self.save(update_fields=...) call in finally
• Added a test that mocks _safe_reset_calculating_state to raise and verifies version is still updated

The concurrency guard (version__lt=pending_version | version__isnull=True) is preserved, lower versions still can't overwrite higher ones.

Tradeoff: If the version .update() DB call itself raises, errors_calculating is incremented even though ClickHouse data was written. This is visible and self-healing (next retry succeeds), far better than the current silent split-brain.

How did you test this code?

• New unit test test_calculate_people_ch_updates_version_even_when_finally_raises
• Existing cohort tests verified for regressions
• ruff check and ruff format pass

[greptile-apps]

Prompt To Fix All With AI

This is a comment left during a code review.
Path: posthog/test/test_cohort_model.py
Line: 513-533

Comment:
**Test only covers one of two `finally` failure paths**

The new test mocks `_safe_reset_calculating_state` to raise, but `self.save()` is called *before* `_safe_reset_calculating_state` in the `finally` block. If `self.save()` raises, `_safe_reset_calculating_state` is never reached yet the fix still protects the version update. Adding a second case (or parameterising the test) would give complete coverage of the original problem statement.

For example, using `@pytest.mark.parametrize` over both `("_safe_reset_calculating_state", "DB connection lost")` and `("save", "DB error")` would demonstrate the fix handles both `finally`-block failure modes, which aligns with the PR description's claim that the fix addresses `self.save()` failures too.

How can I resolve this? If you propose a fix, please make it concise.

_{Reviews (1): Last reviewed commit: "apply fmt" | Re-trigger Greptile}