Add KAFKA_INCOMING_TOPIC and clean up unused code

README.md

@@ -47,6 +47,7 @@ There's a multitude of settings you can use to control the plugin server. Use th
 | CLICKHOUSE_SECURE             | Secure ClickHouse connection                                        | `false`                               |
 | KAFKA_ENABLED                 | use Kafka instead of Celery to ingest events                        | `false`                               |
 | KAFKA_HOSTS                   | comma-delimited Kafka hosts                                         | `null`                                |
+| KAFKA_CONSUMPTION_TOPIC       | Kafka consumption topic override                                    | `null` (automatic)                    |
 | KAFKA_CLIENT_CERT_B64         | Kafka certificate in Base64                                         | `null`                                |
 | KAFKA_CLIENT_CERT_KEY_B64     | Kafka certificate key in Base64                                     | `null`                                |
 | KAFKA_TRUSTED_CERT_B64        | Kafka trusted CA in Base64                                          | `null`                                |

src/config.ts

@@ -20,6 +20,7 @@ export function getDefaultConfig(): PluginsServerConfig {
         KAFKA_CLIENT_CERT_B64: null,
         KAFKA_CLIENT_CERT_KEY_B64: null,
         KAFKA_TRUSTED_CERT_B64: null,
+        KAFKA_CONSUMPTION_TOPIC: null,
         PLUGIN_SERVER_INGESTION: false,
         PLUGINS_CELERY_QUEUE: 'posthog-plugins',
         REDIS_URL: 'redis://127.0.0.1',
@@ -62,6 +63,7 @@ export function getConfigHelp(): Record<keyof PluginsServerConfig, string> {
         LOG_LEVEL: 'minimum log level',
         KAFKA_ENABLED: 'use Kafka instead of Celery to ingest events',
         KAFKA_HOSTS: 'comma-delimited Kafka hosts',
+        KAFKA_CONSUMPTION_TOPIC: 'Kafka consumption topic override',
         KAFKA_CLIENT_CERT_B64: 'Kafka certificate in Base64',
         KAFKA_CLIENT_CERT_KEY_B64: 'Kafka certificate key in Base64',
         KAFKA_TRUSTED_CERT_B64: 'Kafka trusted CA in Base64',

src/extensions/posthog.ts

@@ -1,4 +1,3 @@
-import { KAFKA_EVENTS_INGESTION_HANDOFF } from '../ingestion/topics'
 import { Properties } from '@posthog/plugin-scaffold'
 import { DateTime } from 'luxon'
 import { PluginsServer, PluginConfig, RawEventMessage } from 'types'
@@ -31,7 +30,7 @@ export function createPosthog(server: PluginsServer, pluginConfig: PluginConfig)
                 throw new Error('kafkaProducer not configured!')
             }
             server.kafkaProducer.send({
-                topic: KAFKA_EVENTS_INGESTION_HANDOFF,
+                topic: server.KAFKA_CONSUMPTION_TOPIC!,
                 messages: [
                     {
                         key: data.uuid,

src/ingestion/kafka-queue.ts

@@ -1,7 +1,6 @@
 import * as Sentry from '@sentry/node'
 import { Kafka, Consumer, Message, EachBatchPayload } from 'kafkajs'
 import { PluginsServer, Queue, RawEventMessage } from 'types'
-import { KAFKA_EVENTS_INGESTION_HANDOFF } from './topics'
 import { PluginEvent } from '@posthog/plugin-scaffold'
 import { status } from '../status'
 import { killGracefully } from '../utils'
@@ -85,7 +84,7 @@ export class KafkaQueue implements Queue {
             this.consumer.on(this.consumer.events.CRASH, ({ payload: { error } }) => reject(error))
             status.info('⏬', `Connecting Kafka consumer to ${this.pluginsServer.KAFKA_HOSTS}...`)
             this.wasConsumerRan = true
-            await this.consumer.subscribe({ topic: KAFKA_EVENTS_INGESTION_HANDOFF })
+            await this.consumer.subscribe({ topic: this.pluginsServer.KAFKA_CONSUMPTION_TOPIC! })
             // KafkaJS batching: https://kafka.js.org/docs/consuming#a-name-each-batch-a-eachbatch
             await this.consumer.run({
                 // TODO: eachBatchAutoResolve: false, // don't autoresolve whole batch in case we exit it early
@@ -104,7 +103,7 @@ export class KafkaQueue implements Queue {
             return
         }
         status.info('⏳', 'Pausing Kafka consumer...')
-        await this.consumer.pause([{ topic: KAFKA_EVENTS_INGESTION_HANDOFF }])
+        await this.consumer.pause([{ topic: this.pluginsServer.KAFKA_CONSUMPTION_TOPIC! }])
         status.info('⏸', 'Kafka consumer paused!')
     }
 
@@ -113,12 +112,12 @@ export class KafkaQueue implements Queue {
             return
         }
         status.info('⏳', 'Resuming Kafka consumer...')
-        await this.consumer.resume([{ topic: KAFKA_EVENTS_INGESTION_HANDOFF }])
+        await this.consumer.resume([{ topic: this.pluginsServer.KAFKA_CONSUMPTION_TOPIC! }])
         status.info('▶️', 'Kafka consumer resumed!')
     }
 
     isPaused(): boolean {
-        return this.consumer.paused().some(({ topic }) => topic === KAFKA_EVENTS_INGESTION_HANDOFF)
+        return this.consumer.paused().some(({ topic }) => topic === this.pluginsServer.KAFKA_CONSUMPTION_TOPIC)
     }
 
     async stop(): Promise<void> {

src/ingestion/utils.ts

@@ -1,20 +1,6 @@
-import { DateTime } from 'luxon'
-import { Element, BaseEventMessage, RawEventMessage, EventMessage, BasePerson, RawPerson, Person } from '../types'
+import { Element, BasePerson, RawPerson, Person } from '../types'
 import crypto from 'crypto'
 
-export function parseRawEventMessage(message: RawEventMessage): EventMessage {
-    return {
-        ...(message as BaseEventMessage),
-        data: JSON.parse(message.data),
-        now: DateTime.fromISO(message.now),
-        sent_at: DateTime.fromISO(message.sent_at),
-    }
-}
-
-export function parseRawPerson(rawPerson: RawPerson): Person {
-    return { ...(rawPerson as BasePerson), created_at: DateTime.fromISO(rawPerson.created_at) }
-}
-
 export function unparsePersonPartial(person: Partial<Person>): Partial<RawPerson> {
     return { ...(person as BasePerson), ...(person.created_at ? { created_at: person.created_at.toISO() } : {}) }
 }

src/server.ts

@@ -20,6 +20,7 @@ import { startSchedule } from './services/schedule'
 import { ConnectionOptions } from 'tls'
 import { DB } from './db'
 import { DateTime } from 'luxon'
+import { KAFKA_EVENTS_INGESTION_HANDOFF, KAFKA_EVENTS_WAL } from './ingestion/topics'
 
 export async function createServer(
     config: Partial<PluginsServerConfig> = {},
@@ -80,6 +81,15 @@ export async function createServer(
                 database: serverConfig.CLICKHOUSE_DATABASE,
             },
         })
+        await clickhouse.query('SELECT 1') // test that the connection works
+
+        if (!serverConfig.KAFKA_CONSUMPTION_TOPIC) {
+            // When ingesting events, listen to the "INGESTION_HANDOFF" topic, otherwise listen to the "WAL" and discard
+            serverConfig.KAFKA_CONSUMPTION_TOPIC = serverConfig.PLUGIN_SERVER_INGESTION
+                ? KAFKA_EVENTS_INGESTION_HANDOFF
+                : KAFKA_EVENTS_WAL
+        }
+
         kafka = new Kafka({
             clientId: `plugin-server-v${version}-${new UUIDT()}`,
             brokers: serverConfig.KAFKA_HOSTS.split(','),

src/types.ts

@@ -34,6 +34,7 @@ export interface PluginsServerConfig extends Record<string, any> {
     KAFKA_CLIENT_CERT_B64: string | null
     KAFKA_CLIENT_CERT_KEY_B64: string | null
     KAFKA_TRUSTED_CERT_B64: string | null
+    KAFKA_CONSUMPTION_TOPIC: string | null
     PLUGINS_CELERY_QUEUE: string
     REDIS_URL: string
     BASE_DIR: string

src/worker/queue.ts

@@ -92,16 +92,18 @@ async function startQueueKafka(
             status.error('❓', 'UUID missing in event received from Kafka!')
             return
         }
-        await server.eventsProcessor.processEvent(
-            distinct_id,
-            ip,
-            site_url,
-            event,
-            team_id,
-            DateTime.fromISO(now),
-            sent_at ? DateTime.fromISO(sent_at) : null,
-            uuid
-        )
+        if (server.PLUGIN_SERVER_INGESTION) {
+            await server.eventsProcessor.processEvent(
+                distinct_id,
+                ip,
+                site_url,
+                event,
+                team_id,
+                DateTime.fromISO(now),
+                sent_at ? DateTime.fromISO(sent_at) : null,
+                uuid
+            )
+        }
     })
 
     await kafkaQueue.start()

tests/clickhouse/e2e.test.ts

@@ -9,6 +9,7 @@ import { delay, UUIDT } from '../../src/utils'
 import { resetTestDatabaseClickhouse } from '../helpers/clickhouse'
 import { resetKafka } from '../helpers/kafka'
 import { delayUntilEventIngested } from '../shared/process-event'
+import { KAFKA_EVENTS_INGESTION_HANDOFF } from '../../src/ingestion/topics'
 
 jest.setTimeout(60000) // 60 sec timeout
 
@@ -17,6 +18,7 @@ const extraServerConfig: Partial<PluginsServerConfig> = {
     KAFKA_HOSTS: process.env.KAFKA_HOSTS || 'kafka:9092',
     WORKER_CONCURRENCY: 2,
     PLUGIN_SERVER_INGESTION: true,
+    KAFKA_CONSUMPTION_TOPIC: KAFKA_EVENTS_INGESTION_HANDOFF,
     LOG_LEVEL: LogLevel.Log,
 }
 

tests/clickhouse/process-event.test.ts

@@ -2,12 +2,15 @@ import { PluginsServerConfig, Event } from '../../src/types'
 import { resetTestDatabaseClickhouse } from '../helpers/clickhouse'
 import { resetKafka } from '../helpers/kafka'
 import { createProcessEventTests } from '../shared/process-event'
+import { KAFKA_EVENTS_INGESTION_HANDOFF } from '../../src/ingestion/topics'
 
 jest.setTimeout(180_000) // 3 minute timeout
 
 const extraServerConfig: Partial<PluginsServerConfig> = {
     KAFKA_ENABLED: true,
     KAFKA_HOSTS: process.env.KAFKA_HOSTS || 'kafka:9092',
+    PLUGIN_SERVER_INGESTION: true,
+    KAFKA_CONSUMPTION_TOPIC: KAFKA_EVENTS_INGESTION_HANDOFF,
 }
 
 describe('process event (clickhouse)', () => {

tests/helpers/kafka.ts

@@ -1,116 +1,8 @@
-import { EventEmitter } from 'events'
-import { Kafka, Consumer, logLevel, EachMessagePayload, Producer } from 'kafkajs'
-import {
-    KAFKA_EVENTS,
-    KAFKA_EVENTS_INGESTION_HANDOFF,
-    KAFKA_SESSION_RECORDING_EVENTS,
-} from '../../src/ingestion/topics'
-import { parseRawEventMessage } from '../../src/ingestion/utils'
-import { EventMessage, PluginsServerConfig } from '../../src/types'
+import { Kafka, logLevel } from 'kafkajs'
+import { PluginsServerConfig } from '../../src/types'
 import { delay, UUIDT } from '../../src/utils'
 import { defaultConfig, overrideWithEnv } from '../../src/config'
-
-export class KafkaObserver extends EventEmitter {
-    public kafka: Kafka
-    public producer: Producer
-    public consumer: Consumer
-
-    private isStarted: boolean
-
-    constructor(extraServerConfig: Partial<PluginsServerConfig>) {
-        super()
-        const config = { ...overrideWithEnv(defaultConfig, process.env), ...extraServerConfig }
-        this.kafka = new Kafka({
-            clientId: `plugin-server-test-${new UUIDT()}`,
-            brokers: (config.KAFKA_HOSTS || '').split(','),
-            logLevel: logLevel.WARN,
-        })
-        this.producer = this.kafka.producer()
-        this.consumer = this.kafka.consumer({
-            groupId: 'clickhouse-ingestion-test',
-        })
-        this.isStarted = false
-    }
-
-    public async start(): Promise<void> {
-        console.info('observer started!')
-        if (this.isStarted) {
-            return
-        }
-        this.isStarted = true
-        return await new Promise<void>(async (resolve, reject) => {
-            console.info('connecting producer')
-            await this.producer.connect()
-            console.info('subscribing consumer')
-            await this.consumer.subscribe({ topic: KAFKA_EVENTS })
-            console.info('running consumer')
-            await this.consumer.run({
-                eachMessage: async (payload) => {
-                    console.info('message received!')
-                    this.emit('message', payload)
-                },
-            })
-            console.info('setting group join and crash listeners')
-            const { CONNECT, GROUP_JOIN, CRASH } = this.consumer.events
-            this.consumer.on(CONNECT, () => {
-                console.log('consumer connected to kafka')
-            })
-            this.consumer.on(GROUP_JOIN, () => {
-                console.log('joined group')
-                resolve()
-            })
-            this.consumer.on(CRASH, ({ payload: { error } }) => reject(error))
-        })
-    }
-
-    public async stop(): Promise<void> {
-        this.removeAllListeners()
-        console.info('disconnecting producer')
-        await this.producer.disconnect()
-        console.info('stopping consumer')
-        await this.consumer.stop()
-        console.info('disconnecting consumer')
-        await this.consumer.disconnect()
-    }
-
-    public async handOffMessage(message: EventMessage): Promise<void> {
-        console.info('producing message')
-        await this.producer.send({
-            topic: KAFKA_EVENTS_INGESTION_HANDOFF,
-            messages: [{ value: Buffer.from(JSON.stringify(message)) }],
-        })
-    }
-}
-
-export class KafkaCollector extends EventEmitter {
-    collection: EventMessage[]
-    kafkaObserver: KafkaObserver
-
-    constructor(kafkaObserver: KafkaObserver) {
-        super()
-        this.collection = []
-        this.kafkaObserver = kafkaObserver
-        kafkaObserver.addListener('message', (payload: EachMessagePayload) => {
-            console.info('message received')
-            this.collection.push(parseRawEventMessage(JSON.parse(payload.message.value!.toString())))
-            this.emit('message')
-        })
-    }
-
-    async collect(numberOfMessages: number): Promise<EventMessage[]> {
-        return await new Promise((resolve) => {
-            const resolveIfCollectedEnough = () => {
-                console.log('collection:', this.collection)
-                if (this.collection.length >= numberOfMessages) {
-                    this.removeListener('message', resolveIfCollectedEnough)
-                    resolve(this.collection)
-                }
-            }
-            this.addListener('message', resolveIfCollectedEnough)
-            resolveIfCollectedEnough()
-        })
-    }
-}
+import { KAFKA_EVENTS_INGESTION_HANDOFF, KAFKA_SESSION_RECORDING_EVENTS } from '../../src/ingestion/topics'
 
 /** Clear the kafka queue */
 export async function resetKafka(extraServerConfig: Partial<PluginsServerConfig>, delayMs = 2000) {