Plugin storage

.github/workflows/ci.yml

@@ -65,7 +65,7 @@ jobs:
               env:
                   DATABASE_URL: 'postgres://postgres:postgres@localhost:${{ job.services.postgres.ports[5432] }}/postgres'
                   REDIS_URL: 'redis://localhost'
-              run: yarn task:db-init
+              run: yarn task:init-test-db
 
             - name: Test with Jest
               env:
@@ -113,7 +113,7 @@ jobs:
               env:
                   DATABASE_URL: 'postgres://postgres:postgres@localhost:${{ job.services.postgres.ports[5432] }}/postgres'
                   REDIS_URL: 'redis://localhost'
-              run: yarn task:db-init
+              run: yarn task:init-test-db
 
             - name: Run benchmarks
               env:

benchmarks/worker.benchmark.ts

@@ -2,9 +2,9 @@ import { makePiscina } from '../src/worker/piscina'
 import { defaultConfig } from '../src/config'
 import { PluginEvent } from '@posthog/plugin-scaffold/src/types'
 import { performance } from 'perf_hooks'
-import { mockJestWithIndex } from '../tests/helpers/plugins'
 import * as os from 'os'
 import { LogLevel } from '../src/types'
+import { resetTestDatabase } from '../tests/helpers/sql'
 
 jest.mock('../src/sql')
 jest.setTimeout(600000) // 600 sec timeout
@@ -61,13 +61,12 @@ async function processCountEvents(piscina: ReturnType<typeof makePiscina>, count
     }
 }
 
