Kafka e2e tests

.github/workflows/ci.yml

@@ -122,29 +122,13 @@ jobs:
                     --health-interval 10s
                     --health-timeout 5s
                     --health-retries 5
-            clickhouse:
-                image: yandex/clickhouse-server
-                ports:
-                    - '8123:8123'
-                    - '9000:9000'
-                    - '9440:9440'
-                    - '9009:9009'
-            zookeeper:
-                image: wurstmeister/zookeeper
-            kafka:
-                image: wurstmeister/kafka
-                ports:
-                    - '9092:9092'
-                env:
-                    KAFKA_ADVERTISED_HOST_NAME: localhost
-                    KAFKA_ZOOKEEPER_CONNECT: zookeeper:2181
 
         env:
             REDIS_URL: 'redis://localhost'
             CLICKHOUSE_HOST: 'localhost'
             CLICKHOUSE_DATABASE: 'posthog_test'
             KAFKA_ENABLED: 'true'
-            KAFKA_HOSTS: 'localhost:9092'
+            KAFKA_HOSTS: 'kafka:9092'
 
         steps:
             - name: Check out Django server for database setup
@@ -158,6 +142,17 @@ jobs:
               with:
                   path: 'plugin-server'
 
+            - name: Fix Kafka Hostname
+              run: |
+                  sudo bash -c 'echo "127.0.0.1 kafka zookeeper" >> /etc/hosts'
+                  ping -c 1 kafka
+                  ping -c 1 zookeeper
+
+            - name: Start Kafka, Clickhouse, Zookeeper
+              run: |
+                  cd posthog/ee
+                  docker-compose -f docker-compose.ch.yml up -d zookeeper kafka clickhouse
+
             - name: Set up Python
               uses: actions/setup-python@v2
               with:

package.json

@@ -12,7 +12,7 @@
         "start": "yarn start:dev",
         "start:dist": "node dist/src/index.js --base-dir ../posthog",
         "start:dev": "ts-node-dev --exit-child src/index.ts --base-dir ../posthog",
-        "start:dev:ee": "DATABASE_URL=postgres://posthog:posthog@localhost:5439/posthog KAFKA_ENABLED=true KAFKA_HOSTS=localhost:9092 yarn start:dev",
+        "start:dev:ee": "KAFKA_ENABLED=true KAFKA_HOSTS=localhost:9092 yarn start:dev",
         "build": "yarn clean && yarn compile",
         "clean": "rimraf dist/*",
         "compile:protobuf": "cd src/idl/ && rimraf protos.* && pbjs -t static-module -w commonjs -o protos.js *.proto && pbts -o protos.d.ts protos.js && eslint --fix . && prettier --write .",
@@ -24,8 +24,10 @@
         "prettier:check": "prettier --check .",
         "prepare": "yarn compile:protobuf",
         "prepublishOnly": "yarn build",
