refactor: unify source types and remove ambiguity

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,132 @@
 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,
+  }),
 );

packages/api/src/routers/api/__tests__/sources.test.ts

@@ -1,10 +1,10 @@
-import { SourceKind, TSourceUnion } from '@hyperdx/common-utils/dist/types';
+import { SourceKind, TSource } from '@hyperdx/common-utils/dist/types';
 import { Types } from 'mongoose';
 
 import { getLoggedInAgent, getServer } from '@/fixtures';
-import { Source } from '@/models/source';
+import { LogSource, Source } from '@/models/source';
 
-const MOCK_SOURCE: Omit<Extract<TSourceUnion, { kind: 'log' }>, 'id'> = {
+const MOCK_SOURCE: Omit<Extract<TSource, { kind: 'log' }>, 'id'> = {
   kind: SourceKind.Log,
   name: 'Test Source',
   connection: new Types.ObjectId().toString(),
@@ -35,7 +35,7 @@ describe('sources router', () => {
     const { agent, team } = await getLoggedInAgent(server);
 
     // Create test source
-    await Source.create({
+    await LogSource.create({
       ...MOCK_SOURCE,
       team: team._id,
     });
@@ -93,7 +93,7 @@ describe('sources router', () => {
     const { agent, team } = await getLoggedInAgent(server);
 
     // Create test source
-    const source = await Source.create({
+    const source = await LogSource.create({
       ...MOCK_SOURCE,
       team: team._id,
     });
@@ -129,7 +129,7 @@ describe('sources router', () => {
     const { agent, team } = await getLoggedInAgent(server);
 
     // Create test source
-    const source = await Source.create({
+    const source = await LogSource.create({
       ...MOCK_SOURCE,
       team: team._id,
     });

packages/api/src/routers/api/sources.ts

@@ -1,6 +1,7 @@
 import {
+  SourceKind,
   SourceSchema,
-  sourceSchemaWithout,
+  SourceSchemaNoId,
 } from '@hyperdx/common-utils/dist/types';
 import express from 'express';
 import { z } from 'zod';
@@ -23,14 +24,28 @@ router.get('/', async (req, res, next) => {
 
     const sources = await getSources(teamId.toString());
 
-    return res.json(sources.map(s => s.toJSON({ getters: true })));
+    const out = sources.map(s => {
+      // Typescript gets confused about calling toJSON on the union type, but
+      // breaking it out into a switch statement keeps TS happy
+      switch (s.kind) {
+        case SourceKind.Log:
+          return s.toJSON({ getters: true });
+        case SourceKind.Trace:
+          return s.toJSON({ getters: true });
+        case SourceKind.Metric:
+          return s.toJSON({ getters: true });
+        case SourceKind.Session:
+          return s.toJSON({ getters: true });
+        default:
+          throw new Error(`Source found with invalid kind ${(s as any).kind}`);
+      }
+    });
+    return res.json(out);
   } catch (e) {
     next(e);
   }
 });
 
-const SourceSchemaNoId = sourceSchemaWithout({ id: true });
-
 router.post(
   '/',
   validateRequest({
@@ -40,11 +55,10 @@ router.post(
     try {
       const { teamId } = getNonNullUserWithTeam(req);
 
-      // TODO: HDX-1768 Eliminate type assertion
       const source = await createSource(teamId.toString(), {
         ...req.body,
         team: teamId,
-      } as any);
+      });
 
       res.json(source);
     } catch (e) {
@@ -65,11 +79,10 @@ router.put(
     try {
       const { teamId } = getNonNullUserWithTeam(req);
 
-      // TODO: HDX-1768 Eliminate type assertion
       const source = await updateSource(teamId.toString(), req.params.id, {
         ...req.body,
         team: teamId,
-      } as any);
+      });
 
       if (!source) {
         res.status(404).send('Source not found');