feat: count replay events in ClickHouse as we ingest them

plugin-server/src/worker/ingestion/process-event.ts

@@ -281,6 +281,7 @@ export interface SummarizedSessionRecordingEvent {
     console_warn_count: number
     console_error_count: number
     size: number
+    event_count: number
 }
 
 export const createSessionReplayEvent = (
@@ -355,6 +356,7 @@ export const createSessionReplayEvent = (
         console_warn_count: consoleWarnCount,
         console_error_count: consoleErrorCount,
         size: Buffer.byteLength(JSON.stringify(events), 'utf8'),
+        event_count: events.length,
     }
 
     return data

plugin-server/tests/main/process-event.test.ts

@@ -1236,6 +1236,7 @@ const sessionReplayEventTestCases: {
         | 'console_warn_count'
         | 'console_error_count'
         | 'size'
+        | 'event_count'
     >
 }[] = [
     {
@@ -1252,6 +1253,7 @@ const sessionReplayEventTestCases: {
             console_warn_count: 0,
             console_error_count: 0,
             size: 73,
+            event_count: 1,
         },
     },
     {
@@ -1268,6 +1270,7 @@ const sessionReplayEventTestCases: {
             console_warn_count: 0,
             console_error_count: 0,
             size: 73,
+            event_count: 1,
         },
     },
     {
@@ -1324,6 +1327,7 @@ const sessionReplayEventTestCases: {
             console_warn_count: 3,
             console_error_count: 1,
             size: 762,
+            event_count: 7,
         },
     },
     {
@@ -1362,6 +1366,7 @@ const sessionReplayEventTestCases: {
             console_warn_count: 0,
             console_error_count: 0,
             size: 213,
+            event_count: 2,
         },
     },
     {
@@ -1389,6 +1394,7 @@ const sessionReplayEventTestCases: {
             console_warn_count: 0,
             console_error_count: 0,
             size: 433,
+            event_count: 6,
         },
     },
 ]

posthog/clickhouse/migrations/0048_session_replay_events_count.py

@@ -0,0 +1,26 @@
+from posthog.clickhouse.client.migration_tools import run_sql_with_exceptions
+from posthog.models.session_replay_event.migrations_sql import (
+    DROP_SESSION_REPLAY_EVENTS_TABLE_MV_SQL,
+    DROP_KAFKA_SESSION_REPLAY_EVENTS_TABLE_SQL,
+    ADD_EVENT_COUNT_WRITABLE_SESSION_REPLAY_EVENTS_TABLE_SQL,
+    ADD_EVENT_COUNT_DISTRIBUTED_SESSION_REPLAY_EVENTS_TABLE_SQL,
+    ADD_EVENT_COUNT_SESSION_REPLAY_EVENTS_TABLE_SQL,
+)
+from posthog.models.session_replay_event.sql import (
+    SESSION_REPLAY_EVENTS_TABLE_MV_SQL,
+    KAFKA_SESSION_REPLAY_EVENTS_TABLE_SQL,
+)
+
+operations = [
+    # we have to drop materialized view first so that we're no longer pulling from kakfa
+    # then we drop the kafka table
+    run_sql_with_exceptions(DROP_SESSION_REPLAY_EVENTS_TABLE_MV_SQL()),
+    run_sql_with_exceptions(DROP_KAFKA_SESSION_REPLAY_EVENTS_TABLE_SQL()),
+    # now we can alter the target tables
+    run_sql_with_exceptions(ADD_EVENT_COUNT_WRITABLE_SESSION_REPLAY_EVENTS_TABLE_SQL()),
+    run_sql_with_exceptions(ADD_EVENT_COUNT_DISTRIBUTED_SESSION_REPLAY_EVENTS_TABLE_SQL()),
+    run_sql_with_exceptions(ADD_EVENT_COUNT_SESSION_REPLAY_EVENTS_TABLE_SQL()),
+    # and then recreate the materialized views and kafka tables
+    run_sql_with_exceptions(KAFKA_SESSION_REPLAY_EVENTS_TABLE_SQL()),
+    run_sql_with_exceptions(SESSION_REPLAY_EVENTS_TABLE_MV_SQL()),
+]

posthog/clickhouse/test/__snapshots__/test_schema.ambr

@@ -336,7 +336,8 @@
       console_log_count Int64,
       console_warn_count Int64,
       console_error_count Int64,
-      size Int64
+      size Int64,
+      event_count Int64
   ) ENGINE = Kafka('test.kafka.broker:9092', 'clickhouse_session_replay_events_test', 'group1', 'JSONEachRow')
 
   '
