refactor: unify source types and remove ambiguity

.github/workflows/claude-code-review.yml

@@ -17,14 +17,14 @@ jobs:
     #   github.event.pull_request.user.login == 'external-contributor' ||
     #   github.event.pull_request.user.login == 'new-developer' ||
     #   github.event.pull_request.author_association == 'FIRST_TIME_CONTRIBUTOR'
-    
+
     runs-on: ubuntu-latest
     permissions:
       contents: read
       pull-requests: read
       issues: read
       id-token: write
-    
+
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -46,12 +46,14 @@ jobs:
             - Performance considerations
             - Security concerns
             - Test coverage
-            
+
             Use the repository's CLAUDE.md for guidance on style and conventions. Be constructive and helpful in your feedback.
 
             Use `gh pr comment` with your Bash tool to leave your review as a comment on the PR.
-          
+
           # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
           # or https://docs.claude.com/en/docs/claude-code/sdk#command-line for available options
-          claude_args: '--allowed-tools "Bash(gh issue view:*),Bash(gh search:*),Bash(gh issue list:*),Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr view:*),Bash(gh pr list:*)"'
-
+          claude_args:
+            '--allowed-tools "Bash(gh issue view:*),Bash(gh search:*),Bash(gh
+            issue list:*),Bash(gh pr comment:*),Bash(gh pr diff:*),Bash(gh pr
+            view:*),Bash(gh pr list:*)"'

.github/workflows/claude.yml

@@ -35,7 +35,7 @@ jobs:
         uses: anthropics/claude-code-action@v1
         with:
           anthropic_api_key: ${{ secrets.ANTHROPIC_API_KEY }}
-          
+
           # This is an optional setting that allows Claude to read CI results on PRs
           additional_permissions: |
             actions: read
@@ -47,4 +47,3 @@ jobs:
           # See https://github.com/anthropics/claude-code-action/blob/main/docs/usage.md
           # or https://docs.claude.com/en/docs/claude-code/sdk#command-line for available options
           # claude_args: '--model claude-opus-4-1-20250805 --allowed-tools Bash(gh pr:*)'
-

packages/api/src/clickhouse/__tests__/renderChartConfig.test.ts

@@ -26,7 +26,7 @@ import {
   getServer,
 } from '@/fixtures';
 import Connection from '@/models/connection';
