fix: better source validation and refine required source fields

.changeset/six-ways-sell.md

@@ -0,0 +1,7 @@
+---
+"@hyperdx/common-utils": patch
+"@hyperdx/api": patch
+"@hyperdx/app": patch
+---
+
+fix: change sources to discriminated union

packages/api/src/controllers/sources.ts

@@ -1,21 +1,31 @@
-import { SourceKind } from '@hyperdx/common-utils/dist/types';
-
-import { ISource, Source } from '@/models/source';
-
-/**
- * Clean up metricTables property when changing source type away from Metric.
- * This prevents metric-specific configuration from persisting when switching
- * to Log, Trace, or Session sources.
- */
-function cleanSourceData(source: Omit<ISource, 'id'>): Omit<ISource, 'id'> {
-  // Only clean metricTables if the source is not a Metric type
-  if (source.kind !== SourceKind.Metric) {
-    // explicitly setting to null for mongoose to clear column
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
-    source.metricTables = null as any;
-  }
+import { SourceKind, SourceSchema } from '@hyperdx/common-utils/dist/types';
+
+import {
+  ISourceInput,
+  LogSource,
+  MetricSource,
+  SessionSource,
+  Source,
+  TraceSource,
+} from '@/models/source';
 
-  return source;
+// Returns the discriminator model for the given source kind.
+// Updates must go through the correct discriminator model so Mongoose
+// recognises kind-specific fields (e.g. metricTables on MetricSource).
+function getModelForKind(kind: SourceKind) {
+  switch (kind) {
+    case SourceKind.Log:
+      return LogSource;
+    case SourceKind.Trace:
+      return TraceSource;
+    case SourceKind.Session:
+      return SessionSource;
+    case SourceKind.Metric:
+      return MetricSource;
+    default:
+      kind satisfies never;
+      throw new Error(`${kind} is not a valid SourceKind`);
+  }
 }
 
 export function getSources(team: string) {
@@ -26,19 +36,56 @@ export function getSource(team: string, sourceId: string) {
   return Source.findOne({ _id: sourceId, team });
 }
 
-export function createSource(team: string, source: Omit<ISource, 'id'>) {
-  return Source.create({ ...source, team });
+type DistributiveOmit<T, K extends PropertyKey> = T extends T
+  ? Omit<T, K>
+  : never;
+
+export function createSource(
+  team: string,
+  source: DistributiveOmit<ISourceInput, 'id'>,
+) {
+  // @ts-expect-error The create method has incompatible type signatures but is actually safe
+  return getModelForKind(source.kind)?.create({ ...source, team });
 }
 
-export function updateSource(
+export async function updateSource(
   team: string,
   sourceId: string,
-  source: Omit<ISource, 'id'>,
+  source: DistributiveOmit<ISourceInput, 'id'>,
 ) {
-  const cleanedSource = cleanSourceData(source);
-  return Source.findOneAndUpdate({ _id: sourceId, team }, cleanedSource, {
-    new: true,
-  });
+  const existing = await Source.findOne({ _id: sourceId, team });
+  if (!existing) return null;
+
+  // Same kind: simple update through the discriminator model
+  if (existing.kind === source.kind) {
+    // @ts-expect-error The findOneAndUpdate method has incompatible type signatures but is actually safe
+    return getModelForKind(source.kind)?.findOneAndUpdate(
+      { _id: sourceId, team },
+      source,
+      { new: true },
+    );
+  }
+
+  // Kind changed: validate through Zod before writing since the raw
+  // collection bypass skips Mongoose's discriminator validation.
+  const parseResult = SourceSchema.safeParse(source);
+  if (!parseResult.success) {
+    throw new Error(
+      `Invalid source data: ${parseResult.error.errors.map(e => e.message).join(', ')}`,
+    );
+  }
+
+  // Use replaceOne on the raw collection to swap the entire document
+  // in place (including the discriminator key). This is a single atomic
+  // write — the document is never absent from the collection.
+  const replacement = {
+    ...parseResult.data,
+    _id: existing._id,
+    team: existing.team,
+    updatedAt: new Date(),
+  };
+  await Source.collection.replaceOne({ _id: existing._id }, replacement);
+  return getModelForKind(replacement.kind)?.hydrate(replacement);
 }
 
 export function deleteSource(team: string, sourceId: string) {