From 428542a78ac9be0662e2f6c0a57d94d61b0c7a5b Mon Sep 17 00:00:00 2001 From: Brett Hoerner Date: Thu, 26 Oct 2023 10:02:52 -0600 Subject: [PATCH] =?UTF-8?q?refactor(plugin-server):=20combine=202=20person?= =?UTF-8?q?=20merge=20queries=20into=20a=20single=E2=80=A6=20(#18127)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit refactor(plugin-server): combine 2 person merge queries into a single round trip to PG --- plugin-server/src/utils/db/db.ts | 45 +++++++++++-------- .../src/worker/ingestion/person-state.ts | 26 +++-------- plugin-server/tests/main/db.test.ts | 10 ++--- posthog/api/test/test_decide.py | 2 +- posthog/models/feature_flag/flag_matching.py | 2 +- 5 files changed, 39 insertions(+), 46 deletions(-) diff --git a/plugin-server/src/utils/db/db.ts b/plugin-server/src/utils/db/db.ts index 33c384e28fb03..f92a8a4404d44 100644 --- a/plugin-server/src/utils/db/db.ts +++ b/plugin-server/src/utils/db/db.ts @@ -965,36 +965,43 @@ export class DB { return insertResult.rows[0] } - // Feature Flag Hash Key overrides - public async addFeatureFlagHashKeysForMergedPerson( + public async updateCohortsAndFeatureFlagsForMerge( teamID: Team['id'], sourcePersonID: Person['id'], targetPersonID: Person['id'], tx?: TransactionClient ): Promise { - // Delete and insert in a single query to ensure - // this function is safe wherever it is run. - // The CTE helps make this happen. - // - // Every override is unique for a team-personID-featureFlag combo. - // In case we run into a conflict we would ideally use the override from most recent - // personId used, so the user experience is consistent, however that's tricky to figure out - // this also happens rarely, so we're just going to do the performance optimal thing - // i.e. do nothing on conflicts, so we keep using the value that the person merged into had + // When personIDs change, update places depending on a person_id foreign key + await this.postgres.query( tx ?? PostgresUse.COMMON_WRITE, - ` - WITH deletions AS ( - DELETE FROM posthog_featureflaghashkeyoverride WHERE team_id = $1 AND person_id = $2 + // Do two high level things in a single round-trip to the DB. + // + // 1. Update cohorts. + // 2. Update (delete+insert) feature flags. + // + // NOTE: Every override is unique for a team-personID-featureFlag combo. In case we run + // into a conflict we would ideally use the override from most recent personId used, so + // the user experience is consistent, however that's tricky to figure out this also + // happens rarely, so we're just going to do the performance optimal thing i.e. do + // nothing on conflicts, so we keep using the value that the person merged into had + `WITH cohort_update AS ( + UPDATE posthog_cohortpeople + SET person_id = $1 + WHERE person_id = $2 + RETURNING person_id + ), + deletions AS ( + DELETE FROM posthog_featureflaghashkeyoverride + WHERE team_id = $3 AND person_id = $2 RETURNING team_id, person_id, feature_flag_key, hash_key ) INSERT INTO posthog_featureflaghashkeyoverride (team_id, person_id, feature_flag_key, hash_key) - SELECT team_id, $3, feature_flag_key, hash_key + SELECT team_id, $1, feature_flag_key, hash_key FROM deletions - ON CONFLICT DO NOTHING - `, - [teamID, sourcePersonID, targetPersonID], - 'addFeatureFlagHashKeysForMergedPerson' + ON CONFLICT DO NOTHING`, + [targetPersonID, sourcePersonID, teamID], + 'updateCohortAndFeatureFlagsPeople' ) } diff --git a/plugin-server/src/worker/ingestion/person-state.ts b/plugin-server/src/worker/ingestion/person-state.ts index 9476a55369acd..a65ada8e072f9 100644 --- a/plugin-server/src/worker/ingestion/person-state.ts +++ b/plugin-server/src/worker/ingestion/person-state.ts @@ -506,7 +506,12 @@ export class PersonState { // Merge the distinct IDs // TODO: Doesn't this table need to add updates to CH too? - await this.handleTablesDependingOnPersonID(otherPerson, mergeInto, tx) + await this.db.updateCohortsAndFeatureFlagsForMerge( + otherPerson.team_id, + otherPerson.id, + mergeInto.id, + tx + ) const distinctIdMessages = await this.db.moveDistinctIds(otherPerson, mergeInto, tx) @@ -681,25 +686,6 @@ export class PersonState { return id } - - private async handleTablesDependingOnPersonID( - sourcePerson: Person, - targetPerson: Person, - tx: TransactionClient - ): Promise { - // When personIDs change, update places depending on a person_id foreign key - - // For Cohorts - await this.db.postgres.query( - tx, - 'UPDATE posthog_cohortpeople SET person_id = $1 WHERE person_id = $2', - [targetPerson.id, sourcePerson.id], - 'updateCohortPeople' - ) - - // For FeatureFlagHashKeyOverrides - await this.db.addFeatureFlagHashKeysForMergedPerson(sourcePerson.team_id, sourcePerson.id, targetPerson.id, tx) - } } export function ageInMonthsLowCardinality(timestamp: DateTime): number { diff --git a/plugin-server/tests/main/db.test.ts b/plugin-server/tests/main/db.test.ts index 5050dac1be4f9..e165584e5ee20 100644 --- a/plugin-server/tests/main/db.test.ts +++ b/plugin-server/tests/main/db.test.ts @@ -845,7 +845,7 @@ describe('DB', () => { }) }) - describe('addFeatureFlagHashKeysForMergedPerson()', () => { + describe('updateCohortsAndFeatureFlagsForMerge()', () => { let team: Team let sourcePersonID: Person['id'] let targetPersonID: Person['id'] @@ -889,7 +889,7 @@ describe('DB', () => { }) it("doesn't fail on empty data", async () => { - await db.addFeatureFlagHashKeysForMergedPerson(team.id, sourcePersonID, targetPersonID) + await db.updateCohortsAndFeatureFlagsForMerge(team.id, sourcePersonID, targetPersonID) }) it('updates all valid keys when target person had no overrides', async () => { @@ -906,7 +906,7 @@ describe('DB', () => { hash_key: 'override_value_for_beta_feature', }) - await db.addFeatureFlagHashKeysForMergedPerson(team.id, sourcePersonID, targetPersonID) + await db.updateCohortsAndFeatureFlagsForMerge(team.id, sourcePersonID, targetPersonID) const result = await getAllHashKeyOverrides() @@ -947,7 +947,7 @@ describe('DB', () => { hash_key: 'existing_override_value_for_beta_feature', }) - await db.addFeatureFlagHashKeysForMergedPerson(team.id, sourcePersonID, targetPersonID) + await db.updateCohortsAndFeatureFlagsForMerge(team.id, sourcePersonID, targetPersonID) const result = await getAllHashKeyOverrides() @@ -982,7 +982,7 @@ describe('DB', () => { hash_key: 'override_value_for_beta_feature', }) - await db.addFeatureFlagHashKeysForMergedPerson(team.id, sourcePersonID, targetPersonID) + await db.updateCohortsAndFeatureFlagsForMerge(team.id, sourcePersonID, targetPersonID) const result = await getAllHashKeyOverrides() diff --git a/posthog/api/test/test_decide.py b/posthog/api/test/test_decide.py index 804e8e8d86460..ac7049b98ed98 100644 --- a/posthog/api/test/test_decide.py +++ b/posthog/api/test/test_decide.py @@ -1221,7 +1221,7 @@ def test_feature_flags_v2_consistent_flags_with_merged_persons(self, *args): new_person_id = person.id old_person_id = person2.id # this happens in the plugin server - # https://github.com/PostHog/posthog/blob/master/plugin-server/src/worker/ingestion/person-state.ts#L696 (addFeatureFlagHashKeysForMergedPerson) + # https://github.com/PostHog/posthog/blob/master/plugin-server/src/utils/db/db.ts (updateCohortsAndFeatureFlagsForMerge) # at which point we run the query query = f""" WITH deletions AS ( diff --git a/posthog/models/feature_flag/flag_matching.py b/posthog/models/feature_flag/flag_matching.py index 05c5bbccb8f63..41da1ee13ce7d 100644 --- a/posthog/models/feature_flag/flag_matching.py +++ b/posthog/models/feature_flag/flag_matching.py @@ -717,7 +717,7 @@ def get_all_feature_flags( # 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) + # https://github.com/PostHog/posthog/blob/master/plugin-server/src/utils/db/db.ts (updateCohortsAndFeatureFlagsForMerge) writing_hash_key_override = set_feature_flag_hash_key_overrides( team_id, [distinct_id, hash_key_override], hash_key_override