From a12977ebae998643230a010c17bd4b321e4db9bc Mon Sep 17 00:00:00 2001 From: Neil Kakkar Date: Fri, 7 Apr 2023 16:11:43 +0100 Subject: [PATCH] fix(flags): don't enclose in overall transaction so we get latest reads (#15003) * fix(flags): don't enclose in overall transaction so we get latest reads * fix tests * Update query snapshots --------- Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com> --- .../test/__snapshots__/test_feature_flag.ambr | 40 ++++- posthog/api/test/test_decide.py | 14 +- posthog/api/test/test_feature_flag.py | 2 +- posthog/models/feature_flag/flag_matching.py | 39 ++--- .../test/__snapshots__/test_feature_flag.ambr | 162 ++++++++++++------ posthog/test/test_feature_flag.py | 37 ++-- 6 files changed, 183 insertions(+), 111 deletions(-) diff --git a/posthog/api/test/__snapshots__/test_feature_flag.ambr b/posthog/api/test/__snapshots__/test_feature_flag.ambr index 1ce586bf89e2b..48b8dda1323ac 100644 --- a/posthog/api/test/__snapshots__/test_feature_flag.ambr +++ b/posthog/api/test/__snapshots__/test_feature_flag.ambr @@ -500,7 +500,45 @@ ' SELECT pg_sleep(1); - SAVEPOINT _snapshot_ + WITH target_person_ids AS + (SELECT team_id, + person_id + FROM posthog_persondistinctid + WHERE team_id = 2 + AND distinct_id IN ('example_id', + 'random') ), + existing_overrides AS + (SELECT team_id, + person_id, + feature_flag_key, + hash_key + FROM posthog_featureflaghashkeyoverride + WHERE team_id = 2 + AND person_id IN + (SELECT person_id + FROM target_person_ids) ), + flags_to_override AS + (SELECT key + FROM posthog_featureflag + WHERE team_id = 2 + AND ensure_experience_continuity = TRUE + AND active = TRUE + AND deleted = FALSE + AND key NOT IN + (SELECT feature_flag_key + FROM existing_overrides) ) + INSERT INTO posthog_featureflaghashkeyoverride (team_id, person_id, feature_flag_key, hash_key) + SELECT team_id, + person_id, + key, + 'random' + FROM flags_to_override, + target_person_ids + WHERE EXISTS + (SELECT 1 + FROM posthog_person + WHERE id = person_id + AND team_id = 2) ON CONFLICT DO NOTHING ' --- # name: TestResiliency.test_feature_flags_v3_with_experience_continuity_working_slow_db.5 diff --git a/posthog/api/test/test_decide.py b/posthog/api/test/test_decide.py index 599286b99a38b..5b5ef58be6a57 100644 --- a/posthog/api/test/test_decide.py +++ b/posthog/api/test/test_decide.py @@ -589,7 +589,7 @@ def test_feature_flags_v2_consistent_flags(self): # person2 = Person.objects.create(team=self.team, distinct_ids=["example_id", "other_id"], properties={"email": "tim@posthog.com"}) person.add_distinct_id("other_id") - with self.assertNumQueries(11): + with self.assertNumQueries(9): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": "other_id", "$anon_distinct_id": "example_id"}, @@ -629,7 +629,7 @@ def test_feature_flags_v3_consistent_flags_with_numeric_distinct_ids(self): self.assertTrue(response.json()["featureFlags"]["beta-feature"]) self.assertTrue(response.json()["featureFlags"]["default-flag"]) - with self.assertNumQueries(11): + with self.assertNumQueries(9): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": 12345, "$anon_distinct_id": "example_id"}, @@ -637,7 +637,7 @@ def test_feature_flags_v3_consistent_flags_with_numeric_distinct_ids(self): self.assertTrue(response.json()["featureFlags"]["beta-feature"]) self.assertTrue(response.json()["featureFlags"]["default-flag"]) - with self.assertNumQueries(11): + with self.assertNumQueries(9): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": "xyz", "$anon_distinct_id": 12345}, @@ -645,7 +645,7 @@ def test_feature_flags_v3_consistent_flags_with_numeric_distinct_ids(self): self.assertTrue(response.json()["featureFlags"]["beta-feature"]) self.assertTrue(response.json()["featureFlags"]["default-flag"]) - with self.assertNumQueries(11): + with self.assertNumQueries(9): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": 5, "$anon_distinct_id": 12345}, @@ -705,7 +705,7 @@ def test_feature_flags_v2_consistent_flags_with_ingestion_delays(self): # identify event is sent, but again, ingestion delays, so no entry in personDistinctID table # person.add_distinct_id("other_id") # in which case, we're pretty much trashed - with self.assertNumQueries(10): + with self.assertNumQueries(8): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": "other_id", "$anon_distinct_id": "example_id"}, @@ -773,7 +773,7 @@ def test_feature_flags_v2_consistent_flags_with_merged_persons(self): ) # caching flag definitions in the above mean fewer queries - with self.assertNumQueries(11): + with self.assertNumQueries(9): response = self._post_decide( api_version=2, data={"token": self.team.api_token, "distinct_id": "other_id", "$anon_distinct_id": "example_id"}, @@ -868,7 +868,7 @@ def test_feature_flags_v2_consistent_flags_with_delayed_new_identified_person(se # new person with "other_id" is yet to be created # caching flag definitions in the above mean fewer queries - with self.assertNumQueries(11): + with self.assertNumQueries(9): # one extra query to find person_id for $anon_distinct_id response = self._post_decide( api_version=2, diff --git a/posthog/api/test/test_feature_flag.py b/posthog/api/test/test_feature_flag.py index bdad0984b3e6b..94637ed0c2b4e 100644 --- a/posthog/api/test/test_feature_flag.py +++ b/posthog/api/test/test_feature_flag.py @@ -3092,7 +3092,7 @@ def test_feature_flags_v3_with_experience_continuity_working_slow_db(self): self.assertTrue(serialized_data.is_valid()) serialized_data.save() - with snapshot_postgres_queries_context(self), self.assertNumQueries(9): + with snapshot_postgres_queries_context(self), self.assertNumQueries(7): all_flags, _, _, errors = get_all_feature_flags(team_id, "example_id", hash_key_override="random") self.assertTrue(all_flags["property-flag"]) diff --git a/posthog/models/feature_flag/flag_matching.py b/posthog/models/feature_flag/flag_matching.py index 7802d628867c0..accbe482f0035 100644 --- a/posthog/models/feature_flag/flag_matching.py +++ b/posthog/models/feature_flag/flag_matching.py @@ -5,9 +5,8 @@ import structlog from typing import Any, Dict, List, Optional, Tuple, Union -from django.db import DatabaseError, IntegrityError, transaction +from django.db import DatabaseError, IntegrityError from django.db.models.expressions import ExpressionWrapper, RawSQL -from django.db.backends.utils import CursorWrapper from django.db.models.fields import BooleanField from django.db.models.query import QuerySet from sentry_sdk.api import capture_exception @@ -487,22 +486,18 @@ def get_all_feature_flags( if hash_key_override is not None: try: hash_key_override = str(hash_key_override) - # make the entire hash key override logic a single transaction - # with a small timeout - with execute_with_timeout(FLAG_MATCHING_QUERY_TIMEOUT_MS) as timeout_cursor: - # :TRICKY: There are a few cases for write we need to handle: - # 1. Ingestion delay causing the person to not have been created yet or the distinct_id not yet associated - # 2. Merging of two different already existing persons, which results in 1 person_id being deleted and ff hash key overrides to be moved. - # 3. Person being deleted via UI or API (this is rare) - # - # In all cases, we simply try to find all personIDs associated with the distinct_id - # and the hash_key_override, and add overrides for all these personIDs. - # On merge, if a person is deleted, it is fine because the below line in plugin-server will take care of it. - # https://github.com/PostHog/posthog/blob/master/plugin-server/src/worker/ingestion/person-state.ts#L696 (addFeatureFlagHashKeysForMergedPerson) - - set_feature_flag_hash_key_overrides( - timeout_cursor, team_id, [distinct_id, hash_key_override], hash_key_override - ) + + # :TRICKY: There are a few cases for write we need to handle: + # 1. Ingestion delay causing the person to not have been created yet or the distinct_id not yet associated + # 2. Merging of two different already existing persons, which results in 1 person_id being deleted and ff hash key overrides to be moved. + # 3. Person being deleted via UI or API (this is rare) + # + # In all cases, we simply try to find all personIDs associated with the distinct_id + # and the hash_key_override, and add overrides for all these personIDs. + # On merge, if a person is deleted, it is fine because the below line in plugin-server will take care of it. + # https://github.com/PostHog/posthog/blob/master/plugin-server/src/worker/ingestion/person-state.ts#L696 (addFeatureFlagHashKeysForMergedPerson) + + set_feature_flag_hash_key_overrides(team_id, [distinct_id, hash_key_override], hash_key_override) except Exception as e: # If the database is in read-only mode, we can't handle experience continuity flags, @@ -545,9 +540,7 @@ def get_all_feature_flags( ) -def set_feature_flag_hash_key_overrides( - cursor: CursorWrapper, team_id: int, distinct_ids: List[str], hash_key_override: str -) -> None: +def set_feature_flag_hash_key_overrides(team_id: int, distinct_ids: List[str], hash_key_override: str) -> None: # As a product decision, the first override wins, i.e consistency matters for the first walkthrough. # Thus, we don't need to do upserts here. @@ -558,7 +551,9 @@ def set_feature_flag_hash_key_overrides( for _ in range(max_retries): try: - with transaction.atomic(): + # make the entire hash key override logic a single transaction + # with a small timeout + with execute_with_timeout(FLAG_MATCHING_QUERY_TIMEOUT_MS) as cursor: query = """ WITH target_person_ids AS ( SELECT team_id, person_id FROM posthog_persondistinctid WHERE team_id = %(team_id)s AND distinct_id IN %(distinct_ids)s diff --git a/posthog/test/__snapshots__/test_feature_flag.ambr b/posthog/test/__snapshots__/test_feature_flag.ambr index 8a96fead716eb..0440208f33e2d 100644 --- a/posthog/test/__snapshots__/test_feature_flag.ambr +++ b/posthog/test/__snapshots__/test_feature_flag.ambr @@ -248,33 +248,6 @@ ' --- # name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.1 - 'SAVEPOINT _snapshot_' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.10 - ' - SELECT "posthog_persondistinctid"."person_id", - "posthog_persondistinctid"."distinct_id" - FROM "posthog_persondistinctid" - WHERE ("posthog_persondistinctid"."distinct_id" IN ('other_id', - 'example_id') - AND "posthog_persondistinctid"."team_id" = 2) - ' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.11 - ' - SELECT "posthog_featureflaghashkeyoverride"."feature_flag_key", - "posthog_featureflaghashkeyoverride"."hash_key", - "posthog_featureflaghashkeyoverride"."person_id" - FROM "posthog_featureflaghashkeyoverride" - WHERE ("posthog_featureflaghashkeyoverride"."person_id" IN (1, - 2, - 3, - 4, - 5 /* ... */) - AND "posthog_featureflaghashkeyoverride"."team_id" = 2) - ' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.2 ' WITH target_person_ids AS (SELECT team_id, @@ -317,16 +290,37 @@ AND team_id = 2) ON CONFLICT DO NOTHING ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.3 - 'ROLLBACK TO SAVEPOINT _snapshot_' +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.10 + ' + SELECT "posthog_persondistinctid"."person_id", + "posthog_persondistinctid"."distinct_id" + FROM "posthog_persondistinctid" + WHERE ("posthog_persondistinctid"."distinct_id" IN ('other_id', + 'example_id') + AND "posthog_persondistinctid"."team_id" = 2) + ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.4 - 'RELEASE SAVEPOINT _snapshot_' +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.11 + ' + SELECT "posthog_featureflaghashkeyoverride"."feature_flag_key", + "posthog_featureflaghashkeyoverride"."hash_key", + "posthog_featureflaghashkeyoverride"."person_id" + FROM "posthog_featureflaghashkeyoverride" + WHERE ("posthog_featureflaghashkeyoverride"."person_id" IN (1, + 2, + 3, + 4, + 5 /* ... */) + AND "posthog_featureflaghashkeyoverride"."team_id" = 2) + ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.5 - 'SAVEPOINT _snapshot_' +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.2 + ' + + SET LOCAL statement_timeout = 2 + ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.6 +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.3 ' WITH target_person_ids AS (SELECT team_id, @@ -369,28 +363,23 @@ AND team_id = 2) ON CONFLICT DO NOTHING ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.7 - 'ROLLBACK TO SAVEPOINT _snapshot_' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.8 - 'RELEASE SAVEPOINT _snapshot_' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.9 +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.4 ' SET LOCAL statement_timeout = 2 ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.5 ' - - SET LOCAL statement_timeout = 2 + SELECT "posthog_persondistinctid"."person_id", + "posthog_persondistinctid"."distinct_id" + FROM "posthog_persondistinctid" + WHERE ("posthog_persondistinctid"."distinct_id" IN ('other_id', + 'example_id') + AND "posthog_persondistinctid"."team_id" = 2) ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.1 - 'SAVEPOINT _snapshot_' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.10 +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.6 ' SELECT "posthog_featureflaghashkeyoverride"."feature_flag_key", "posthog_featureflaghashkeyoverride"."hash_key", @@ -404,7 +393,25 @@ AND "posthog_featureflaghashkeyoverride"."team_id" = 2) ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.2 +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.7 + 'ROLLBACK TO SAVEPOINT _snapshot_' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.8 + 'RELEASE SAVEPOINT _snapshot_' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_error_race_conditions_on_person_merging.9 + ' + + SET LOCAL statement_timeout = 2 + ' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging + ' + + SET LOCAL statement_timeout = 2 + ' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.1 ' WITH target_person_ids AS (SELECT team_id, @@ -447,16 +454,27 @@ AND team_id = 2) ON CONFLICT DO NOTHING ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.3 - 'ROLLBACK TO SAVEPOINT _snapshot_' ---- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.4 - 'RELEASE SAVEPOINT _snapshot_' +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.10 + ' + SELECT "posthog_featureflaghashkeyoverride"."feature_flag_key", + "posthog_featureflaghashkeyoverride"."hash_key", + "posthog_featureflaghashkeyoverride"."person_id" + FROM "posthog_featureflaghashkeyoverride" + WHERE ("posthog_featureflaghashkeyoverride"."person_id" IN (1, + 2, + 3, + 4, + 5 /* ... */) + AND "posthog_featureflaghashkeyoverride"."team_id" = 2) + ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.5 - 'SAVEPOINT _snapshot_' +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.2 + ' + + SET LOCAL statement_timeout = 2 + ' --- -# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.6 +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.3 ' WITH target_person_ids AS (SELECT team_id, @@ -499,6 +517,36 @@ AND team_id = 2) ON CONFLICT DO NOTHING ' --- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.4 + ' + + SET LOCAL statement_timeout = 2 + ' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.5 + ' + SELECT "posthog_persondistinctid"."person_id", + "posthog_persondistinctid"."distinct_id" + FROM "posthog_persondistinctid" + WHERE ("posthog_persondistinctid"."distinct_id" IN ('other_id', + 'example_id') + AND "posthog_persondistinctid"."team_id" = 2) + ' +--- +# name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.6 + ' + SELECT "posthog_featureflaghashkeyoverride"."feature_flag_key", + "posthog_featureflaghashkeyoverride"."hash_key", + "posthog_featureflaghashkeyoverride"."person_id" + FROM "posthog_featureflaghashkeyoverride" + WHERE ("posthog_featureflaghashkeyoverride"."person_id" IN (1, + 2, + 3, + 4, + 5 /* ... */) + AND "posthog_featureflaghashkeyoverride"."team_id" = 2) + ' +--- # name: TestHashKeyOverridesRaceConditions.test_hash_key_overrides_with_simulated_race_conditions_on_person_merging.7 'RELEASE SAVEPOINT _snapshot_' --- diff --git a/posthog/test/test_feature_flag.py b/posthog/test/test_feature_flag.py index 87c3d315dc81b..28398a16e2d54 100644 --- a/posthog/test/test_feature_flag.py +++ b/posthog/test/test_feature_flag.py @@ -1826,10 +1826,9 @@ def setUpTestData(cls): def test_setting_overrides(self): - with connection.cursor() as cursor: - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" - ) + set_feature_flag_hash_key_overrides( + team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" + ) with connection.cursor() as cursor: cursor.execute( @@ -1841,10 +1840,9 @@ def test_setting_overrides(self): def test_retrieving_hash_key_overrides(self): - with connection.cursor() as cursor: - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" - ) + set_feature_flag_hash_key_overrides( + team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" + ) hash_keys = get_feature_flag_hash_key_overrides(self.team.pk, ["example_id"]) @@ -1860,13 +1858,8 @@ def test_hash_key_overrides_for_multiple_ids_when_people_are_not_merged(self): team=self.team, distinct_ids=["2"], properties={"email": "beuk2@posthog.com", "team": "posthog"} ) - with connection.cursor() as cursor: - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=["1"], hash_key_override="other_id1" - ) - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=["2"], hash_key_override="aother_id2" - ) + set_feature_flag_hash_key_overrides(team_id=self.team.pk, distinct_ids=["1"], hash_key_override="other_id1") + set_feature_flag_hash_key_overrides(team_id=self.team.pk, distinct_ids=["2"], hash_key_override="aother_id2") hash_keys = get_feature_flag_hash_key_overrides(self.team.pk, ["1", "2"]) @@ -1888,10 +1881,9 @@ def test_setting_overrides_doesnt_balk_with_existing_overrides(self): ) # and now we come to get new overrides - with connection.cursor() as cursor: - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" - ) + set_feature_flag_hash_key_overrides( + team_id=self.team.pk, distinct_ids=self.person.distinct_ids, hash_key_override="other_id" + ) with connection.cursor() as cursor: cursor.execute( @@ -1903,10 +1895,9 @@ def test_setting_overrides_doesnt_balk_with_existing_overrides(self): def test_setting_overrides_when_persons_dont_exist(self): - with connection.cursor() as cursor: - set_feature_flag_hash_key_overrides( - cursor, team_id=self.team.pk, distinct_ids=["1", "2", "3", "4"], hash_key_override="other_id" - ) + set_feature_flag_hash_key_overrides( + team_id=self.team.pk, distinct_ids=["1", "2", "3", "4"], hash_key_override="other_id" + ) with connection.cursor() as cursor: cursor.execute(