@@ -922,7 +923,8 @@
       console_log_count Int64,
       console_warn_count Int64,
       console_error_count Int64,
-      size Int64
+      size Int64,
+      event_count Int64
   ) ENGINE = Kafka('kafka:9092', 'clickhouse_session_replay_events_test', 'group1', 'JSONEachRow')
 
   '
@@ -1344,7 +1346,14 @@
       console_warn_count SimpleAggregateFunction(sum, Int64),
       console_error_count SimpleAggregateFunction(sum, Int64),
       -- this column allows us to estimate the amount of data that is being ingested
-      size SimpleAggregateFunction(sum, Int64)
+      size SimpleAggregateFunction(sum, Int64),
+      -- this allows us to count the number of messages received in a session
+      -- often very useful in incidents or debugging
+      message_count SimpleAggregateFunction(sum, Int64)
+      -- this allows us to count the number of snapshot events received in a session
+      -- often very useful in incidents or debugging
+      -- because we batch events we expect message_count to be lower than event_count
+      event_count SimpleAggregateFunction(sum, Int64)
   ) ENGINE = Distributed('posthog', 'posthog_test', 'sharded_session_replay_events', sipHash64(distinct_id))
 
   '
@@ -1377,7 +1386,10 @@
   sum(console_log_count) as console_log_count,
   sum(console_warn_count) as console_warn_count,
   sum(console_error_count) as console_error_count,
-  sum(size) as size
+  sum(size) as size,
+  -- we can count the number of kafka messages instead of sending it explicitly
+  count(*) as message_count,
+  sum(event_count) as event_count
   FROM posthog_test.kafka_session_replay_events
   group by session_id, team_id
 
@@ -1608,7 +1620,14 @@
       console_warn_count SimpleAggregateFunction(sum, Int64),
       console_error_count SimpleAggregateFunction(sum, Int64),
       -- this column allows us to estimate the amount of data that is being ingested
-      size SimpleAggregateFunction(sum, Int64)
+      size SimpleAggregateFunction(sum, Int64),
+      -- this allows us to count the number of messages received in a session
+      -- often very useful in incidents or debugging
+      message_count SimpleAggregateFunction(sum, Int64)
+      -- this allows us to count the number of snapshot events received in a session
+      -- often very useful in incidents or debugging
+      -- because we batch events we expect message_count to be lower than event_count
+      event_count SimpleAggregateFunction(sum, Int64)
   ) ENGINE = ReplicatedAggregatingMergeTree('/clickhouse/tables/77f1df52-4b43-11e9-910f-b8ca3a9b9f3e_{shard}/posthog.session_replay_events', '{replica}')
 
       PARTITION BY toYYYYMM(min_first_timestamp)
@@ -2226,7 +2245,14 @@
       console_warn_count SimpleAggregateFunction(sum, Int64),
       console_error_count SimpleAggregateFunction(sum, Int64),
       -- this column allows us to estimate the amount of data that is being ingested
-      size SimpleAggregateFunction(sum, Int64)
+      size SimpleAggregateFunction(sum, Int64),
+      -- this allows us to count the number of messages received in a session
+      -- often very useful in incidents or debugging
+      message_count SimpleAggregateFunction(sum, Int64)
+      -- this allows us to count the number of snapshot events received in a session
+      -- often very useful in incidents or debugging
+      -- because we batch events we expect message_count to be lower than event_count
+      event_count SimpleAggregateFunction(sum, Int64)
   ) ENGINE = ReplicatedAggregatingMergeTree('/clickhouse/tables/77f1df52-4b43-11e9-910f-b8ca3a9b9f3e_{shard}/posthog.session_replay_events', '{replica}')
 
       PARTITION BY toYYYYMM(min_first_timestamp)

posthog/hogql/database/schema/session_replay_events.py

@@ -31,6 +31,7 @@
     "console_warn_count": IntegerDatabaseField(name="console_warn_count"),
     "console_error_count": IntegerDatabaseField(name="console_error_count"),
     "size": IntegerDatabaseField(name="size"),