-        "setup:dev": "cd ../posthog && (dropdb test_posthog || echo 'no db to drop') && createdb test_posthog && source env/bin/activate && DATABASE_URL=postgres://localhost:5432/test_posthog DEBUG=1 python manage.py migrate",
-        "setup:dev:ee": "export DEBUG=1 PRIMARY_DB=clickhouse DATABASE_URL=postgres://posthog:posthog@localhost:5439/test_posthog PGPASSWORD=posthog && cd ../posthog && (dropdb -p 5439 -h localhost -U posthog test_posthog || echo 'no db to drop') && createdb -p 5439 -h localhost -U posthog test_posthog && source env/bin/activate && python manage.py migrate && python manage.py migrate_clickhouse"
+        "setup:dev:clickhouse": "cd ../posthog && export DEBUG=1 PRIMARY_DB=clickhouse && source env/bin/activate && python manage.py migrate_clickhouse",
+        "setup:test:ee": "yarn setup:test:postgres && yarn setup:test:clickhouse",
+        "setup:test:postgres": "cd ../posthog && (dropdb test_posthog || echo 'no db to drop') && createdb test_posthog && source env/bin/activate && DATABASE_URL=postgres://localhost:5432/test_posthog DEBUG=1 python manage.py migrate",
+        "setup:test:clickhouse": "cd ../posthog && export TEST=1 PRIMARY_DB=clickhouse CLICKHOUSE_DATABASE=posthog_test && source env/bin/activate && python manage.py migrate_clickhouse"
     },
     "bin": {
         "posthog-plugin-server": "bin/posthog-plugin-server"

src/celery/worker.ts

@@ -171,14 +171,14 @@ export class Worker extends Base implements Queue {
                 throw new Error(`Missing process handler for task ${taskName}`)
             }
 
-            console.info(
+            console.debug(
                 `celery.node Received task: ${taskName}[${taskId}], args: ${args}, kwargs: ${JSON.stringify(kwargs)}`
             )
 
             const timeStart = process.hrtime()
             const taskPromise = handler(...args, kwargs).then((result) => {
                 const diff = process.hrtime(timeStart)
-                console.info(
+                console.debug(
                     `celery.node Task ${taskName}[${taskId}] succeeded in ${diff[0] + diff[1] / 1e9}s: ${result}`
                 )
                 this.activeTasks.delete(taskPromise)

src/config.ts

@@ -10,7 +10,7 @@ export function getDefaultConfig(): PluginsServerConfig {
         CELERY_DEFAULT_QUEUE: 'celery',
         DATABASE_URL: isTestEnv ? 'postgres://localhost:5432/test_posthog' : 'postgres://localhost:5432/posthog',
         CLICKHOUSE_HOST: 'localhost',
-        CLICKHOUSE_DATABASE: 'default',
+        CLICKHOUSE_DATABASE: isTestEnv ? 'posthog_test' : 'default',
         CLICKHOUSE_USERNAME: 'default',
         CLICKHOUSE_PASSWORD: null,
         CLICKHOUSE_CA: null,

src/db.ts

@@ -1,20 +1,24 @@
 import { Properties } from '@posthog/plugin-scaffold'
-import { ClickHouse } from 'clickhouse'
+import { ClickHouse, QueryCursor } from 'clickhouse'
 import { Producer } from 'kafkajs'
 import { DateTime } from 'luxon'
 import { Pool, QueryConfig, QueryResult, QueryResultRow } from 'pg'
+import { string } from 'yargs'
 import { KAFKA_PERSON, KAFKA_PERSON_UNIQUE_ID } from './ingestion/topics'
-import { unparsePersonPartial } from './ingestion/utils'
+import { chainToElements, hashElements, unparsePersonPartial } from './ingestion/utils'
 import {
     Person,
     PersonDistinctId,
     RawPerson,
     RawOrganization,
-    Team,
     PostgresSessionRecordingEvent,
     Event,
+    ClickHouseEvent,
+    Element,
+    SessionRecordingEvent,
+    ElementGroup,
 } from './types'
-import { castTimestampOrNow, sanitizeSqlIdentifier, UUIDT } from './utils'
+import { castTimestampOrNow, clickHouseTimestampToISO, sanitizeSqlIdentifier } from './utils'
 
 /** The recommended way of accessing the database. */
 export class DB {
@@ -31,13 +35,22 @@ export class DB {
         this.clickhouse = clickhouse
     }
 
+    // Direct queries
+
     public async postgresQuery<R extends QueryResultRow = any, I extends any[] = any[]>(
         queryTextOrConfig: string | QueryConfig<I>,
         values?: I
     ): Promise<QueryResult<R>> {
         return this.postgres.query(queryTextOrConfig, values)
     }
 
+    public async clickhouseQuery(query: string, reqParams?: Record<string, any>): Promise<Record<string, any>> {
+        if (!this.clickhouse) {
+            throw new Error('ClickHouse connection has not been provided to this DB instance!')
+        }
+        return this.clickhouse.query(query, reqParams).toPromise()
+    }
+
     // Person
 
     public async fetchPersons(): Promise<Person[]> {
@@ -106,7 +119,7 @@ export class DB {
         const values = [...Object.values(unparsePersonPartial(update)), person.id]
         await this.postgresQuery(
             `UPDATE posthog_person SET ${Object.keys(update).map(
-                (field, index) => sanitizeSqlIdentifier(field) + ' = $' + (index + 1)
+                (field, index) => `"${sanitizeSqlIdentifier(field)}" = $${index + 1}`
             )} WHERE id = $${Object.values(update).length + 1}`,
             values
         )
@@ -130,10 +143,8 @@ export class DB {
         await this.postgresQuery('DELETE FROM posthog_persondistinctid WHERE person_id = $1', [personId])
         await this.postgresQuery('DELETE FROM posthog_person WHERE id = $1', [personId])
         if (this.clickhouse) {
-            await this.clickhouse.query(`ALTER TABLE person DELETE WHERE id = ${personId}`).toPromise()
-            await this.clickhouse
-                .query(`ALTER TABLE person_distinct_id DELETE WHERE person_id = ${personId}`)
-                .toPromise()
+            await this.clickhouseQuery(`ALTER TABLE person DELETE WHERE id = ${personId}`)
+            await this.clickhouseQuery(`ALTER TABLE person_distinct_id DELETE WHERE person_id = ${personId}`)
         }
     }
 
@@ -168,7 +179,7 @@ export class DB {
         const updatedPersonDistinctId: PersonDistinctId = { ...personDistinctId, ...update }
         await this.postgresQuery(
             `UPDATE posthog_persondistinctid SET ${Object.keys(update).map(
-                (field, index) => sanitizeSqlIdentifier(field) + ' = $' + (index + 1)
+                (field, index) => `"${sanitizeSqlIdentifier(field)}" = $${index + 1}`
             )} WHERE id = $${Object.values(update).length + 1}`,
             [...Object.values(update), personDistinctId.id]
         )
@@ -192,21 +203,96 @@ export class DB {
 
     // Event
 
-    public async fetchEvents(): Promise<Event[]> {
-        const result = await this.postgresQuery('SELECT * FROM posthog_event')
-        return result.rows as Event[]
+    public async fetchEvents(): Promise<Event[] | ClickHouseEvent[]> {
+        if (this.kafkaProducer) {
+            const events = (await this.clickhouseQuery(`SELECT * FROM events`)) as ClickHouseEvent[]
+            return (
+                events?.map(
+                    (event) =>
+                        ({
+                            ...event,
+                            ...(typeof event['properties'] === 'string'
+                                ? { properties: JSON.parse(event.properties) }
+                                : {}),
+                            timestamp: clickHouseTimestampToISO(event.timestamp),
+                        } as ClickHouseEvent)
+                ) || []
+            )
+        } else {
+            const result = await this.postgresQuery('SELECT * FROM posthog_event')
+            return result.rows as Event[]
+        }
     }
 
     // SessionRecordingEvent
 
-    public async fetchSessionRecordingEvents(): Promise<PostgresSessionRecordingEvent[]> {
-        const result = await this.postgresQuery('SELECT * FROM posthog_sessionrecordingevent')
-        return result.rows as PostgresSessionRecordingEvent[]
+    public async fetchSessionRecordingEvents(): Promise<PostgresSessionRecordingEvent[] | SessionRecordingEvent[]> {
+        if (this.kafkaProducer) {
+            const events = ((await this.clickhouseQuery(
+                `SELECT * FROM session_recording_events`
+            )) as SessionRecordingEvent[]).map((event) => {
+                return {
+                    ...event,
+                    snapshot_data: event.snapshot_data ? JSON.parse(event.snapshot_data) : null,
+                }
+            })
+            return events
+        } else {
+            const result = await this.postgresQuery('SELECT * FROM posthog_sessionrecordingevent')
+            return result.rows as PostgresSessionRecordingEvent[]
+        }
     }
 
     // Element
 
-    public async fetchElements(): Promise<Element[]> {
-        return (await this.postgresQuery('SELECT * FROM posthog_element')).rows
+    public async fetchElements(event?: Event): Promise<Element[]> {
+        if (this.kafkaProducer) {
+            const events = (await this.clickhouseQuery(
+                `SELECT elements_chain FROM events WHERE uuid='${sanitizeSqlIdentifier((event as any).uuid)}'`
+            )) as ClickHouseEvent[]
+            const chain = events?.[0]?.elements_chain
+            return chainToElements(chain)
+        } else {
+            return (await this.postgresQuery('SELECT * FROM posthog_element')).rows
+        }
+    }
+
+    public async createElementGroup(elements: Element[], teamId: number): Promise<string> {
+        const cleanedElements = elements.map((element, index) => ({ ...element, order: index }))
+        const hash = hashElements(cleanedElements)
+
+        try {
+            const insertResult = await this.postgresQuery(
+                'INSERT INTO posthog_elementgroup (hash, team_id) VALUES ($1, $2) RETURNING *',
+                [hash, teamId]
+            )
+            const elementGroup = insertResult.rows[0] as ElementGroup
+            for (const element of cleanedElements) {
+                await this.postgresQuery(
+                    'INSERT INTO posthog_element (text, tag_name, href, attr_id, nth_child, nth_of_type, attributes, "order", event_id, attr_class, group_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)',
+                    [
+                        element.text,
+                        element.tag_name,
+                        element.href,
+                        element.attr_id,
+                        element.nth_child,
+                        element.nth_of_type,
+                        element.attributes || '{}',
+                        element.order,
+                        element.event_id,
+                        element.attr_class,
+                        elementGroup.id,
+                    ]
+                )
+            }
+        } catch (error) {
+            // Throw further if not postgres error nr "23505" == "unique_violation"
+            // https://www.postgresql.org/docs/12/errcodes-appendix.html
+            if (error.code !== '23505') {
+                throw error
+            }
+        }
+
+        return hash
     }
 }

src/extensions/posthog.ts

@@ -26,10 +26,11 @@ export function createPosthog(server: PluginsServer, pluginConfig: PluginConfig)
 
     if (server.KAFKA_ENABLED) {
         // Sending event to our Kafka>ClickHouse pipeline
-        const producer = server.kafka!.producer()
         sendEvent = async (data) => {
-            await producer.connect()
-            producer!.send({
+            if (!server.kafkaProducer) {
+                throw new Error('kafkaProducer not configured!')
+            }
+            server.kafkaProducer.send({
                 topic: KAFKA_EVENTS_INGESTION_HANDOFF,
                 messages: [
                     {

src/ingestion/kafka-queue.ts

@@ -47,11 +47,16 @@ export class KafkaQueue implements Queue {
             ...rawEvent,
             data: JSON.parse(rawEvent.data),
         }))
-        const pluginEvents: PluginEvent[] = parsedEvents.map((parsedEvent) => ({
-            ...parsedEvent,
-            event: parsedEvent.data.event,
-            properties: parsedEvent.data.properties,
-        }))
+        const pluginEvents: PluginEvent[] = rawEvents.map((rawEvent) => {
+            const { data: dataStr, ...restOfRawEvent } = rawEvent
+            const event = { ...restOfRawEvent, ...JSON.parse(dataStr) }
+            return {
+                ...event,
+                kafka_offset: restOfRawEvent.kafka_offset,
+                site_url: event.site_url || null,
+                ip: event.ip || null,
+            }
+        })
         const processedEvents: PluginEvent[] = (
             await this.processEventBatch(pluginEvents)
         ).filter((event: PluginEvent[] | false | null | undefined) => Boolean(event))