-import { Source } from '@/models/source';
+import { LogSource, MetricSource } from '@/models/source';
 
 const TEST_METRIC_TABLES = {
   sum: DEFAULT_METRICS_TABLE.SUM,
@@ -77,7 +77,7 @@ describe('renderChartConfig', () => {
       username: config.CLICKHOUSE_USER,
       password: config.CLICKHOUSE_PASSWORD,
     });
-    logSource = await Source.create({
+    logSource = await LogSource.create({
       kind: 'log',
       team: team._id,
       from: {
@@ -88,7 +88,7 @@ describe('renderChartConfig', () => {
       connection: connection.id,
       name: 'Logs',
     });
-    metricSource = await Source.create({
+    metricSource = await MetricSource.create({
       kind: 'metric',
       team: team._id,
       from: {

packages/api/src/controllers/sources.ts

@@ -1,4 +1,12 @@
-import { ISource, Source } from '@/models/source';
+import { SourceKind } from '@/../../common-utils/dist/types';
+import {
+  ISource,
+  LogSource,
+  MetricSource,
+  SessionSource,
+  Source,
+  TraceSource,
+} from '@/models/source';
 
 export function getSources(team: string) {
   return Source.find({ team });
@@ -9,17 +17,45 @@ export function getSource(team: string, sourceId: string) {
 }
 
 export function createSource(team: string, source: Omit<ISource, 'id'>) {
-  return Source.create({ ...source, team });
+  switch (source.kind) {
+    case SourceKind.Log:
+      return LogSource.create({ ...source, team });
+    case SourceKind.Trace:
+      return TraceSource.create({ ...source, team });
+    case SourceKind.Metric:
+      return MetricSource.create({ ...source, team });
+    case SourceKind.Session:
+      return SessionSource.create({ ...source, team });
+    default:
+      throw new Error(`Source with kind ${source.kind} is invalid`);
+  }
 }
 
 export function updateSource(
   team: string,
   sourceId: string,
   source: Omit<ISource, 'id'>,
 ) {
-  return Source.findOneAndUpdate({ _id: sourceId, team }, source, {
-    new: true,
-  });
+  switch (source.kind) {
+    case SourceKind.Log:
+      return LogSource.findOneAndUpdate({ _id: sourceId, team }, source, {
+        new: true,
+      });
+    case SourceKind.Trace:
+      return TraceSource.findOneAndUpdate({ _id: sourceId, team }, source, {
+        new: true,
+      });
+    case SourceKind.Metric:
+      return MetricSource.findOneAndUpdate({ _id: sourceId, team }, source, {
+        new: true,
+      });
+    case SourceKind.Session:
+      return SessionSource.findOneAndUpdate({ _id: sourceId, team }, source, {
+        new: true,
+      });
+    default:
+      throw new Error(`Source with kind ${source.kind} is invalid`);
+  }
 }
 
 export function deleteSource(team: string, sourceId: string) {

packages/api/src/models/source.ts

@@ -1,86 +1,133 @@
 import {
+  LogSourceSchema,
   MetricsDataType,
+  MetricSourceSchema,
+  SessionSourceSchema,
   SourceKind,
+  TraceSourceSchema,
   TSource,
 } from '@hyperdx/common-utils/dist/types';
 import mongoose, { Schema } from 'mongoose';
+import { z } from 'zod';
 
-type ObjectId = mongoose.Types.ObjectId;
-
-export interface ISource extends Omit<TSource, 'connection'> {
-  team: ObjectId;
-  connection: ObjectId | string;
-}
+import { objectIdSchema } from '@/utils/zod';
 
+const sourceExtension = {
+  team: objectIdSchema.or(z.instanceof(mongoose.Types.ObjectId)),
+  connection: objectIdSchema.or(z.instanceof(mongoose.Types.ObjectId)),
+};
+const SourceModelSchema = z.discriminatedUnion('kind', [
+  LogSourceSchema.extend(sourceExtension),
+  TraceSourceSchema.extend(sourceExtension),
+  SessionSourceSchema.extend(sourceExtension),
+  MetricSourceSchema.extend(sourceExtension),
+]);
+export type ISource = z.infer<typeof SourceModelSchema>;
 export type SourceDocument = mongoose.HydratedDocument<ISource>;
 
 export const Source = mongoose.model<ISource>(
   'Source',
-  new Schema<ISource>(
-    {
-      kind: {
-        type: String,
-        enum: Object.values(SourceKind),
-        required: true,
-      },
-      team: {
-        type: mongoose.Schema.Types.ObjectId,
-        required: true,
-        ref: 'Team',
-      },
-      from: {
-        databaseName: String,
-        tableName: String,
-      },
-      timestampValueExpression: String,
-      connection: {
-        type: mongoose.Schema.Types.ObjectId,
-        required: true,
-        ref: 'Connection',
-      },
-
-      name: String,
-      displayedTimestampValueExpression: String,
-      implicitColumnExpression: String,
-      serviceNameExpression: String,
-      bodyExpression: String,
-      tableFilterExpression: String,
-      eventAttributesExpression: String,
-      resourceAttributesExpression: String,
-      defaultTableSelectExpression: String,
-      uniqueRowIdExpression: String,
-      severityTextExpression: String,
-      traceIdExpression: String,
-      spanIdExpression: String,
-      traceSourceId: String,
-      sessionSourceId: String,
-      metricSourceId: String,
+  new Schema<ISource>({
+    name: String,
+    kind: {
+      type: String,
+      enum: Object.values(SourceKind),
+      required: true,
+    },
+    from: {
+      databaseName: String,
+      tableName: String,
+    },
+    team: {
+      type: mongoose.Schema.Types.ObjectId,
+      required: true,
+      ref: 'Team',
+    },
+    connection: {
+      type: mongoose.Schema.Types.ObjectId,
+      required: true,
+      ref: 'Connection',
+    },
+  }),
+);
 
-      durationExpression: String,
-      durationPrecision: Number,
-      parentSpanIdExpression: String,
-      spanNameExpression: String,
+export const LogSource = Source.discriminator<
+  Extract<TSource, { kind: SourceKind.Log }>
+>(
+  SourceKind.Log,
+  new mongoose.Schema<Extract<TSource, { kind: SourceKind.Log }>>({
+    timestampValueExpression: String,
+    defaultTableSelectExpression: String,
+    serviceNameExpression: String,
+    severityTextExpression: String,
+    bodyExpression: String,
+    eventAttributesExpression: String,
+    resourceAttributesExpression: String,
+    displayedTimestampValueExpression: String,
+    metricSourceId: String,
+    traceSourceId: String,
+    traceIdExpression: String,
+    spanIdExpression: String,
+    implicitColumnExpression: String,
+    uniqueRowIdExpression: String,
+    tableFilterExpression: String,
+  }),
+);
 
-      logSourceId: String,
-      spanKindExpression: String,
-      statusCodeExpression: String,
-      statusMessageExpression: String,
-      spanEventsValueExpression: String,
+export const TraceSource = Source.discriminator<
+  Extract<TSource, { kind: SourceKind.Trace }>
+>(
+  SourceKind.Trace,
+  new mongoose.Schema<Extract<TSource, { kind: SourceKind.Trace }>>({
+    defaultTableSelectExpression: String,
+    timestampValueExpression: String,
+    durationExpression: String,
+    durationPrecision: Number,
+    traceIdExpression: String,
+    spanIdExpression: String,
+    parentSpanIdExpression: String,
+    spanNameExpression: String,
+    spanKindExpression: String,
+    logSourceId: String,
+    sessionSourceId: String,
+    metricSourceId: String,
+    statusCodeExpression: String,
+    statusMessageExpression: String,
+    serviceNameExpression: String,
+    resourceAttributesExpression: String,
+    eventAttributesExpression: String,
+    spanEventsValueExpression: String,
+    implicitColumnExpression: String,
+  }),
+);
 
-      metricTables: {
-        type: {
-          [MetricsDataType.Gauge]: String,
-          [MetricsDataType.Histogram]: String,
-          [MetricsDataType.Sum]: String,
-          [MetricsDataType.Summary]: String,
-          [MetricsDataType.ExponentialHistogram]: String,
-        },
-        default: undefined,
+export const MetricSource = Source.discriminator<
+  Extract<TSource, { kind: SourceKind.Metric }>
+>(
+  SourceKind.Metric,
+  new mongoose.Schema<Extract<TSource, { kind: SourceKind.Metric }>>({
+    metricTables: {
+      type: {
+        [MetricsDataType.Gauge]: String,
+        [MetricsDataType.Histogram]: String,
+        [MetricsDataType.Sum]: String,
+        [MetricsDataType.Summary]: String,
+        [MetricsDataType.ExponentialHistogram]: String,
       },
+      default: undefined,
     },
-    {
-      toJSON: { virtuals: true },
-      timestamps: true,
-    },
-  ),
+    timestampValueExpression: String,
+    resourceAttributesExpression: String,
+    logSourceId: String,
+  }),
+);
+
+export const SessionSource = Source.discriminator<
+  Extract<TSource, { kind: SourceKind.Session }>
+>(
+  SourceKind.Session,
+  new mongoose.Schema<Extract<TSource, { kind: SourceKind.Session }>>({
+    traceSourceId: String,
+    timestampValueExpression: String,
+  }),
 );