feat(feature-flags): Enable experience continuity with feature flags

latest_migrations.manifest

@@ -3,7 +3,7 @@ auth: 0012_alter_user_first_name_max_length
 axes: 0006_remove_accesslog_trusted
 contenttypes: 0002_remove_content_type_name
 ee: 0012_migrate_tags_v2
-posthog: 0240_organizationinvite_message
+posthog: 0241_feature_flag_hash_key_overrides
 rest_hooks: 0002_swappable_hook_model
 sessions: 0001_initial
 social_django: 0010_uid_db_index

plugin-server/src/worker/ingestion/process-event.ts

@@ -5,7 +5,7 @@ import crypto from 'crypto'
 import equal from 'fast-deep-equal'
 import { ProducerRecord } from 'kafkajs'
 import { DateTime, Duration } from 'luxon'
-import { DatabaseError } from 'pg'
+import { DatabaseError, PoolClient } from 'pg'
 
 import { Event as EventProto, IEvent } from '../../config/idl/protos'
 import { KAFKA_EVENTS, KAFKA_SESSION_RECORDING_EVENTS } from '../../config/kafka-topics'
@@ -476,13 +476,7 @@ export class EventsProcessor {
                     client
                 )
 
-                // Merge the distinct IDs
-                await this.db.postgresQuery(
-                    'UPDATE posthog_cohortpeople SET person_id = $1 WHERE person_id = $2',
-                    [mergeInto.id, otherPerson.id],
-                    'updateCohortPeople',
-                    client
-                )
+                await this.handleTablesDependingOnPersonID(otherPerson, mergeInto, client)
 
                 const distinctIdMessages = await this.db.moveDistinctIds(otherPerson, mergeInto, client)
 
@@ -798,4 +792,62 @@ export class EventsProcessor {
             )
         }
     }
+
+    private async handleTablesDependingOnPersonID(
+        sourcePerson: Person,
+        targetPerson: Person,
+        client?: PoolClient
+    ): Promise<undefined> {
+        // When personIDs change, update places depending on a person_id foreign key
+
+        // For Cohorts
+        await this.db.postgresQuery(
+            'UPDATE posthog_cohortpeople SET person_id = $1 WHERE person_id = $2',
+            [targetPerson.id, sourcePerson.id],
+            'updateCohortPeople',
+            client
+        )
+
+        // For FeatureFlagHashKeyOverrides
+        const allOverrides = await this.db.postgresQuery(
+            'SELECT id, person_id, feature_flag_key FROM posthog_featureflaghashkeyoverride WHERE team_id = $1 AND person_id = ANY($2)',
+            [sourcePerson.team_id, [sourcePerson.id, targetPerson.id]],
+            'selectFeatureFlagHashKeyOverride'
+        )
+
+        // Update where feature_flag_key exists for sourcePerson but not for targetPerson
+        const sourceOverrides = allOverrides.rows.filter((override) => override.person_id == sourcePerson.id)
+        const targetOverrideKeys = allOverrides.rows
+            .filter((override) => override.person_id == targetPerson.id)
+            .map((override) => override.feature_flag_key)
+
+        const itemsToUpdate = sourceOverrides
+            .filter((override) => !targetOverrideKeys.includes(override.feature_flag_key))
+            .map((override) => override.id)
+
+        if (allOverrides.rowCount === 0 || sourceOverrides.length === 0) {
+            return
+        }
+
+        if (itemsToUpdate.length !== 0) {
+            await this.db.postgresQuery(
+                `UPDATE posthog_featureflaghashkeyoverride SET person_id = $1 WHERE person_id = $2 AND id = ANY($3)
+                `,
+                [targetPerson.id, sourcePerson.id, itemsToUpdate],
+                'updateFeatureFlagHashKeyOverride',
+                client
+            )
+        }
+
+        // delete all other instances
+        // necessary to make sure person can then be deleted
+        if (sourceOverrides.length !== 0) {
+            await this.db.postgresQuery(
+                'DELETE FROM posthog_featureflaghashkeyoverride WHERE person_id = $1',
+                [sourcePerson.id],
+                'deleteFeatureFlagHashKeyOverride',
+                client
+            )
+        }
+    }
 }

plugin-server/tests/helpers/sql.ts

