feat: add wide table query support

moonrailgun · moonrailgun · commit b8fa8988344b · 2025-08-09T00:08:13.000+08:00
diff --git a/src/server/model/insights/index.ts b/src/server/model/insights/index.ts
@@ -8,8 +8,10 @@ import { insightsSurvey } from './survey.js';
 import { insightsAIGateway } from './aiGateway.js';
 import { compact, omit } from 'lodash-es';
 import { prisma } from '../_client.js';
-import { insightsWarehouse } from './warehouse/longTable.js';
+import { insightsLongTableWarehouse } from './warehouse/longTable.js';
 import { INIT_WORKSPACE_ID } from '../../utils/const.js';
+import { findWarehouseApplication } from './warehouse/utils.js';
+import { insightsWideTableWarehouse } from './warehouse/wideTable.js';
 
 export function queryInsight(
   query: z.infer<typeof insightsQuerySchema>,
@@ -30,7 +32,12 @@ export function queryInsight(
   }
 
   if (insightType === 'warehouse' && query.workspaceId === INIT_WORKSPACE_ID) {
-    return insightsWarehouse(query, context);
+    const application = findWarehouseApplication(query.insightId);
+    if (application?.type === 'wideTable') {
+      return insightsWideTableWarehouse(query, context);
+    }
+
+    return insightsLongTableWarehouse(query, context);
   }
 
   throw new Error('Unknown Insight Type');
diff --git a/src/server/model/insights/warehouse/__snapshots__/wideTable.spec.ts.snap b/src/server/model/insights/warehouse/__snapshots__/wideTable.spec.ts.snap
@@ -0,0 +1,21 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`WarehouseWideTableInsightsSqlBuilder > default > sql 1`] = `
+"select
+        DATE_FORMAT("event_date", '%Y-%m-%d') date , count(1) as "$all_event"
+      from "events" 
+      where "event_date" BETWEEN '2025-08-01' AND '2025-08-02'
+      group by 1
+      order by 1 desc
+      limit 1000"
+`;
+
+exports[`WarehouseWideTableInsightsSqlBuilder > group by > sql 1`] = `
+"select
+        DATE_FORMAT("event_date", '%Y-%m-%d') date , is_public as "%is_public" , count(1) as "$all_event"
+      from "events" 
+      where "event_date" BETWEEN '2025-08-01' AND '2025-08-08'
+      group by 1 , 2
+      order by 1 desc
+      limit 1000"
+`;
diff --git a/src/server/model/insights/warehouse/longTable.ts b/src/server/model/insights/warehouse/longTable.ts
@@ -15,16 +15,16 @@ import {
   getWarehouseApplications,
   getWarehouseConnection,
   MYSQL_DATE_FORMATS,
-  WarehouseInsightsApplication,
+  WarehouseLongTableInsightsApplication,
 } from './utils.js';
 
 export class WarehouseLongTableInsightsSqlBuilder extends InsightsSqlBuilder {
-  getApplication(): WarehouseInsightsApplication {
+  getApplication(): WarehouseLongTableInsightsApplication {
     const name = this.query.insightId;
 
     const application = getWarehouseApplications().find(
-      (app) => app.name === name
-    );
+      (app) => app.name === name && app.type === 'longTable'
+    ) as WarehouseLongTableInsightsApplication;
 
     if (!application) {
       throw new Error(`Application ${name} not found`);
@@ -352,7 +352,8 @@ export class WarehouseLongTableInsightsSqlBuilder extends InsightsSqlBuilder {
   }
 
   async executeQuery(sql: Prisma.Sql): Promise<any[]> {
-    const connection = getWarehouseConnection();
+    const application = this.getApplication();
+    const connection = getWarehouseConnection(application.databaseUrl);
 
     const [rows] = await connection.query(
       sql.sql.replaceAll('"', '`'), // avoid mysql and pg sql syntax error about double quote
@@ -363,7 +364,7 @@ export class WarehouseLongTableInsightsSqlBuilder extends InsightsSqlBuilder {
   }
 }
 
-export async function insightsWarehouse(
+export async function insightsLongTableWarehouse(
   query: z.infer<typeof insightsQuerySchema>,
   context: { timezone: string }
 ) {
@@ -377,18 +378,20 @@ export async function insightsWarehouse(
   return result;
 }
 
-export async function insightsWarehouseEvents(
+export async function insightsLongTableWarehouseEvents(
   applicationId: string
 ): Promise<string[]> {
-  const connection = getWarehouseConnection();
-  const eventTable = getWarehouseApplications().find(
-    (app) => app.name === applicationId
-  )?.eventTable;
+  const application = getWarehouseApplications().find(
+    (app) => app.name === applicationId && app.type === 'longTable'
+  ) as WarehouseLongTableInsightsApplication;
+  const connection = getWarehouseConnection(application.databaseUrl);
 
-  if (!eventTable) {
+  if (!application) {
     throw new Error(`Event table not found for application ${applicationId}`);
   }
 
+  const eventTable = application.eventTable;
+
   const [rows] = await connection.query(`
 SELECT DISTINCT event_name
 FROM (
@@ -405,18 +408,20 @@ LIMIT 100;`);
   return compact(rows.map((row) => get(row, 'event_name', '')));
 }
 
-export async function insightsWarehouseFilterParams(
+export async function insightsLongTableWarehouseFilterParams(
   applicationId: string
 ): Promise<string[]> {
-  const connection = getWarehouseConnection();
-  const eventParametersTable = getWarehouseApplications().find(
-    (app) => app.name === applicationId
-  )?.eventParametersTable;
+  const application = getWarehouseApplications().find(
+    (app) => app.name === applicationId && app.type === 'longTable'
+  ) as WarehouseLongTableInsightsApplication;
+  const connection = getWarehouseConnection(application.databaseUrl);
 
-  if (!eventParametersTable) {
+  if (!application) {
     throw new Error(`Event table not found for application ${applicationId}`);
   }
 
+  const eventParametersTable = application.eventParametersTable;
+
   const [rows] = await connection.query(`
 SELECT DISTINCT param_value
 FROM (
diff --git a/src/server/model/insights/warehouse/utils.ts b/src/server/model/insights/warehouse/utils.ts
@@ -18,9 +18,10 @@ export const dateTypeSchema = z.enum([
   'datetime', // for example: 2025-08-01 00:00:00
 ]);
 
-export const warehouseInsightsApplicationSchema = z.object({
+const warehouseLongTableInsightsApplicationSchema = z.object({
+  databaseUrl: z.string().optional(),
   name: z.string(),
-  type: z.enum(['longTable', 'wideTable']).default('longTable'), // long table or wide table, TODO: implement wide table support
+  type: z.literal('longTable').optional(), // long table, default
   eventTable: z.object({
     name: z.string(),
     eventNameField: z.string(),
@@ -42,10 +43,40 @@ export const warehouseInsightsApplicationSchema = z.object({
   }),
 });
 
+const warehouseWideTableInsightsApplicationSchema = z.object({
+  databaseUrl: z.string().optional(),
+  name: z.string(),
+  type: z.literal('wideTable'), // wide table
+  tableName: z.string(),
+  fields: z.array(
+    z.object({
+      name: z.string(),
+      type: z.string().default('string'),
+    })
+  ),
+  distinctField: z.string(),
+  createdAtField: z.string(),
+  createdAtFieldType: dateTypeSchema.default('timestampMs'),
+  dateBasedCreatedAtField: z.string().optional(), // for improve performance, treat as date type
+});
+
+export const warehouseInsightsApplicationSchema = z.union([
+  warehouseLongTableInsightsApplicationSchema,
+  warehouseWideTableInsightsApplicationSchema,
+]);
+
 export type WarehouseInsightsApplication = z.infer<
   typeof warehouseInsightsApplicationSchema
 >;
 
+export type WarehouseLongTableInsightsApplication = z.infer<
+  typeof warehouseLongTableInsightsApplicationSchema
+>;
+
+export type WarehouseWideTableInsightsApplication = z.infer<
+  typeof warehouseWideTableInsightsApplicationSchema
+>;
+
 let applications: WarehouseInsightsApplication[] | null = null;
 export function getWarehouseApplications() {
   if (!applications) {
@@ -57,14 +88,24 @@ export function getWarehouseApplications() {
   return applications;
 }
 
-let connection: Pool | null = null;
-export function getWarehouseConnection() {
-  if (!env.insights.warehouse.enable || !env.insights.warehouse.url) {
+export function findWarehouseApplication(name: string) {
+  return getWarehouseApplications().find((app) => app.name === name);
+}
+
+const connections = new Map<string, Pool>();
+export function getWarehouseConnection(url = env.insights.warehouse.url) {
+  if (!env.insights.warehouse.enable) {
     throw new Error('Warehouse is not enabled');
   }
 
+  if (!url) {
+    throw new Error('Warehouse url is not set');
+  }
+
+  let connection = connections.get(url);
   if (!connection) {
-    connection = createPool(env.insights.warehouse.url);
+    connection = createPool(url);
+    connections.set(url, connection);
   }
 
   return connection;
diff --git a/src/server/model/insights/warehouse/wideTable.spec.ts b/src/server/model/insights/warehouse/wideTable.spec.ts
@@ -0,0 +1,93 @@
+import { beforeAll, describe, expect, test } from 'vitest';
+import { WarehouseWideTableInsightsSqlBuilder } from './wideTable.js';
+import { unwrapSQL } from '../../../utils/prisma.js';
+import dayjs from 'dayjs';
+import { env } from '../../../utils/env.js';
+
+describe('WarehouseWideTableInsightsSqlBuilder', () => {
+  const insightId = 'wide_table_test';
+  const insightType = 'warehouse';
+
+  beforeAll(() => {
+    env.insights.warehouse.applicationsJson = JSON.stringify([
+      {
+        name: 'wide_table_test',
+        type: 'wideTable',
+        tableName: 'events',
+        fields: [
+          {
+            name: 'event_name',
+            type: 'string',
+          },
+        ],
+        distinctField: 'user_id',
+        createdAtField: 'event_timestamp',
+        dateBasedCreatedAtField: 'event_date',
+      },
+    ]);
+  });
+
+  test('default', () => {
+    const builder = new WarehouseWideTableInsightsSqlBuilder(
+      {
+        insightId,
+        insightType,
+        workspaceId: '',
+        metrics: [
+          {
+            name: '$all_event',
+            math: 'events',
+          },
+        ],
+        filters: [],
+        time: {
+          startAt: dayjs('2025-08-01').valueOf(),
+          endAt: dayjs('2025-08-02').valueOf(),
+          unit: 'day',
+        },
+        groups: [],
+      },
+      {
+        timezone: 'UTC',
+      }
+    );
+
+    const sql = builder.build();
+    expect(unwrapSQL(sql)).toMatchSnapshot('sql');
+  });
+
+  test('group by', () => {
+    const builder = new WarehouseWideTableInsightsSqlBuilder(
+      {
+        workspaceId: '',
+        insightId: 'wide_table_test',
+        insightType: 'warehouse',
+        metrics: [
+          {
+            math: 'events',
+            name: '$all_event',
+          },
+        ],
+        filters: [],
+        groups: [
+          {
+            value: 'is_public',
+            type: 'string',
+          },
+        ],
+        time: {
+          startAt: 1753977600000,
+          endAt: 1754668799999,
+          unit: 'day',
+          timezone: 'Asia/Shanghai',
+        },
+      },
+      {
+        timezone: 'UTC',
+      }
+    );
+
+    const sql = builder.build();
+    expect(unwrapSQL(sql)).toMatchSnapshot('sql');
+  });
+});
diff --git a/src/server/model/insights/warehouse/wideTable.ts b/src/server/model/insights/warehouse/wideTable.ts
diff --git a/src/server/trpc/routers/insights.ts b/src/server/trpc/routers/insights.ts