-function setupPiscina(workers: number, code: string, tasksPerWorker: number) {
+function setupPiscina(workers: number, tasksPerWorker: number) {
     return makePiscina({
         ...defaultConfig,
         WORKER_CONCURRENCY: workers,
         TASKS_PER_WORKER: tasksPerWorker,
         LOG_LEVEL: LogLevel.Log,
-        __jestMock: mockJestWithIndex(code),
     })
 }
 
@@ -119,6 +118,7 @@ test('piscina worker benchmark', async () => {
 
     const results: Array<Record<string, string | number>> = []
     for (const { testName, events: _events, testCode } of tests) {
+        await resetTestDatabase(testCode)
         const events = isLightDevRun ? _events / 10 : _events
         for (const batchSize of [1, 10, 100].filter((size) => size <= events)) {
             const result: Record<string, any> = {
@@ -128,7 +128,7 @@ test('piscina worker benchmark', async () => {
                 batchSize,
             }
             for (const threads of workerThreads) {
-                const piscina = setupPiscina(threads, testCode, 100)
+                const piscina = setupPiscina(threads, 100)
 
                 // warmup
                 await processCountEvents(piscina, threads * 4)

jest.config.js

@@ -3,6 +3,6 @@ module.exports = {
     testEnvironment: 'node',
     clearMocks: true,
     coverageProvider: 'v8',
-    setupFilesAfterEnv: ['./jest.setup.pg-mock.js', './jest.setup.kafka-mock.js', './jest.setup.fetch-mock.js'],
+    setupFilesAfterEnv: ['./jest.setup.kafka-mock.js', './jest.setup.fetch-mock.js'],
     testMatch: ['<rootDir>/tests/**/*.test.ts', '<rootDir>/benchmarks/**/*.benchmark.ts'],
 }

package.json

@@ -1,11 +1,11 @@
 {
     "name": "posthog-plugin-server",
-    "version": "0.6.0",
+    "version": "0.6.1",
     "description": "PostHog Plugin Server",
     "types": "dist/src/index.d.ts",
     "main": "dist/src/index.js",
     "scripts": {
-        "test": "jest --testPathIgnorePatterns='benchmarks/'",
+        "test": "jest --runInBand tests/*.test.ts",
         "benchmark": "node --expose-gc node_modules/.bin/jest --runInBand benchmarks/",
         "start": "yarn start:dev",
         "start:dist": "node dist/src/index.js --base-dir ../posthog",
@@ -18,7 +18,7 @@
         "prettier": "prettier --write .",
         "prettier:check": "prettier --check .",
         "prepublishOnly": "yarn test && yarn build",
-        "task:db-init": "ts-node ./tasks/db-init.ts"
+        "task:init-test-db": "NODE_ENV=test ts-node ./tasks/init-test-db.ts"
     },
     "bin": {
         "posthog-plugin-server": "bin/posthog-plugin-server"
@@ -52,7 +52,7 @@
         "@babel/core": "^7.0.0",
         "@babel/preset-env": "^7.0.0",
         "@babel/preset-typescript": "^7.8.3",
-        "@posthog/plugin-scaffold": "0.2.5",
+        "@posthog/plugin-scaffold": "0.2.6",
         "@types/adm-zip": "^0.4.33",
         "@types/ioredis": "^4.17.7",
         "@types/jest": "^26.0.15",

src/celery/broker.ts

@@ -168,7 +168,6 @@ export default class RedisBroker implements Pausable {
                 if (body) {
                     callback(body)
                 }
-                Promise.resolve()
                 return body
             })
             .then((body) => {

src/celery/worker.ts

@@ -205,7 +205,6 @@ export class Worker extends Base implements Queue {
      *
      * @todo implement here
      */
-    // eslint-disable-next-line class-methods-use-this
     public async stop(): Promise<void> {
         const taskCount = this.activeTasks.size
         if (taskCount > 0) {

src/config.ts

@@ -4,9 +4,11 @@ export const defaultConfig = overrideWithEnv(getDefaultConfig())
 export const configHelp = getConfigHelp()
 
 export function getDefaultConfig(): PluginsServerConfig {
+    const isTestEnv = process.env.NODE_ENV === 'test'
+
     return {
         CELERY_DEFAULT_QUEUE: 'celery',
-        DATABASE_URL: 'postgres://localhost:5432/posthog',
+        DATABASE_URL: isTestEnv ? 'postgres://localhost:5432/posthog_test' : 'postgres://localhost:5432/posthog',
         KAFKA_ENABLED: false,
         KAFKA_HOSTS: null,
         KAFKA_CLIENT_CERT_B64: null,

src/extensions/storage.ts

@@ -0,0 +1,35 @@
+import { PluginConfig, PluginsServer } from '../types'
+import { StorageExtension } from '@posthog/plugin-scaffold'
+
+export function createStorage(server: PluginsServer, pluginConfig: PluginConfig): StorageExtension {
+    const get = async function (key: string, defaultValue: unknown): Promise<unknown> {
+        const result = await server.db.query(
+            'SELECT * FROM posthog_pluginstorage WHERE "plugin_config_id"=$1 AND "key"=$2 LIMIT 1',
+            [pluginConfig.id, key]
+        )
+        return result?.rows.length === 1 ? JSON.parse(result.rows[0].value) : defaultValue
+    }
+    const set = async function (key: string, value: unknown): Promise<void> {
+        if (typeof value === 'undefined') {
+            await server.db.query('DELETE FROM posthog_pluginstorage WHERE "plugin_config_id"=$1 AND "key"=$2', [
+                pluginConfig.id,
+                key,
+            ])
+        } else {
+            await server.db.query(
+                `
+                    INSERT INTO posthog_pluginstorage ("plugin_config_id", "key", "value") 
+                    VALUES ($1, $2, $3)
+                    ON CONFLICT ("plugin_config_id", "key") 
+                    DO UPDATE SET value = $3
+                `,
+                [pluginConfig.id, key, JSON.stringify(value)]
+            )
+        }
+    }
+
+    return {
+        get,
+        set,
+    }
+}

src/server.ts

@@ -11,7 +11,7 @@ import { PluginEvent } from '@posthog/plugin-scaffold'
 import { defaultConfig } from './config'
 import Piscina from 'piscina'
 import * as Sentry from '@sentry/node'
-import { areWeTestingWithJest, delay } from './utils'
+import { delay } from './utils'
 import { StatsD } from 'hot-shots'
 import { EventsProcessor } from './ingestion/process-event'
 import { status } from './status'
@@ -34,7 +34,7 @@ export async function createServer(
             console.error(error)
         })
         .on('ready', () => {
-            if (!areWeTestingWithJest()) {
+            if (process.env.NODE_ENV !== 'test') {
                 status.info('✅', 'Connected to Redis!')
             }
         })

src/sql.ts

@@ -1,33 +1,20 @@
-import { areWeTestingWithJest } from './utils'
 import { Plugin, PluginAttachmentDB, PluginConfig, PluginConfigId, PluginError, PluginsServer } from './types'
 
-// This nice "mocking" system with areWeTestingWithJest is used since we want to mock data in worker threads.
-// Jest mocks don't penetrate that far. Improvements welcome.
-
 export async function getPluginRows(server: PluginsServer): Promise<Plugin[]> {
-    if (areWeTestingWithJest() && server.__jestMock?.getPluginRows) {
-        return server.__jestMock?.getPluginRows
-    }
     const { rows: pluginRows }: { rows: Plugin[] } = await server.db.query(
         "SELECT * FROM posthog_plugin WHERE id in (SELECT plugin_id FROM posthog_pluginconfig WHERE enabled='t' GROUP BY plugin_id)"
     )
     return pluginRows
 }
 
 export async function getPluginAttachmentRows(server: PluginsServer): Promise<PluginAttachmentDB[]> {
-    if (areWeTestingWithJest() && server.__jestMock?.getPluginAttachmentRows) {
-        return server.__jestMock?.getPluginAttachmentRows
-    }
     const { rows }: { rows: PluginAttachmentDB[] } = await server.db.query(
         "SELECT * FROM posthog_pluginattachment WHERE plugin_config_id in (SELECT id FROM posthog_pluginconfig WHERE enabled='t')"
     )
     return rows
 }
 
 export async function getPluginConfigRows(server: PluginsServer): Promise<PluginConfig[]> {
-    if (areWeTestingWithJest() && server.__jestMock?.getPluginConfigRows) {
-        return server.__jestMock?.getPluginConfigRows
-    }
     const { rows }: { rows: PluginConfig[] } = await server.db.query(
         "SELECT * FROM posthog_pluginconfig WHERE enabled='t'"
     )

src/types.ts

@@ -40,12 +40,6 @@ export interface PluginsServerConfig extends Record<string, any> {
     STATSD_PORT: number
     STATSD_PREFIX: string
     SCHEDULE_LOCK_TTL: number
-
-    __jestMock?: {
-        getPluginRows: Plugin[]
-        getPluginConfigRows: PluginConfig[]
-        getPluginAttachmentRows: PluginAttachmentDB[]
-    }
 }
 
 export interface PluginsServer extends PluginsServerConfig {
@@ -90,6 +84,8 @@ export interface Plugin {
     archive: Buffer | null
     source?: string
     error?: PluginError
+    from_json?: boolean
+    from_web?: boolean
 }
 
 export interface PluginConfig {

src/utils.ts

@@ -127,10 +127,6 @@ export function cloneObject<T extends any | any[]>(obj: T): T {
     return clone as T
 }
 
-export function areWeTestingWithJest(): boolean {
-    return Boolean(process.env.JEST_WORKER_ID)
-}
-
 /** LUT of byte value to hexadecimal representation. For UUID stringification. */
 const byteToHex: string[] = []
 for (let i = 0; i < 256; i++) {

src/vm.ts

@@ -5,7 +5,7 @@ import { createCache } from './extensions/cache'
 import { createPosthog } from './extensions/posthog'
 import { createGoogle } from './extensions/google'
 import { PluginsServer, PluginConfig, PluginConfigVMReponse } from './types'
-import { areWeTestingWithJest } from './utils'
+import { createStorage } from './extensions/storage'
 
 export async function createPluginConfigVM(
     server: PluginsServer,
@@ -25,7 +25,7 @@ export async function createPluginConfigVM(
     vm.freeze(fetch, 'fetch')
     vm.freeze(createGoogle(), 'google')
 
-    if (areWeTestingWithJest()) {
+    if (process.env.NODE_ENV === 'test') {
         vm.freeze(setTimeout, '__jestSetTimeout')
     }
     vm.freeze(
@@ -37,6 +37,7 @@ export async function createPluginConfigVM(
             ),
             config: pluginConfig.config,
             attachments: pluginConfig.attachments,
+            storage: createStorage(server, pluginConfig),
         },
         '__pluginHostMeta'
     )

src/worker/piscina.js

@@ -1,10 +1,6 @@
 const Sentry = require('@sentry/node')
 const { isMainThread, threadId } = require('worker_threads')
 
-function areWeTestingWithJest() {
-    return Boolean(process.env.JEST_WORKER_ID)
-}
-
 if (isMainThread) {
     const Piscina = require('piscina')
     const { createConfig } = require('./config')
@@ -20,7 +16,7 @@ if (isMainThread) {
         },
     }
 } else {
-    if (areWeTestingWithJest()) {
+    if (process.env.NODE_ENV === 'test') {
         require('ts-node').register()
     }
 

tasks/db-init.ts → tasks/init-test-db.ts

@@ -7,23 +7,23 @@ async function task() {
     await server.db.query(createPlugin)
     await server.db.query(createPluginConfig)
     await server.db.query(createPluginAttachment)
+    await server.db.query(createPluginStorage)
 
     await closeServer()
 }
 
-const ifNotExists = 'IF NOT EXISTS'
-
 const createTeam = `
-    CREATE TABLE ${ifNotExists} posthog_team
+    CREATE TABLE posthog_team
     (
         id            serial  NOT NULL
             CONSTRAINT posthog_team_pkey
-                PRIMARY KEY
+                PRIMARY KEY,
+        name          varchar(200)
     );
 `
 
 const createPlugin = `
-    CREATE TABLE ${ifNotExists} posthog_plugin
+    CREATE TABLE posthog_plugin
     (
         id            serial  NOT NULL
             CONSTRAINT posthog_plugin_pkey
@@ -36,12 +36,14 @@ const createPlugin = `
         archive       bytea,
         from_json     boolean NOT NULL,
         from_web      boolean NOT NULL,
-        error         jsonb
+        error         jsonb,
+        plugin_type     varchar(200),
+        source          text
     );
 `
 
 const createPluginAttachment = `
-    CREATE TABLE ${ifNotExists} posthog_pluginattachment
+    CREATE TABLE posthog_pluginattachment
     (
         id               serial       NOT NULL
             CONSTRAINT posthog_pluginattachment_pkey
@@ -61,15 +63,15 @@ const createPluginAttachment = `
                 DEFERRABLE INITIALLY DEFERRED
     );
 
-    CREATE INDEX ${ifNotExists} posthog_pluginattachment_plugin_config_id_cc94a1b9
+    CREATE INDEX posthog_pluginattachment_plugin_config_id_cc94a1b9
         ON posthog_pluginattachment (plugin_config_id);
 
-    CREATE INDEX ${ifNotExists} posthog_pluginattachment_team_id_415eacc7
+    CREATE INDEX posthog_pluginattachment_team_id_415eacc7
         ON posthog_pluginattachment (team_id);
 `
 
 const createPluginConfig = `
-    CREATE TABLE ${ifNotExists} posthog_pluginconfig
+    CREATE TABLE posthog_pluginconfig
     (
         id        serial  NOT NULL
             CONSTRAINT posthog_pluginconfig_pkey
@@ -88,11 +90,25 @@ const createPluginConfig = `
         error     jsonb
     );
 
-    CREATE INDEX ${ifNotExists} posthog_pluginconfig_team_id_71185766
+    CREATE INDEX posthog_pluginconfig_team_id_71185766
         ON posthog_pluginconfig (team_id);
 
-    CREATE INDEX ${ifNotExists} posthog_pluginconfig_plugin_id_d014ca1c
+    CREATE INDEX posthog_pluginconfig_plugin_id_d014ca1c
         ON posthog_pluginconfig (plugin_id);
 `
 
+const createPluginStorage = `
+    CREATE TABLE posthog_pluginstorage (
+        id serial NOT NULL CONSTRAINT posthog_pluginstorage_pkey PRIMARY KEY,
+        key varchar(200) NOT NULL,
+        value text,
+        plugin_config_id integer NOT NULL CONSTRAINT posthog_pluginstorag_plugin_config_id_6744363a_fk_posthog_p
+        REFERENCES posthog_pluginconfig DEFERRABLE INITIALLY DEFERRED
+    );
+
+    CREATE INDEX posthog_pluginstorage_plugin_config_id_6744363a ON posthog_pluginstorage (plugin_config_id);
+
+    CREATE UNIQUE INDEX posthog_unique_plugin_storage_key ON posthog_pluginstorage (plugin_config_id int4_ops, KEY text_ops);
+`
+
 task()