@@ -31,6 +31,7 @@ export const POSTGRES_TRUNCATE_TABLES_QUERY = `
 TRUNCATE TABLE
     posthog_personalapikey,
     posthog_featureflag,
+    posthog_featureflaghashkeyoverride,
     posthog_annotation,
     posthog_dashboarditem,
     posthog_dashboard,

plugin-server/tests/main/process-event.test.ts

@@ -32,7 +32,8 @@ export async function createPerson(
     server: Hub,
     team: Team,
     distinctIds: string[],
-    properties: Record<string, any> = {}
+    properties: Record<string, any> = {},
+    identified = false
 ): Promise<Person> {
     return server.db.createPerson(
         DateTime.utc(),
@@ -41,7 +42,7 @@ export async function createPerson(
         {},
         team.id,
         null,
-        false,
+        identified,
         new UUIDT().toString(),
         distinctIds
     )
@@ -1475,6 +1476,184 @@ test('distinct with anonymous_id which was already created', async () => {
     expect(person.is_identified).toEqual(true)
 })
 
+test('distinct with anonymous_id which was already created triggers foreign key updates', async () => {
+    const anonPerson = await createPerson(hub, team, ['anonymous_id'])
+    const identifiedPerson = await createPerson(hub, team, ['new_distinct_id'], { email: 'someone@gmail.com' })
+
+    // existing overrides
+    await hub.db.postgresQuery(
+        `INSERT INTO "posthog_featureflaghashkeyoverride" ("feature_flag_key", "person_id", "team_id", "hash_key")
+            VALUES
+                ('beta-feature', $1, $2, 'example_id'),
+                ('multivariate-flag', $3, $4, 'example_id')`,
+        [anonPerson.id, team.id, anonPerson.id, team.id],
+        'testHashKeyOverride'
+    )
+
+    // this event means the `anonPerson` will be deleted
+    // so hashkeyoverride should be updated to `identifiedPerson`'s id
+    await processEvent(
+        'new_distinct_id',
+        '',
+        '',
+        {
+            event: '$identify',
+            properties: {
+                $anon_distinct_id: 'anonymous_id',
+                token: team.api_token,
+                distinct_id: 'new_distinct_id',
+            },
+        } as any as PluginEvent,
+        team.id,
+        now,
+        now,
+        new UUIDT().toString()
+    )
+
+    const [person] = await hub.db.fetchPersons()
+    expect(person.id).toEqual(identifiedPerson.id)
+    expect(await hub.db.fetchDistinctIdValues(person)).toEqual(['anonymous_id', 'new_distinct_id'])
+    expect(person.is_identified).toEqual(true)
+
+    const result = await hub.db.postgresQuery(
+        `SELECT "feature_flag_key", "person_id", "hash_key" FROM "posthog_featureflaghashkeyoverride" WHERE "team_id" = $1`,
+        [team.id],
+        'testQueryHashKeyOverride'
+    )
+    expect(result.rows.sort()).toEqual([
+        {
+            feature_flag_key: 'beta-feature',
+            person_id: identifiedPerson.id,
+            hash_key: 'example_id',
+        },
+        {
+            feature_flag_key: 'multivariate-flag',
+            person_id: identifiedPerson.id,
+            hash_key: 'example_id',
+        },
+    ])
+})
+
+test('distinct with conflicting keys triggers foreign key updates and deletes gracefully', async () => {
+    const anonPerson = await createPerson(hub, team, ['anon_id'])
+    const identifiedPerson = await createPerson(hub, team, ['new_distinct_id'], { email: 'someone@gmail.com' }, true)
+
+    // existing overrides for both anonPerson and identifiedPerson
+    // which implies a clash when anonPerson is deleted
+    await hub.db.postgresQuery(
+        `INSERT INTO "posthog_featureflaghashkeyoverride" ("feature_flag_key", "person_id", "team_id", "hash_key")
+            VALUES
+                ('beta-feature', $1, $2, 'example_id'),
+                ('beta-feature', $3, $4, 'different_id'),
+                ('multivariate-flag', $5, $6, 'other_different_id')
+        `,
+        [anonPerson.id, team.id, identifiedPerson.id, team.id, anonPerson.id, team.id],
+        'testHashKeyOverride'
+    )
+
+    // this event means the `anonPerson` will be deleted
+    // so hashkeyoverride should be updated to `identifiedPerson`'s id
+    await processEvent(
+        'new_distinct_id',
+        '',
+        '',
+        {
+            event: '$identify',
+            properties: {
+                $anon_distinct_id: 'anon_id',
+                token: team.api_token,
+                distinct_id: 'new_distinct_id',
+            },
+        } as any as PluginEvent,
+        team.id,
+        now,
+        now,
+        new UUIDT().toString()
+    )
+
+    const [person] = await hub.db.fetchPersons()
+    expect(person.id).toEqual(identifiedPerson.id)
+    expect(await hub.db.fetchDistinctIdValues(person)).toEqual(['anon_id', 'new_distinct_id'])
+    expect(person.is_identified).toEqual(true)
+
+    const result = await hub.db.postgresQuery(
+        `SELECT "feature_flag_key", "person_id", "hash_key" FROM "posthog_featureflaghashkeyoverride" WHERE "team_id" = $1`,
+        [team.id],
+        'testQueryHashKeyOverride'
+    )
+    expect(result.rows.sort()).toEqual([
+        {
+            feature_flag_key: 'beta-feature',
+            person_id: identifiedPerson.id,
+            hash_key: 'different_id', // wasn't overriden from anon flag, because override already exists
+        },
+        {
+            feature_flag_key: 'multivariate-flag',
+            person_id: identifiedPerson.id,
+            hash_key: 'other_different_id',
+        },
+    ])
+})
+
+test('distinct with no conflicting keys doesnt trigger foreign key updates', async () => {
+    await createPerson(hub, team, ['anon_id'])
+    const identifiedPerson = await createPerson(hub, team, ['new_distinct_id'], { email: 'someone@gmail.com' }, true)
+
+    // existing overrides for both anonPerson and identifiedPerson
+    // which implies a clash when anonPerson is deleted
+    await hub.db.postgresQuery(
+        `INSERT INTO "posthog_featureflaghashkeyoverride" ("feature_flag_key", "person_id", "team_id", "hash_key")
+            VALUES
+                ('beta-feature', $1, $2, 'example_id'),
+                ('multivariate-flag', $3, $4, 'different_id')`,
+        [identifiedPerson.id, team.id, identifiedPerson.id, team.id],
+        'testHashKeyOverride'
+    )
+
+    // this event means the `anonPerson` will be deleted
+    // so hashkeyoverride should be updated to `identifiedPerson`'s id
+    await processEvent(
+        'new_distinct_id',
+        '',
+        '',
+        {
+            event: '$identify',
+            properties: {
+                $anon_distinct_id: 'anon_id',
+                token: team.api_token,
+                distinct_id: 'new_distinct_id',
+            },
+        } as any as PluginEvent,
+        team.id,
+        now,
+        now,
+        new UUIDT().toString()
+    )
+
+    const [person] = await hub.db.fetchPersons()
+    expect(person.id).toEqual(identifiedPerson.id)
+    expect(await hub.db.fetchDistinctIdValues(person)).toEqual(['anon_id', 'new_distinct_id'])
+    expect(person.is_identified).toEqual(true)
+
+    const result = await hub.db.postgresQuery(
+        `SELECT "feature_flag_key", "person_id", "hash_key" FROM "posthog_featureflaghashkeyoverride" WHERE "team_id" = $1`,
+        [team.id],
+        'testQueryHashKeyOverride'
+    )
+    expect(result.rows.sort()).toEqual([
+        {
+            feature_flag_key: 'beta-feature',
+            person_id: identifiedPerson.id,
+            hash_key: 'example_id',
+        },
+        {
+            feature_flag_key: 'multivariate-flag',
+            person_id: identifiedPerson.id,
+            hash_key: 'different_id',
+        },
+    ])
+})
+
 test('identify with the same distinct_id as anon_distinct_id', async () => {
     await createPerson(hub, team, ['anonymous_id'])
 

posthog/api/decide.py

@@ -150,7 +150,9 @@ def get_decide(request: HttpRequest):
             team = user.teams.get(id=project_id)
 
         if team:
-            feature_flags = get_overridden_feature_flags(team.pk, data["distinct_id"], data.get("groups", {}))
+            feature_flags = get_overridden_feature_flags(
+                team.pk, data["distinct_id"], data.get("groups", {}), hash_key_override=data.get("$anon_distinct_id")
+            )
             response["featureFlags"] = feature_flags if api_version >= 2 else list(feature_flags.keys())
 
             if team.session_recording_opt_in and (on_permitted_domain(team, request) or len(team.app_urls) == 0):

posthog/api/feature_flag.py

@@ -54,6 +54,7 @@ class Meta:
             "created_at",
             "is_simple_flag",
             "rollout_percentage",
+            "ensure_experience_continuity",
         ]
 
     # Simple flags are ones that only have rollout_percentage