Skip to content

Commit

Permalink
refactor(plugin-server): combine 2 person merge queries into a single… (
Browse files Browse the repository at this point in the history 
#18127)

refactor(plugin-server): combine 2 person merge queries into a single round trip to PG
  • Loading branch information
@bretthoerner
bretthoerner authored Oct 26, 2023
1 parent 345a043 commit 428542a
Show file tree
Hide file tree
Showing 5 changed files with 39 additions and 46 deletions.
45 changes: 26 additions & 19 deletions plugin-server/src/utils/db/db.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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<void> {
// 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'
)
}

Expand Down
26 changes: 6 additions & 20 deletions plugin-server/src/worker/ingestion/person-state.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -681,25 +686,6 @@ export class PersonState {

return id
}

private async handleTablesDependingOnPersonID(
sourcePerson: Person,
targetPerson: Person,
tx: TransactionClient
): Promise<void> {
// 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 {
Expand Down
10 changes: 5 additions & 5 deletions plugin-server/tests/main/db.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -845,7 +845,7 @@ describe('DB', () => {
})
})

describe('addFeatureFlagHashKeysForMergedPerson()', () => {
describe('updateCohortsAndFeatureFlagsForMerge()', () => {
let team: Team
let sourcePersonID: Person['id']
let targetPersonID: Person['id']
Expand Down Expand Up @@ -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 () => {
Expand All @@ -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()

Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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()

Expand Down
2 changes: 1 addition & 1 deletion posthog/api/test/test_decide.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
2 changes: 1 addition & 1 deletion posthog/models/feature_flag/flag_matching.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 428542a

Please sign in to comment.