move team and group index resolution to Managers; batch team + group fetches; batch all event and prop writes

Author: eli-r-ph

plugin-server/src/property-defs/property-defs-consumer.ts

@@ -10,17 +10,20 @@ import {
     ClickHouseEvent,
     EventDefinitionType,
     EventPropertyType,
+    GroupTypeToColumnIndex,
     Hub,
     PluginServerService,
     PropertyDefinitionType,
     PropertyDefinitionTypeEnum,
     PropertyType,
     RawClickHouseEvent,
-    ResolvedGroups,
+    TeamId,
 } from '../types'
 import { parseRawClickHouseEvent } from '../utils/event'
 import { status } from '../utils/status'
 import { UUIDT } from '../utils/utils'
+import { GroupTypeManager } from '../worker/ingestion/group-type-manager'
+import { TeamManager } from '../worker/ingestion/team-manager'
 import { PropertyDefsDB } from './services/property-defs-db'
 import {
     getPropertyType,
@@ -32,6 +35,8 @@ import {
 // Must require as `tsc` strips unused `import` statements and just requiring this seems to init some globals
 require('@sentry/tracing')
 
+const BATCH_SIZE = 100
+
 // TODO(eli): wire up LOTS more metrics ASAP!
 
 const propertyDefTypesCounter = new Counter({
@@ -59,8 +64,8 @@ const propDefDroppedCounter = new Counter({
 export type CollectedPropertyDefinitions = {
     // looked up prior to event/prop extraction
     knownTeamIds: Set<number>
-    // looked up prior to event/prop extraction
-    resolvedTeamGroups: Record<number, ResolvedGroups>
+    // looked up prior to event/prop extraction. map of project_id => group_type => group_index
+    resolvedTeamGroups: Record<number, GroupTypeToColumnIndex>
     // known team ID => resolved group_type & group_type_index
     eventDefinitionsById: Record<number, Record<string, EventDefinitionType>>
     // known team ID => deduped properties
@@ -79,6 +84,8 @@ export class PropertyDefsConsumer {
 
     private batchConsumer?: BatchConsumer
     private propertyDefsDB: PropertyDefsDB
+    private teamManager: TeamManager
+    private groupTypeManager: GroupTypeManager
     private isStopping = false
     protected heartbeat = () => {}
     protected promises: Set<Promise<any>> = new Set()
@@ -87,6 +94,8 @@ export class PropertyDefsConsumer {
         this.groupId = hub.PROPERTY_DEFS_CONSUMER_GROUP_ID
         this.topic = hub.PROPERTY_DEFS_CONSUMER_CONSUME_TOPIC
         this.propertyDefsDB = new PropertyDefsDB(hub)
+        this.teamManager = new TeamManager(hub.postgres)
+        this.groupTypeManager = new GroupTypeManager(hub.postgres, this.teamManager)
     }
 
     public get service(): PluginServerService {
@@ -133,7 +142,9 @@ export class PropertyDefsConsumer {
     }
 
     public async handleKafkaBatch(messages: Message[]) {
-        const parsedMessages = await this.runInstrumented('parseKafkaMessages', () => this.parseKafkaBatch(messages))
+        const parsedMessages: ClickHouseEvent[] = await this.runInstrumented('parseKafkaMessages', () =>
+            this.parseKafkaBatch(messages)
+        )
 
         // used to filter and dedup to minimum batch of writable records
         const collected: CollectedPropertyDefinitions = {
@@ -144,18 +155,18 @@ export class PropertyDefsConsumer {
             eventPropertiesById: {},
         }
 
-        const teamsInBatch = this.extractTeamIds(parsedMessages)
-        const groupTeamsInBatch = this.extractGroupTeamIds(parsedMessages, collected.knownTeamIds)
+        const eventTeamIds = parsedMessages.map((msg) => msg.team_id as TeamId)
+        const groupTeamIds = parsedMessages.filter((msg) => msg.event == '$groupidentify').map((msg) => msg.team_id)
 
-        const [knownTeamIds, resolvedTeamGroups] = await Promise.all([
-            this.runInstrumented('resolveTeams', () => this.resolveTeams(this.propertyDefsDB, teamsInBatch)),
-            this.runInstrumented('resolveGroupsForTeams', () =>
-                this.resolveGroupsForTeams(this.propertyDefsDB, groupTeamsInBatch)
+        const [knownTeamIds, resolvedProjectGroups] = await Promise.all([
+            this.runInstrumented('resolveTeams', () => this.teamManager.validateTeamIds(eventTeamIds)),
+            this.runInstrumented('resolveProjectGroupTypeIndices', () =>
+                this.groupTypeManager.fetchGroupTypesIndicesForTeams(groupTeamIds)
             ),
         ])
 
-        collected.knownTeamIds = knownTeamIds
-        collected.resolvedTeamGroups = resolvedTeamGroups
+        collected.knownTeamIds = new Set(knownTeamIds)
+        collected.resolvedTeamGroups = resolvedProjectGroups
 
         console.log('🔁', `Event batch teams and group indices resolved`)
 
@@ -167,38 +178,47 @@ export class PropertyDefsConsumer {
         console.log('🔁', `Property definitions collected`, JSON.stringify(collected, null, 2))
 
         for (const knownTeamId in collected.eventDefinitionsById) {
+            let buffer: EventDefinitionType[] = []
             for (const key in collected.eventDefinitionsById[knownTeamId]) {
                 const eventDef = collected.eventDefinitionsById[knownTeamId][key]
-
+                buffer.push(eventDef)
                 eventDefTypesCounter.inc()
-                status.info('🔁', `Writing event definition`, { eventDef })
 
-                // TODO: Batch all these DB writes
-                void this.scheduleWork(this.propertyDefsDB.writeEventDefinition(eventDef))
+                if (buffer.length === BATCH_SIZE) {
+                    status.info('🔁', `Writing event definition batch of size ${buffer.length}`)
+                    void this.scheduleWork(this.propertyDefsDB.writeEventDefinitionsBatch(buffer))
+                    buffer = []
+                }
             }
         }
 
         for (const knownTeamId in collected.propertyDefinitionsById) {
+            let buffer: PropertyDefinitionType[] = []
             for (const key in collected.propertyDefinitionsById[knownTeamId]) {
                 const propDef: PropertyDefinitionType = collected.propertyDefinitionsById[knownTeamId][key]
-
+                buffer.push(propDef)
                 propertyDefTypesCounter.inc({ type: propDef.property_type?.toString() ?? 'unknown' })
-                status.info('🔁', `Writing property definition`, { propDef })
 
-                // TODO: Batch all these DB writes
-                void this.scheduleWork(this.propertyDefsDB.writePropertyDefinition(propDef))
+                if (buffer.length === BATCH_SIZE) {
+                    status.info('🔁', `Writing property definitions batch of size ${buffer.length}`)
+                    void this.scheduleWork(this.propertyDefsDB.writePropertyDefinitionsBatch(buffer))
+                    buffer = []
+                }
             }
         }
 
         for (const knownTeamId in collected.eventPropertiesById) {
+            let buffer: EventPropertyType[] = []
             for (const key in collected.eventPropertiesById[knownTeamId]) {
                 const eventProp = collected.eventPropertiesById[knownTeamId][key]
-
                 eventPropTypesCounter.inc()
-                status.info('🔁', `Writing event property`, { eventProp })
+                buffer.push(eventProp)
 
-                // TODO: Batch all these DB writes
-                void this.scheduleWork(this.propertyDefsDB.writeEventProperty(eventProp))
+                if (buffer.length === BATCH_SIZE) {
+                    status.info('🔁', `Writing event properties batch of size ${buffer.length}`)
+                    void this.scheduleWork(this.propertyDefsDB.writeEventPropertiesBatch(buffer))
+                    buffer = []
+                }
             }
         }
 
@@ -207,45 +227,6 @@ export class PropertyDefsConsumer {
         status.debug('🔁', `Processed batch`)
     }
 
-    private async resolveTeams(db: PropertyDefsDB, teamIdsInBatch: Set<number>): Promise<Set<number>> {
-        const teamsFound = await db.findTeamIds(Array.from(teamIdsInBatch))
-        return new Set(teamsFound.filter((row) => !teamIdsInBatch.has(row.teamId)).map((row) => row.teamId))
-    }
-
-    private async resolveGroupsForTeams(
-        db: PropertyDefsDB,
-        knownTeamIdsWithGroup: Set<number>
-    ): Promise<Record<number, ResolvedGroups>> {
-        const result = await db.resolveGroupsForTeams(Array.from(knownTeamIdsWithGroup))
-
-        const out: Record<number, ResolvedGroups> = {}
-        result.forEach((row) => {
-            let resolved: ResolvedGroups
-            if (out[row.teamId]) {
-                resolved = out[row.teamId]
-            } else {
-                resolved = {}
-            }
-
-            resolved[row.groupName] = row.groupIndex
-            out[row.teamId] = resolved
-        })
-
-        return out
-    }
-
-    private extractTeamIds(events: ClickHouseEvent[]): Set<number> {
-        return new Set(events.map((event) => event.team_id))
-    }
-
-    private extractGroupTeamIds(events: ClickHouseEvent[], knownTeamIds: Set<number>): Set<number> {
-        return new Set(
-            events
-                .filter((event) => event.event == '$groupidentify' && knownTeamIds.has(event.team_id))
-                .map((event) => event.team_id)
-        )
-    }
-
     private extractPropertyDefinitions(events: ClickHouseEvent[], collected: CollectedPropertyDefinitions) {
         for (const event of events) {
             if (!collected.knownTeamIds.has(event.team_id)) {

plugin-server/src/property-defs/services/property-defs-db.ts

@@ -4,16 +4,6 @@ import { EventDefinitionType, EventPropertyType, Hub, PropertyDefinitionType } f
 import { PostgresUse } from '../../utils/db/postgres'
 import { status } from '../../utils/status'
 
-export interface TeamIdRow {
-    teamId: number
-}
-
-export interface TeamGroupRow {
-    teamId: number
-    groupName: string
-    groupIndex: number
-}
-
 export class PropertyDefsDB {
     constructor(private hub: Hub) {}
 
@@ -34,7 +24,7 @@ export class PropertyDefsDB {
             })
     }
 
-    async writeEventProperiesBatch(eventProperties: EventPropertyType[]) {
+    async writeEventPropertiesBatch(eventProperties: EventPropertyType[]) {
         const values = Array(eventProperties.length).fill('($1, $2, $3, $4)').join(', ')
         const query = `INSERT INTO posthog_eventproperty (event, property, team_id, project_id)
                     VALUES ${values}
@@ -76,7 +66,7 @@ export class PropertyDefsDB {
             )
             .catch((e) => {
                 status.error('🔁', `Error writing property definitions batch`, {
-                    propertyDefinition: propertyDefinition,
+                    propertyDefinition,
                     error: e.message,
                 })
                 throw e
@@ -165,28 +155,6 @@ export class PropertyDefsDB {
             })
     }
 
-    async findTeamIds(teamIds: number[]): Promise<TeamIdRow[]> {
-        const result = await this.hub.postgres.query<TeamIdRow>(
-            PostgresUse.COMMON_READ,
-            `SELECT id AS team_id FROM posthog_team WHERE id = ANY ($1)`,
-            [teamIds],
-            'findTeamIds'
-        )
-
-        return result.rows
-    }
-
-    async resolveGroupsForTeams(teamIds: number[]): Promise<TeamGroupRow[]> {
-        const result = await this.hub.postgres.query<TeamGroupRow>(
-            PostgresUse.COMMON_READ,
-            `SELECT team_ids, group_type_index FROM posthog_grouptype WHERE team_ids = ANY ($1)`,
-            [teamIds],
-            'findTeamIds'
-        )
-
-        return result.rows
-    }
-
     async listPropertyDefinitions(teamId: number): Promise<PropertyDefinitionType[]> {
         const result = await this.hub.postgres.query<PropertyDefinitionType>(
             PostgresUse.COMMON_READ,

plugin-server/src/types.ts

@@ -1186,6 +1186,16 @@ export interface EventPropertyType {
     project_id: number | null
 }
 
+export interface TeamIdRow {
+    teamId: number
+}
+
+export interface TeamGroupRow {
+    teamId: number
+    groupName: string
+    groupIndex: number
+}
+
 export type GroupTypeToColumnIndex = Record<string, GroupTypeIndex>
 
 export enum PropertyUpdateOperation {

plugin-server/src/worker/ingestion/group-type-manager.ts

@@ -1,13 +1,16 @@
-import { GroupTypeIndex, GroupTypeToColumnIndex, ProjectId, Team, TeamId } from '../../types'
+import { GroupTypeIndex, GroupTypeToColumnIndex, ProjectId, Team, TeamGroupRow, TeamId } from '../../types'
 import { PostgresRouter, PostgresUse } from '../../utils/db/postgres'
 import { timeoutGuard } from '../../utils/db/utils'
 import { captureTeamEvent } from '../../utils/posthog'
+import { status } from '../../utils/status'
 import { getByAge } from '../../utils/utils'
 import { TeamManager } from './team-manager'
 
 /** How many unique group types to allow per team */
 export const MAX_GROUP_TYPES_PER_TEAM = 5
 
+const CHUNK_SIZE = 100
+
 export class GroupTypeManager {
     private groupTypesCache: Map<ProjectId, [GroupTypeToColumnIndex, number]>
     private instanceSiteUrl: string
@@ -17,6 +20,77 @@ export class GroupTypeManager {
         this.instanceSiteUrl = instanceSiteUrl || 'unknown'
     }
 
+    // BaseEvent doesn't have project_id yet so still using team_id here.
+    // TODO(eli): see if the underlying Kafka event payload does and we can add it (cc Ben)
+    public async fetchGroupTypesIndicesForTeams(
+        groupTeamIds: TeamId[]
+    ): Promise<Record<number, GroupTypeToColumnIndex>> {
+        const out: Record<number, GroupTypeToColumnIndex> = {}
+        if (groupTeamIds.length === 0) {
+            return out
+        }
+
+        const dedupedTeamIds = new Set(groupTeamIds)
+
+        // first, capture already cached group types and their indexes
+        const cachedTeamIds = new Set(
+            Array.from(dedupedTeamIds).filter((teamId) => {
+                const gtci = this.groupTypesCache.get(teamId as ProjectId) // HACK ALERT!! no way this is a good idea (cc Ben)
+                if (gtci) {
+                    gtci.forEach((entry) => {
+                        // if it's a GroupTypeToColumnIndex, not a number, we can use it
+                        if (typeof entry !== 'number') {
+                            if (!out[teamId]) {
+                                out[teamId] = entry
+                            }
+                        }
+                    })
+                    return true
+                }
+                return false
+            })
+        )
+
+        // finally, figure out what we need to fetch from the DB, then do so in batches
+        const teamIdsToFetch = Array.from(dedupedTeamIds.difference(cachedTeamIds))
+
+        const handles: Promise<TeamGroupRow[]>[] = []
+        for (let i = 0; i < teamIdsToFetch.length; i += CHUNK_SIZE) {
+            const chunk = teamIdsToFetch.slice(i, i + CHUNK_SIZE)
+            handles.push(this.fetchGroupTypeIndicesForTeams(chunk))
+        }
+
+        await Promise.all(handles).then((results) => {
+            results.forEach((foundGTIs) =>
+                foundGTIs.forEach((teamGroupRow) => {
+                    if (!out[teamGroupRow.teamId]) {
+                        out[teamGroupRow.teamId] = {}
+                    }
+                    out[teamGroupRow.teamId][teamGroupRow.groupName] = teamGroupRow.groupIndex as GroupTypeIndex
+                })
+            )
+        })
+
+        return out
+    }
+
+    // TODO(eli): convert to use ProjectId if we can add to ClickHouseEvent
+    private async fetchGroupTypeIndicesForTeams(teamIds: TeamId[]): Promise<TeamGroupRow[]> {
+        const result = await this.postgres
+            .query<TeamGroupRow>(
+                PostgresUse.COMMON_READ,
+                `SELECT team_id, group_type, group_type_index FROM posthog_grouptypemapping WHERE team_id = ANY ($1)`,
+                [teamIds],
+                'findGroupTypeIndicesForTeams'
+            )
+            .catch((e) => {
+                status.error('🔁', `Error fetching group type mappings`, { error: e.message })
+                throw e
+            })
+
+        return result.rows
+    }
+
     public async fetchGroupTypes(projectId: ProjectId): Promise<GroupTypeToColumnIndex> {
         const cachedGroupTypes = getByAge(this.groupTypesCache, projectId)
         if (cachedGroupTypes) {
@@ -26,7 +100,7 @@ export class GroupTypeManager {
         const timeout = timeoutGuard(`Still running "fetchGroupTypes". Timeout warning after 30 sec!`)
         try {
             const { rows } = await this.postgres.query(
-                PostgresUse.COMMON_WRITE,
+                PostgresUse.COMMON_WRITE, // TODO: can we get away with COMMON_READ here? cc Ben
                 `SELECT * FROM posthog_grouptypemapping WHERE project_id = $1`,
                 [projectId],
                 'fetchGroupTypes'