+    "event_count": IntegerDatabaseField(name="event_count"),
     "pdi": LazyJoin(
         from_field="distinct_id",
         join_table=PersonDistinctIdsTable(),
@@ -77,6 +78,7 @@ def select_from_session_replay_events_table(requested_fields: Dict[str, List[str
         "console_error_count": ast.Call(name="sum", args=[ast.Field(chain=[table_name, "console_error_count"])]),
         "distinct_id": ast.Call(name="any", args=[ast.Field(chain=[table_name, "distinct_id"])]),
         "size": ast.Call(name="sum", args=[ast.Field(chain=[table_name, "size"])]),
+        "event_count": ast.Call(name="sum", args=[ast.Field(chain=[table_name, "event_count"])]),
     }
 
     select_fields: List[ast.Expr] = []

posthog/hogql/database/test/__snapshots__/test_database.ambr

@@ -276,6 +276,10 @@
               "key": "size",
               "type": "integer"
           },
+          {
+              "key": "event_count",
+              "type": "integer"
+          },
           {
               "key": "pdi",
               "type": "lazy_table",
@@ -405,6 +409,10 @@
               "key": "size",
               "type": "integer"
           },
+          {
+              "key": "event_count",
+              "type": "integer"
+          },
           {
               "key": "pdi",
               "type": "lazy_table",
@@ -849,6 +857,10 @@
               "key": "size",
               "type": "integer"
           },
+          {
+              "key": "event_count",
+              "type": "integer"
+          },
           {
               "key": "pdi",
               "type": "lazy_table",
@@ -978,6 +990,10 @@
               "key": "size",
               "type": "integer"
           },
+          {
+              "key": "event_count",
+              "type": "integer"
+          },
           {
               "key": "pdi",
               "type": "lazy_table",

posthog/models/session_replay_event/migrations_sql.py

@@ -65,3 +65,27 @@
     table_name=SESSION_REPLAY_EVENTS_DATA_TABLE(),
     cluster=settings.CLICKHOUSE_CLUSTER,
 )
+
+# migration to add size column to the session replay table
+ALTER_SESSION_REPLAY_ADD_EVENT_COUNT_COLUMN = """
+    ALTER TABLE {table_name} on CLUSTER '{cluster}'
+        ADD COLUMN IF NOT EXISTS message_count SimpleAggregateFunction(sum, Int64),
+        ADD COLUMN IF NOT EXISTS event_count SimpleAggregateFunction(sum, Int64)
+"""
+
+ADD_EVENT_COUNT_DISTRIBUTED_SESSION_REPLAY_EVENTS_TABLE_SQL = (
+    lambda: ALTER_SESSION_REPLAY_ADD_EVENT_COUNT_COLUMN.format(
+        table_name="session_replay_events",
+        cluster=settings.CLICKHOUSE_CLUSTER,
+    )
+)
+
+ADD_EVENT_COUNT_WRITABLE_SESSION_REPLAY_EVENTS_TABLE_SQL = lambda: ALTER_SESSION_REPLAY_ADD_EVENT_COUNT_COLUMN.format(
+    table_name="writable_session_replay_events",
+    cluster=settings.CLICKHOUSE_CLUSTER,
+)
+
+ADD_EVENT_COUNT_SESSION_REPLAY_EVENTS_TABLE_SQL = lambda: ALTER_SESSION_REPLAY_ADD_EVENT_COUNT_COLUMN.format(
+    table_name=SESSION_REPLAY_EVENTS_DATA_TABLE(),
+    cluster=settings.CLICKHOUSE_CLUSTER,
+)

posthog/models/session_replay_event/sql.py

@@ -27,7 +27,8 @@
     console_log_count Int64,
     console_warn_count Int64,
     console_error_count Int64,
-    size Int64
+    size Int64,
+    event_count Int64
 ) ENGINE = {engine}
 """
 
@@ -53,7 +54,14 @@
     console_warn_count SimpleAggregateFunction(sum, Int64),
     console_error_count SimpleAggregateFunction(sum, Int64),
     -- this column allows us to estimate the amount of data that is being ingested
-    size SimpleAggregateFunction(sum, Int64)
+    size SimpleAggregateFunction(sum, Int64),
+    -- this allows us to count the number of messages received in a session
+    -- often very useful in incidents or debugging
+    message_count SimpleAggregateFunction(sum, Int64),
+    -- this allows us to count the number of snapshot events received in a session
+    -- often very useful in incidents or debugging
+    -- because we batch events we expect message_count to be lower than event_count
+    event_count SimpleAggregateFunction(sum, Int64)
 ) ENGINE = {engine}
 """
 
@@ -117,7 +125,10 @@
 sum(console_log_count) as console_log_count,
 sum(console_warn_count) as console_warn_count,
 sum(console_error_count) as console_error_count,
-sum(size) as size
+sum(size) as size,
+-- we can count the number of kafka messages instead of sending it explicitly
+count(*) as message_count,
+sum(event_count) as event_count
 FROM {database}.kafka_session_replay_events
 group by session_id, team_id
 """.format(