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,11 +1,16 @@
 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> {
+export async function loadPlugin(
+    server: PluginsServer,
+    pluginConfig: PluginConfig,
+    prevConfig?: PluginConfig
+): Promise<boolean> {
     const { plugin } = pluginConfig
 
     if (!plugin) {
@@ -53,6 +58,7 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
                 indexJs,
                 `local ${pluginDigest(plugin)} from "${pluginPath}"!`
             )
+            await inferPluginCapabilities(server, pluginConfig, prevConfig)
             return true
         } else if (plugin.archive) {
             let config: PluginJsonConfig = {}
@@ -72,13 +78,15 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
 
             if (indexJs) {
                 void pluginConfig.vm?.initialize!(server, pluginConfig, indexJs, pluginDigest(plugin))
+                await inferPluginCapabilities(server, pluginConfig, prevConfig)
                 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, prevConfig)
             return true
         } else {
             pluginConfig.vm?.failInitialization!()
@@ -94,3 +102,51 @@ export async function loadPlugin(server: PluginsServer, pluginConfig: PluginConf
     }
     return false
 }
+
+async function inferPluginCapabilities(server: PluginsServer, pluginConfig: PluginConfig, prevConfig?: 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)
+            }
+        }
+    }
+
+    const prevCapabilities = prevConfig?.plugin?.capabilities
+    if (
+        prevCapabilities &&
+        capabilities.jobs?.sort().toString() == prevCapabilities.jobs?.sort().toString() &&
+        capabilities.scheduled_tasks?.sort().toString() == prevCapabilities.scheduled_tasks?.sort().toString() &&
+        capabilities.methods?.sort().toString() == prevCapabilities.methods?.sort().toString()
+    ) {
+        // pass - no change in capabilities
+    } else {
+        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')
                 }

src/worker/plugins/setup.ts

@@ -24,7 +24,7 @@ export async function setupPlugins(server: PluginsServer): Promise<void> {
             pluginConfig.vm = prevConfig.vm
         } else {
             pluginConfig.vm = new LazyPluginVM()
-            pluginVMLoadPromises.push(loadPlugin(server, pluginConfig))
+            pluginVMLoadPromises.push(loadPlugin(server, pluginConfig, prevConfig))
 
             if (prevConfig) {
                 void teardownPlugins(server, prevConfig)

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: {}, // 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>>