diff --git a/src/main/pluginsServer.ts b/src/main/pluginsServer.ts index e8610dd5..65155366 100644 --- a/src/main/pluginsServer.ts +++ b/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() diff --git a/src/utils/db/db.ts b/src/utils/db/db.ts index 3006d952..0ac03844 100644 --- a/src/utils/db/db.ts +++ b/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> { + public async fetchAllActionsGroupedByTeam(): Promise>> { 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 = {} + const actions: Record> = {} 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 { @@ -771,7 +774,7 @@ export class DB { // Team Internal Metrics - public async fetchInternalMetricsTeam(): Promise { + public async fetchInternalMetricsTeam(): Promise { const { rows } = await this.postgresQuery( ` SELECT posthog_team.id as team_id diff --git a/src/utils/utils.ts b/src/utils/utils.ts index bc8ba6f2..4fe474ae 100644 --- a/src/utils/utils.ts +++ b/src/utils/utils.ts @@ -567,6 +567,37 @@ export function filterIncrementProperties(incrementProperties: unknown): Record< return filteredIncrementProperties } +export function groupBy, K extends keyof T>( + objects: T[], + key: K, + flat?: false +): Record +export function groupBy, K extends keyof T>( + objects: T[], + key: K, + flat: true +): Record +export function groupBy, K extends keyof T>( + objects: T[], + key: K, + flat = false +): Record { + 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) + : objects.reduce((grouping, currentItem) => { + ;(grouping[currentItem[key]] = grouping[currentItem[key]] || []).push(currentItem) + return grouping + }, {} as Record) +} + export function clamp(value: number, min: number, max: number): number { return value > max ? max : value < min ? min : value } diff --git a/src/worker/ingestion/action-manager.ts b/src/worker/ingestion/action-manager.ts index 33659926..f11820ca 100644 --- a/src/worker/ingestion/action-manager.ts +++ b/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 +export type ActionMap = Record +type ActionCache = Record 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 { - 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 { - const refetchedAction = await this.db.fetchAction(id) + public async reloadAction(teamId: Team['id'], actionId: Action['id']): Promise { + 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` ) } } diff --git a/src/worker/ingestion/process-event.ts b/src/worker/ingestion/process-event.ts index 24722edd..4ce55bcb 100644 --- a/src/worker/ingestion/process-event.ts +++ b/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' diff --git a/src/worker/tasks.ts b/src/worker/tasks.ts index 59f9ba78..823bccfe 100644 --- a/src/worker/tasks.ts +++ b/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 = { 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) diff --git a/tests/postgres/worker.test.ts b/tests/postgres/worker.test.ts index 2b612a25..51534e87 100644 --- a/tests/postgres/worker.test.ts +++ b/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 () => { diff --git a/tests/shared/db.test.ts b/tests/shared/db.test.ts new file mode 100644 index 00000000..55534113 --- /dev/null +++ b/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 + 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'] }, + ], + }, + ], + }, + }, + }) + }) +}) diff --git a/tests/utils.test.ts b/tests/utils.test.ts index 6e853bdc..7904d13c 100644 --- a/tests/utils.test.ts +++ b/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!' + ) + }) +}) diff --git a/tests/worker/ingestion/action-manager.test.ts b/tests/worker/ingestion/action-manager.test.ts index 42b7f866..360fc400 100644 --- a/tests/worker/ingestion/action-manager.test.ts +++ b/tests/worker/ingestion/action-manager.test.ts @@ -1,4 +1,4 @@ -import { Hub, PropertyOperator, RawAction } from '../../../src/types' +import { Hub, PropertyOperator } from '../../../src/types' import { createHub } from '../../../src/utils/db/hub' import { ActionManager } from '../../../src/worker/ingestion/action-manager' import { resetTestDatabase } from '../../helpers/sql' @@ -14,17 +14,20 @@ describe('ActionManager', () => { actionManager = new ActionManager(hub.db) await actionManager.prepare() }) + afterEach(async () => { await closeServer() }) - it('returns the correct action', async () => { - const ACTION_ID = 69 - const ACTION_STEP_ID = 913 + const TEAM_ID = 2 + const ACTION_ID = 69 + const ACTION_STEP_ID = 913 - const action = actionManager.getAction(ACTION_ID) + it('returns the correct actions generally', async () => { + const action = actionManager.getTeamActions(TEAM_ID) - expect(action).toMatchObject({ + expect(Object.values(action!).length).toEqual(1) + expect(action![ACTION_ID]).toMatchObject({ id: ACTION_ID, name: 'Test Action', deleted: false, @@ -55,11 +58,12 @@ describe('ActionManager', () => { ) // This is normally dispatched by Django and broadcasted by Piscina - await actionManager.reloadAction(ACTION_ID) + await actionManager.reloadAction(TEAM_ID, ACTION_ID) - const reloadedAction = actionManager.getAction(ACTION_ID) + const reloadedAction = actionManager.getTeamActions(TEAM_ID) - expect(reloadedAction).toMatchObject({ + expect(Object.values(action!).length).toEqual(1) + expect(reloadedAction![ACTION_ID]).toMatchObject({ id: ACTION_ID, name: 'Test Action', deleted: false, @@ -84,10 +88,48 @@ describe('ActionManager', () => { }) // This is normally dispatched by Django and broadcasted by Piscina - actionManager.dropAction(ACTION_ID) + actionManager.dropAction(TEAM_ID, ACTION_ID) + + const droppedAction = actionManager.getTeamActions(TEAM_ID) + + expect(Object.values(droppedAction!).length).toEqual(0) + }) + + it('returns the correct actions when deleted = TRUE', async () => { + const action = actionManager.getTeamActions(TEAM_ID) + + expect(Object.values(action!).length).toEqual(1) + expect(action![ACTION_ID]).toMatchObject({ + 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'] }], + }, + ], + }) + + await hub.db.postgresQuery(`UPDATE posthog_action SET deleted = TRUE WHERE id = $1`, [ACTION_ID], 'testKey') + + // This is normally dispatched by Django and broadcasted by Piscina + await actionManager.reloadAction(TEAM_ID, ACTION_ID) - const droppedAction = actionManager.getAction(ACTION_ID) + const droppedAction = actionManager.getTeamActions(TEAM_ID) - expect(droppedAction).toBeUndefined() + expect(Object.values(droppedAction!).length).toEqual(0) }) })