From 6b4d1ed0ae5e58a7a71ead14327871d696c98096 Mon Sep 17 00:00:00 2001
From: Justin Bellero <jpbellero@gmail.com>
Date: Wed, 24 Apr 2024 17:23:47 -0400
Subject: [PATCH] feat: add discriminator helper (#47)

* add discriminator.ts helper

Allow extension and inheritance of base models with discriminator models.

* Update index.ts

* fix typo in index.ts
---
 src/runtime/server/services/discriminator.ts | 33 ++++++++++++++++++++
 src/runtime/server/services/index.ts         |  1 +
 2 files changed, 34 insertions(+)
 create mode 100644 src/runtime/server/services/discriminator.ts

diff --git a/src/runtime/server/services/discriminator.ts b/src/runtime/server/services/discriminator.ts
new file mode 100644
index 0000000..5a0c287
--- /dev/null
+++ b/src/runtime/server/services/discriminator.ts
@@ -0,0 +1,33 @@
+import type { Model, SchemaDefinition, SchemaOptions } from 'mongoose'
+import mongoose from 'mongoose'
+
+export function defineMongooseDiscriminatorModel<T>(
+  nameOrOptions: string | {
+    name: string;
+    baseModel: Model<T>;
+    schema: SchemaDefinition<T>;
+    options?: SchemaOptions;
+    hooks?: (schema: mongoose.Schema<T>) => void;
+  },
+  baseModel?: Model<T>,
+  schema?: SchemaDefinition<T>,
+  options?: SchemaOptions,
+  hooks?: (schema: mongoose.Schema<T>) => void,
+): Model<T> {
+  let name: string
+  if (typeof nameOrOptions === 'string') { name = nameOrOptions }
+  else {
+    name = nameOrOptions.name
+    baseModel = nameOrOptions.baseModel
+    schema = nameOrOptions.schema
+    options = nameOrOptions.options
+    hooks = nameOrOptions.hooks
+  }
+
+  const newSchema = new mongoose.Schema<T>(schema, options as any)
+
+  if (hooks)
+    hooks(newSchema)
+
+  return baseModel!.discriminator<T>(name, newSchema)
+}
diff --git a/src/runtime/server/services/index.ts b/src/runtime/server/services/index.ts
index 5e1f2d3..ed9cee9 100644
--- a/src/runtime/server/services/index.ts
+++ b/src/runtime/server/services/index.ts
@@ -1,2 +1,3 @@
 export { defineMongooseConnection } from './connection'
 export { defineMongooseModel } from './model'
+export { defineMongooseDiscriminatorModel } from './discriminator'