Plugin Capabilities implementation

README.md

@@ -132,6 +132,10 @@ Q: Where is teamID populated? At event creation time? (in posthog/posthog? row.p
 
 TODO
 
+#### How VM Extensions Work
+
+TODO
+
 ### End Notes
 
 <a name="f1">1</a>: What are tasks? - TASKS_PER_WORKER - a Piscina setting (https://github.com/piscinajs/piscina#constructor-new-piscinaoptions) -> concurrentTasksPerWorker

src/types.ts

@@ -161,7 +161,13 @@ export interface Plugin {
     from_web?: boolean
     created_at: string
     updated_at: string
-    capabilities?: Record<string, any>
+    capabilities?: Capabilities
+}
+
+export interface Capabilities {
+    jobs: string[]
+    scheduled_tasks: string[]
+    methods: string[]
 }
 
 export interface PluginConfig {
@@ -255,13 +261,13 @@ export type WorkerMethods = {
 export interface PluginConfigVMReponse {
     vm: VM
     methods: {
-        setupPlugin: () => Promise<void>
-        teardownPlugin: () => Promise<void>
-        onEvent: (event: PluginEvent) => Promise<void>
-        onSnapshot: (event: PluginEvent) => Promise<void>
-        processEvent: (event: PluginEvent) => Promise<PluginEvent>
+        setupPlugin: () => Promise<void> | undefined
+        teardownPlugin: () => Promise<void> | undefined
+        onEvent: (event: PluginEvent) => Promise<void> | undefined
+        onSnapshot: (event: PluginEvent) => Promise<void> | undefined
+        processEvent: (event: PluginEvent) => Promise<PluginEvent> | undefined
         // DEPRECATED
-        processEventBatch: (batch: PluginEvent[]) => Promise<PluginEvent[]>
+        processEventBatch: (batch: PluginEvent[]) => Promise<PluginEvent[]> | undefined
     }
     tasks: Record<PluginTaskType, Record<string, PluginTask>>
 }

src/utils/db/sql.ts

@@ -1,4 +1,5 @@
 import {
+    Capabilities,
     Plugin,
     PluginAttachmentDB,
     PluginConfig,
@@ -51,6 +52,25 @@ export async function getPluginConfigRows(server: PluginsServer): Promise<Plugin
     return rows
 }
 
+export async function setPluginCapabilities(
+    server: PluginsServer,
+    pluginConfig: PluginConfig,
+    capabilities: Capabilities
+): Promise<void> {
+    await server.db.postgresQuery(
+        'UPDATE posthog_plugin SET capabilities = ($1) WHERE id = $2',
+        [capabilities, pluginConfig.plugin_id],
+        'setPluginCapabilities'
+    )
+    await server.db.createPluginLogEntry(
+        pluginConfig,
+        PluginLogEntrySource.System,
+        PluginLogEntryType.Info,
+        `Set plugin capabilities (instance ID ${server.instanceId}).`,
+        server.instanceId
+    )
+}
+
 export async function setError(
     server: PluginsServer,
     pluginError: PluginError | null,

src/worker/plugins/loadPlugin.ts

@@ -1,8 +1,9 @@
 import * as fs from 'fs'
 import * as path from 'path'
 
-import { PluginConfig, PluginJsonConfig, PluginsServer } from '../../types'
+import { Capabilities, PluginConfig, PluginJsonConfig, PluginsServer } from '../../types'
 import { processError } from '../../utils/db/error'
+import { setPluginCapabilities } from '../../utils/db/sql'
 import { getFileFromArchive, pluginDigest } from '../../utils/utils'
 
 export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConfig): Promise<boolean> {
@@ -53,6 +54,7 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
                 indexJs,
                 `local ${pluginDigest(plugin)} from "${pluginPath}"!`
             )
+            await inferPluginCapabilities(server, pluginConfig)
             return true
         } else if (plugin.archive) {
             let config: PluginJsonConfig = {}
@@ -72,13 +74,15 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
 
             if (indexJs) {
                 void pluginConfig.vm?.initialize!(server, pluginConfig, indexJs, pluginDigest(plugin))
+                await inferPluginCapabilities(server, pluginConfig)
                 return true
             } else {
                 pluginConfig.vm?.failInitialization!()
                 await processError(server, pluginConfig, `Could not load index.js for ${pluginDigest(plugin)}!`)
             }
         } else if (plugin.plugin_type === 'source' && plugin.source) {
             void pluginConfig.vm?.initialize!(server, pluginConfig, plugin.source, pluginDigest(plugin))
+            await inferPluginCapabilities(server, pluginConfig)
             return true
         } else {
             pluginConfig.vm?.failInitialization!()
@@ -94,3 +98,40 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
     }
     return false
 }
+
+async function inferPluginCapabilities(server: PluginsServer, pluginConfig: PluginConfig) {
+    // infer on load implies there's no lazy loading, but all workers get
+    // these properties loaded
+    const vm = await pluginConfig.vm?.resolveInternalVm
+    const capabilities: Capabilities = { scheduled_tasks: [], jobs: [], methods: [] }
+
+    const tasks = vm?.tasks
+    const methods = vm?.methods
+
+    if (methods) {
+        for (const [key, value] of Object.entries(methods)) {
+            if (value !== undefined) {
+                capabilities.methods.push(key)
+            }
+        }
+    }
+
+    if (tasks?.schedule) {
+        for (const [key, value] of Object.entries(tasks.schedule)) {
+            if (value) {
+                capabilities.scheduled_tasks.push(key)
+            }
+        }
+    }
+
+    if (tasks?.job) {
+        for (const [key, value] of Object.entries(tasks.job)) {
+            if (value) {
+                capabilities.jobs.push(key)
+            }
+        }
+    }
+
+    await setPluginCapabilities(server, pluginConfig, capabilities)
+    pluginConfig.plugin!.capabilities = capabilities
+}

src/worker/plugins/run.ts

@@ -58,7 +58,7 @@ export async function runProcessEvent(server: PluginsServer, event: PluginEvent)
 
             try {
                 returnedEvent = (await processEvent(returnedEvent)) || null
-                if (returnedEvent.team_id != teamId) {
+                if (returnedEvent?.team_id != teamId) {
                     returnedEvent = null // don't try to ingest events with modified teamIDs
                     throw new Error('Illegal Operation: Plugin tried to change teamID')
                 }

tests/helpers/plugins.ts

@@ -48,7 +48,7 @@ export const plugin60: Plugin = {
     is_preinstalled: false,
     created_at: new Date().toISOString(),
     updated_at: new Date().toISOString(),
-    capabilities: {},
+    capabilities: { jobs: [], scheduled_tasks: [], methods: [] }, // inferred on setup
 }
 
 export const pluginAttachment1: PluginAttachmentDB = {
@@ -139,3 +139,11 @@ export function mockPluginTempFolder(indexJs: string, pluginJson?: string): [Plu
         },
     ]
 }
+
+export const mockPluginSourceCode = (indexJs: string): Plugin => ({
+    ...plugin60,
+    archive: null,
+    plugin_type: 'source',
+    url: undefined,
+    source: indexJs,
+})

tests/helpers/sql.ts

@@ -71,8 +71,15 @@ async function insertRow(db: Pool, table: string, object: Record<string, any>):
     const params = Object.keys(object)
         .map((_, i) => `\$${i + 1}`)
         .join(',')
+    const values = Object.values(object).map((value) => {
+        if (Array.isArray(value) && value.length > 0) {
+            return JSON.stringify(value)
+        }
+        return value
+    })
+
     try {
-        await db.query(`INSERT INTO ${table} (${keys}) VALUES (${params})`, Object.values(object))
+        await db.query(`INSERT INTO ${table} (${keys}) VALUES (${params})`, values)
     } catch (error) {
         console.error(`Error on table ${table} when inserting object:\n`, object, '\n', error)
         throw error
@@ -122,11 +129,11 @@ export async function createUserTeamAndOrganization(
         organization_id: organizationId,
         app_urls: [],
         name: 'TEST PROJECT',
-        event_names: JSON.stringify([]),
-        event_names_with_usage: JSON.stringify([]),
-        event_properties: JSON.stringify([]),
-        event_properties_with_usage: JSON.stringify([]),
-        event_properties_numerical: JSON.stringify([]),
+        event_names: [],
+        event_names_with_usage: [],
+        event_properties: [],
+        event_properties_with_usage: [],
+        event_properties_numerical: [],
         created_at: new Date().toISOString(),
         updated_at: new Date().toISOString(),
         anonymize_ips: false,
@@ -140,7 +147,7 @@ export async function createUserTeamAndOrganization(
         api_token: `THIS IS NOT A TOKEN FOR TEAM ${teamId}`,
         test_account_filters: [],
         timezone: 'UTC',
-        data_attributes: JSON.stringify(['data-attr']),
+        data_attributes: ['data-attr'],
     })
 }
 

tests/helpers/sqlMock.ts

@@ -9,5 +9,8 @@ export const getPluginAttachmentRows = (s.getPluginAttachmentRows as unknown) as
 export const getPluginConfigRows = (s.getPluginConfigRows as unknown) as jest.MockedFunction<
     UnPromisify<typeof s.getPluginConfigRows>
 >
+export const setPluginCapabilities = (s.setPluginCapabilities as unknown) as jest.MockedFunction<
+    UnPromisify<typeof s.setPluginCapabilities>
+>
 export const setError = (s.setError as unknown) as jest.MockedFunction<UnPromisify<typeof s.setError>>
 export const disablePlugin = (s.disablePlugin as unknown) as jest.MockedFunction<UnPromisify<void>>

tests/plugins.test.ts

@@ -9,14 +9,15 @@ import { runProcessEvent } from '../src/worker/plugins/run'
 import { loadSchedule, setupPlugins } from '../src/worker/plugins/setup'
 import {
     commonOrganizationId,
+    mockPluginSourceCode,
     mockPluginTempFolder,
     mockPluginWithArchive,
     plugin60,
     pluginAttachment1,
     pluginConfig39,
 } from './helpers/plugins'
 import { resetTestDatabase } from './helpers/sql'
-import { getPluginAttachmentRows, getPluginConfigRows, getPluginRows } from './helpers/sqlMock'
+import { getPluginAttachmentRows, getPluginConfigRows, getPluginRows, setPluginCapabilities } from './helpers/sqlMock'
 
 jest.mock('../src/utils/db/sql')
 jest.mock('../src/utils/status')
@@ -48,8 +49,17 @@ test('setupPlugins and runProcessEvent', async () => {
     expect(getPluginRows).toHaveBeenCalled()
     expect(getPluginAttachmentRows).toHaveBeenCalled()
     expect(getPluginConfigRows).toHaveBeenCalled()
-
-    expect(Array.from(plugins.entries())).toEqual([[60, plugin60]])
+    expect(setPluginCapabilities).toHaveBeenCalled()
+
+    expect(Array.from(plugins.entries())).toEqual([
+        [
+            60,
+            {
+                ...plugin60,
+                capabilities: { jobs: [], scheduled_tasks: [], methods: ['processEvent', 'processEventBatch'] },
+            },
+        ],
+    ])
     expect(Array.from(pluginConfigs.keys())).toEqual([39])
 
     const pluginConfig = pluginConfigs.get(39)!
@@ -445,6 +455,102 @@ test('plugin config order', async () => {
     expect(returnedEvent2!.properties!.plugins).toEqual([61, 60, 62])
 })
 
+test('plugin with archive loads capabilities', async () => {
+    getPluginRows.mockReturnValueOnce([
+        mockPluginWithArchive(`
+            function setupPlugin (meta) { meta.global.key = 'value' }
+            function processEvent (event, meta) { event.properties={"x": 1}; return event }
+        `),
+    ])
+    getPluginConfigRows.mockReturnValueOnce([pluginConfig39])
+    getPluginAttachmentRows.mockReturnValueOnce([pluginAttachment1])
+
+    await setupPlugins(mockServer)
+    const { pluginConfigs } = mockServer
+
+    const pluginConfig = pluginConfigs.get(39)!
+    expect(pluginConfig.plugin!.capabilities!.methods.sort()).toEqual([
+        'processEvent',
+        'processEventBatch',
+        'setupPlugin',
+    ])
+    expect(pluginConfig.plugin!.capabilities!.jobs).toHaveLength(0)
+    expect(pluginConfig.plugin!.capabilities!.scheduled_tasks).toHaveLength(0)
+})
+
+test('plugin with archive loads all capabilities, no random caps', async () => {
+    getPluginRows.mockReturnValueOnce([
+        mockPluginWithArchive(`
+            function processEvent (event, meta) { event.properties={"x": 1}; return event }
+            function randomFunction (event, meta) { return event}
+            function onEvent (event, meta) { return event }
+
+            function runEveryHour(meta) {console.log('1')}
+
+            export const jobs = {
+                x: (event, meta) => console.log(event)
+            }
+        `),
+    ])
+    getPluginConfigRows.mockReturnValueOnce([pluginConfig39])
+    getPluginAttachmentRows.mockReturnValueOnce([pluginAttachment1])
+
+    await setupPlugins(mockServer)
+    const { pluginConfigs } = mockServer
+
+    const pluginConfig = pluginConfigs.get(39)!
+    expect(pluginConfig.plugin!.capabilities!.methods.sort()).toEqual(['onEvent', 'processEvent', 'processEventBatch'])
+    expect(pluginConfig.plugin!.capabilities!.jobs).toEqual(['x'])
+    expect(pluginConfig.plugin!.capabilities!.scheduled_tasks).toEqual(['runEveryHour'])
+})
+
+test('plugin with source file loads capabilities', async () => {
+    const [plugin, unlink] = mockPluginTempFolder(`
+        function processEvent (event, meta) { event.properties={"x": 1}; return event }
+        function randomFunction (event, meta) { return event}
+        function onEvent (event, meta) { return event }
+    `)
+
+    getPluginRows.mockReturnValueOnce([plugin])
+    getPluginConfigRows.mockReturnValueOnce([pluginConfig39])
+    getPluginAttachmentRows.mockReturnValueOnce([pluginAttachment1])
+
+    await setupPlugins(mockServer)
+    const { pluginConfigs } = mockServer
+
+    const pluginConfig = pluginConfigs.get(39)!
+    expect(pluginConfig.plugin!.capabilities!.methods.sort()).toEqual(['onEvent', 'processEvent', 'processEventBatch'])
+    expect(pluginConfig.plugin!.capabilities!.jobs).toEqual([])
+    expect(pluginConfig.plugin!.capabilities!.scheduled_tasks).toEqual([])
+
+    unlink()
+})
+
+test('plugin with source code loads capabilities', async () => {
+    const source_code = `
+        function processEvent (event, meta) { event.properties={"x": 1}; return event }
+        function randomFunction (event, meta) { return event}
+        function onSnapshot (event, meta) { return event }
+    `
+    getPluginRows.mockReturnValueOnce([mockPluginSourceCode(source_code)])
+
+    getPluginConfigRows.mockReturnValueOnce([pluginConfig39])
+    getPluginAttachmentRows.mockReturnValueOnce([pluginAttachment1])
+
+    await setupPlugins(mockServer)
+    const { pluginConfigs } = mockServer
+
+    const pluginConfig = pluginConfigs.get(39)!
+
+    expect(pluginConfig.plugin!.capabilities!.methods.sort()).toEqual([
+        'onSnapshot',
+        'processEvent',
+        'processEventBatch',
+    ])
+    expect(pluginConfig.plugin!.capabilities!.jobs).toEqual([])
+    expect(pluginConfig.plugin!.capabilities!.scheduled_tasks).toEqual([])
+})
+
 test('reloading plugins after config changes', async () => {
     const makePlugin = (id: number, updated_at = '2020-11-02'): any => ({
         ...plugin60,

tests/sql.test.ts

@@ -105,7 +105,7 @@ test('getPluginRows', async () => {
             url: 'https://www.npmjs.com/package/posthog-maxmind-plugin',
             created_at: expect.anything(),
             updated_at: expect.anything(),
-            capabilities: {},
+            capabilities: { jobs: [], scheduled_tasks: [], methods: [] },
         },
     ]