fix(flags): don't enclose in overall transaction so we get latest rea…

…ds (#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>

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

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,23 +629,23 @@ 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"},
             )
             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},
             )
             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,

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"])

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