Reorient ActionManager to group by teamId for practicality

src/main/pluginsServer.ts

@@ -153,9 +153,9 @@ export async function startPluginsServer(
                 await scheduleControl?.reloadSchedule()
             },
             'reload-action': async (message) =>
-                await piscina?.broadcastTask({ task: 'reloadAction', args: { actionId: parseInt(message) } }),
+                await piscina?.broadcastTask({ task: 'reloadAction', args: JSON.parse(message) }),
             'drop-action': async (message) =>
-                await piscina?.broadcastTask({ task: 'dropAction', args: { actionId: parseInt(message) } }),
+                await piscina?.broadcastTask({ task: 'dropAction', args: JSON.parse(message) }),
         })
         await pubSub.start()
 

src/utils/db/db.ts

@@ -32,7 +32,7 @@ import {
     RawOrganization,
     RawPerson,
     SessionRecordingEvent,
-    TeamId,
+    Team,
     TimestampFormat,
 } from '../../types'
 import { instrumentQuery } from '../metrics'
@@ -726,29 +726,32 @@ export class DB {
 
     // Action & ActionStep
 
-    public async fetchAllActionsMap(): Promise<Record<Action['id'], Action>> {
+    public async fetchAllActionsGroupedByTeam(): Promise<Record<Team['id'], Record<Action['id'], Action>>> {
         const rawActions: RawAction[] = (
             await this.postgresQuery(`SELECT * FROM posthog_action WHERE deleted = FALSE`, undefined, 'fetchActions')
         ).rows
-        const actionSteps: ActionStep[] = (
+        const actionSteps: (ActionStep & { team_id: Team['id'] })[] = (
             await this.postgresQuery(
-                `SELECT posthog_actionstep.*, posthog_action.deleted FROM posthog_actionstep
-            JOIN posthog_action ON (posthog_action.id = posthog_actionstep.action_id)
-            WHERE posthog_action.deleted = FALSE`,
+                `SELECT posthog_actionstep.*, posthog_action.team_id
+                    FROM posthog_actionstep JOIN posthog_action ON (posthog_action.id = posthog_actionstep.action_id)
+                    WHERE posthog_action.deleted = FALSE`,
                 undefined,
                 'fetchActionSteps'
             )
         ).rows
-        const actionsMap: Record<Action['id'], Action> = {}
+        const actions: Record<Team['id'], Record<Action['id'], Action>> = {}
         for (const rawAction of rawActions) {
-            actionsMap[rawAction.id] = { ...rawAction, steps: [] }
+            if (!actions[rawAction.team_id]) {
+                actions[rawAction.team_id] = {}
+            }
+            actions[rawAction.team_id][rawAction.id] = { ...rawAction, steps: [] }
         }
         for (const actionStep of actionSteps) {
-            if (actionStep.action_id in actionsMap) {
-                actionsMap[actionStep.action_id].steps.push(actionStep)
+            if (actions[actionStep.team_id]?.[actionStep.action_id]) {
+                actions[actionStep.team_id][actionStep.action_id].steps.push(actionStep)
             }
         }
-        return actionsMap
+        return actions
     }
 
     public async fetchAction(id: Action['id']): Promise<Action | null> {
@@ -771,7 +774,7 @@ export class DB {
 
     // Team Internal Metrics
 
-    public async fetchInternalMetricsTeam(): Promise<TeamId | null> {
+    public async fetchInternalMetricsTeam(): Promise<Team['id'] | null> {
         const { rows } = await this.postgresQuery(
             `
             SELECT posthog_team.id as team_id

src/utils/utils.ts

@@ -567,6 +567,37 @@ export function filterIncrementProperties(incrementProperties: unknown): Record<
     return filteredIncrementProperties
 }
 
+export function groupBy<T extends Record<string, any>, K extends keyof T>(
+    objects: T[],
+    key: K,
+    flat?: false
+): Record<T[K], T[]>
+export function groupBy<T extends Record<string, any>, K extends keyof T>(
+    objects: T[],
+    key: K,
+    flat: true
+): Record<T[K], T>
+export function groupBy<T extends Record<string, any>, K extends keyof T>(
+    objects: T[],
+    key: K,
+    flat = false
+): Record<T[K], T[] | T> {
+    return flat
+        ? objects.reduce((grouping, currentItem) => {
+              if (currentItem[key] in grouping) {
+                  throw new Error(
+                      `Key "${key}" has more than one matching value, which is not allowed in flat groupBy!`
+                  )
+              }
+              grouping[currentItem[key]] = currentItem
+              return grouping
+          }, {} as Record<T[K], T>)
+        : objects.reduce((grouping, currentItem) => {
+              ;(grouping[currentItem[key]] = grouping[currentItem[key]] || []).push(currentItem)
+              return grouping
+          }, {} as Record<T[K], T[]>)
+}
+
 export function clamp(value: number, min: number, max: number): number {
     return value > max ? max : value < min ? min : value
 }

src/worker/ingestion/action-manager.ts

@@ -1,8 +1,9 @@
-import { Action, PluginsServerConfig } from '../../types'
+import { Action, Team } from '../../types'
 import { DB } from '../../utils/db/db'
 import { status } from '../../utils/status'
 
-type ActionCache = Record<Action['id'], Action>
+export type ActionMap = Record<Action['id'], Action>
+type ActionCache = Record<Team['id'], ActionMap>
 
 export class ActionManager {
     private ready: boolean
@@ -20,48 +21,52 @@ export class ActionManager {
         this.ready = true
     }
 
-    public getAction(id: Action['id']): Action | undefined {
+    public getTeamActions(teamId: Team['id']): ActionMap | null {
         if (!this.ready) {
             throw new Error('ActionManager is not ready! Run actionManager.prepare() before this')
         }
-        return this.actionCache[id]
+        return this.actionCache[teamId] || null
     }
 
     public async reloadAllActions(): Promise<void> {
-        this.actionCache = await this.db.fetchAllActionsMap()
+        this.actionCache = await this.db.fetchAllActionsGroupedByTeam()
         status.info('🍿', 'Fetched all actions from DB anew')
     }
 
-    public async reloadAction(id: Action['id']): Promise<void> {
-        const refetchedAction = await this.db.fetchAction(id)
+    public async reloadAction(teamId: Team['id'], actionId: Action['id']): Promise<void> {
+        const refetchedAction = await this.db.fetchAction(actionId)
+        const wasCachedAlready = teamId in this.actionCache && actionId in this.actionCache[teamId]
         if (refetchedAction) {
             status.info(
                 '🍿',
-                id in this.actionCache ? `Refetched action ID ${id} from DB` : `Fetched new action ID ${id} from DB`
+                wasCachedAlready
+                    ? `Refetched action ID ${actionId} (team ID ${teamId}) from DB`
+                    : `Fetched new action ID ${actionId} (team ID ${teamId}) from DB`
             )
-            this.actionCache[id] = refetchedAction
-        } else if (id in this.actionCache) {
+            this.actionCache[teamId][actionId] = refetchedAction
+        } else if (wasCachedAlready) {
             status.info(
                 '🍿',
-                `Tried to fetch action ID ${id} from DB, but it wasn't found in DB, so deleted from cache instead`
+                `Tried to fetch action ID ${actionId} (team ID ${teamId}) from DB, but it wasn't found in DB, so deleted from cache instead`
             )
-            delete this.actionCache[id]
+            delete this.actionCache[teamId][actionId]
         } else {
             status.info(
                 '🍿',
-                `Tried to fetch action ID ${id} from DB, but it wasn't found in DB or cache, so did nothing instead`
+                `Tried to fetch action ID ${actionId} (team ID ${teamId}) from DB, but it wasn't found in DB or cache, so did nothing instead`
             )
         }
     }
 
-    public dropAction(id: Action['id']): void {
-        if (id in this.actionCache) {
-            status.info('🍿', `Deleted action ID ${id} from cache`)
-            delete this.actionCache[id]
+    public dropAction(teamId: Team['id'], actionId: Action['id']): void {
+        const wasCachedAlready = teamId in this.actionCache && actionId in this.actionCache[teamId]
+        if (wasCachedAlready) {
+            status.info('🍿', `Deleted action ID ${actionId} (team ID ${teamId}) from cache`)
+            delete this.actionCache[teamId][actionId]
         } else {
             status.info(
                 '🍿',
-                `Tried to delete action ID ${id} from cache, but it wasn't found in cache, so did nothing instead`
+                `Tried to delete action ID ${actionId} (team ID ${teamId}) from cache, but it wasn't found in cache, so did nothing instead`
             )
         }
     }

src/worker/ingestion/process-event.ts

@@ -24,7 +24,6 @@ import { Client } from '../../utils/celery/client'
 import { DB } from '../../utils/db/db'
 import { KafkaProducerWrapper } from '../../utils/db/kafka-producer-wrapper'
 import { elementsToString, personInitialAndUTMProperties, sanitizeEventName, timeoutGuard } from '../../utils/db/utils'
-import { PubSub } from '../../utils/pubsub'
 import { status } from '../../utils/status'
 import { castTimestampOrNow, filterIncrementProperties, UUID, UUIDT } from '../../utils/utils'
 import { ActionManager } from './action-manager'

src/worker/tasks.ts

@@ -1,6 +1,6 @@
 import { PluginEvent } from '@posthog/plugin-scaffold/src/types'
 
-import { Action, EnqueuedJob, Hub, PluginTaskType } from '../types'
+import { Action, EnqueuedJob, Hub, PluginTaskType, Team } from '../types'
 import { ingestEvent } from './ingestion/ingest-event'
 import { runOnEvent, runOnSnapshot, runPluginTask, runProcessEvent, runProcessEventBatch } from './plugins/run'
 import { loadSchedule, setupPlugins } from './plugins/setup'
@@ -48,11 +48,11 @@ export const workerTasks: Record<string, TaskRunner> = {
     reloadAllActions: async (hub) => {
         return await hub.eventsProcessor.actionManager.reloadAllActions()
     },
-    reloadAction: async (hub, args: { actionId: Action['id'] }) => {
-        return await hub.eventsProcessor.actionManager.reloadAction(args.actionId)
+    reloadAction: async (hub, args: { teamId: Team['id']; actionId: Action['id'] }) => {
+        return await hub.eventsProcessor.actionManager.reloadAction(args.teamId, args.actionId)
     },
-    dropAction: (hub, args: { actionId: Action['id'] }) => {
-        return hub.eventsProcessor.actionManager.dropAction(args.actionId)
+    dropAction: (hub, args: { teamId: Team['id']; actionId: Action['id'] }) => {
+        return hub.eventsProcessor.actionManager.dropAction(args.teamId, args.actionId)
     },
     teardownPlugins: async (hub) => {
         await teardownPlugins(hub)

tests/postgres/worker.test.ts

@@ -320,15 +320,15 @@ describe('createTaskRunner()', () => {
     })
 
     it('handles `reloadAction` task', async () => {
-        await taskRunner({ task: 'reloadAction', args: { actionId: 777 } })
+        await taskRunner({ task: 'reloadAction', args: { teamId: 2, actionId: 777 } })
 
-        expect(hub.eventsProcessor.actionManager.reloadAction).toHaveBeenCalledWith(777)
+        expect(hub.eventsProcessor.actionManager.reloadAction).toHaveBeenCalledWith(2, 777)
     })
 
     it('handles `dropAction` task', async () => {
-        await taskRunner({ task: 'dropAction', args: { actionId: 777 } })
+        await taskRunner({ task: 'dropAction', args: { teamId: 2, actionId: 777 } })
 
-        expect(hub.eventsProcessor.actionManager.dropAction).toHaveBeenCalledWith(777)
+        expect(hub.eventsProcessor.actionManager.dropAction).toHaveBeenCalledWith(2, 777)
     })
 
     it('handles `teardownPlugin` task', async () => {

tests/shared/db.test.ts

@@ -0,0 +1,59 @@
+import { Hub, PropertyOperator } from '../../src/types'
+import { DB } from '../../src/utils/db/db'
+import { createHub } from '../../src/utils/db/hub'
+import { ActionManager } from '../../src/worker/ingestion/action-manager'
+import { resetTestDatabase } from '../helpers/sql'
+
+describe('DB', () => {
+    let hub: Hub
+    let closeServer: () => Promise<void>
+    let db: DB
+
+    beforeEach(async () => {
+        ;[hub, closeServer] = await createHub()
+        await resetTestDatabase()
+        db = hub.db
+    })
+
+    afterEach(async () => {
+        await closeServer()
+    })
+
+    const TEAM_ID = 2
+    const ACTION_ID = 69
+    const ACTION_STEP_ID = 913
+
+    test('fetchAllActionsGroupedByTeam', async () => {
+        const action = await db.fetchAllActionsGroupedByTeam()
+
+        expect(action).toMatchObject({
+            [TEAM_ID]: {
+                [ACTION_ID]: {
+                    id: ACTION_ID,
+                    name: 'Test Action',
+                    deleted: false,
+                    post_to_slack: false,
+                    slack_message_format: '',
+                    is_calculating: false,
+                    steps: [
+                        {
+                            id: ACTION_STEP_ID,
+                            action_id: ACTION_ID,
+                            tag_name: null,
+                            text: null,
+                            href: null,
+                            selector: null,
+                            url: null,
+                            url_matching: null,
+                            name: null,
+                            event: null,
+                            properties: [
+                                { type: 'event', operator: PropertyOperator.Exact, key: 'foo', value: ['bar'] },
+                            ],
+                        },
+                    ],
+                },
+            },
+        })
+    })
+})

tests/utils.test.ts

@@ -8,6 +8,7 @@ import {
     getFileFromArchive,
     getFileFromTGZ,
     getFileFromZip,
+    groupBy,
     sanitizeSqlIdentifier,
     setLogLevel,
     UUID,
@@ -320,3 +321,76 @@ describe('escapeClickHouseString', () => {
         expect(sanitizedString).toStrictEqual("insert\\'escape \\\\")
     })
 })
+
+describe('groupBy', () => {
+    it('groups simple objects', () => {
+        const objects = [
+            { i: 2, foo: 'x' },
+            { i: 2, foo: 'y' },
+            { i: 4, foo: 'x' },
+            { i: 7, foo: 'z' },
+        ]
+
+        const groupingByI = groupBy(objects, 'i')
+        expect(groupingByI).toEqual({
+            2: [
+                { i: 2, foo: 'x' },
+                { i: 2, foo: 'y' },
+            ],
+            4: [{ i: 4, foo: 'x' }],
+            7: [{ i: 7, foo: 'z' }],
+        })
+
+        const groupingByFoo = groupBy(objects, 'foo')
+        expect(groupingByFoo).toEqual({
+            x: [
+                { i: 2, foo: 'x' },
+                { i: 4, foo: 'x' },
+            ],
+            y: [{ i: 2, foo: 'y' }],
+            z: [{ i: 7, foo: 'z' }],
+        })
+    })
+
+    it('handles undefineds', () => {
+        const objects = [{ i: 2, foo: 'x' }, { i: 2, foo: 'y' }, { i: 4, foo: 'x' }, { foo: 'z' }]
+
+        const groupingByI = groupBy(objects, 'i')
+        expect(groupingByI).toEqual({
+            2: [
+                { i: 2, foo: 'x' },
+                { i: 2, foo: 'y' },
+            ],
+            4: [{ i: 4, foo: 'x' }],
+            undefined: [{ foo: 'z' }],
+        })
+    })
+
+    it('works in flat mode', () => {
+        const objects = [
+            { i: 2, foo: 'x' },
+            { i: 4, foo: 'x' },
+            { i: 7, foo: 'z' },
+        ]
+
+        const groupingByI = groupBy(objects, 'i', true)
+        expect(groupingByI).toEqual({
+            2: { i: 2, foo: 'x' },
+            4: { i: 4, foo: 'x' },
+            7: { i: 7, foo: 'z' },
+        })
+    })
+
+    it("doesn't work in flat mode if multiple values match a single key", () => {
+        const objects = [
+            { i: 2, foo: 'x' },
+            { i: 2, foo: 'y' },
+            { i: 4, foo: 'x' },
+            { i: 7, foo: 'z' },
+        ]
+
+        expect(() => groupBy(objects, 'i', true)).toThrowError(
+            'Key "i" has more than one matching value, which is not allowed in flat groupBy!'
+        )
+    })
+})