wip

packages/codegen/src/EntityDbMetadata.ts

@@ -220,6 +220,11 @@ export class EntityDbMetadata {
   updatedAt: PrimitiveField | undefined;
   deletedAt: PrimitiveField | undefined;
   baseClassName: string | undefined;
+  inheritanceType: "sti" | "cti" | undefined;
+  /** This will only be set on the base meta. */
+  stiDiscriminatorField?: string;
+  /** This will only be set on the sub metas. */
+  stiDiscriminatorValue?: number;
   abstract: boolean;
   invalidDeferredFK: boolean;
 
@@ -228,6 +233,7 @@ export class EntityDbMetadata {
 
     if (isSubClassTable(table)) {
       this.baseClassName = tableToEntityName(config, table.columns.get("id").foreignKeys[0].referencedTable);
+      this.inheritanceType = "cti";
     }
 
     this.primaryKey =

packages/codegen/src/config.ts

@@ -9,6 +9,11 @@ import { fail, sortKeys, trueIfResolved } from "./utils";
 
 const jsonFormatter = createFromBuffer(getBuffer());
 
+const stiConfig = z.object({
+  entityName: z.string(),
+  fields: z.array(z.string()),
+});
+
 const fieldConfig = z
   .object({
     derived: z.optional(z.union([z.literal("sync"), z.literal("async")])),
@@ -18,6 +23,7 @@ const fieldConfig = z
     zodSchema: z.optional(z.string()),
     type: z.optional(z.string()),
     serde: z.optional(z.string()),
+    singleTableInheritance: z.optional(z.record(z.string(), stiConfig)),
   })
   .strict();
 

packages/codegen/src/generateEntityCodegenFile.ts

@@ -30,10 +30,10 @@ import {
   OptsOf,
   OrderBy,
   PartialOrNull,
-  ReactiveField,
   PersistedAsyncReference,
   PolymorphicReference,
   ProjectEntity,
+  ReactiveField,
   SSAssert,
   TaggedId,
   ValueFilter,
@@ -371,7 +371,7 @@ export function generateEntityCodegenFile(config: Config, dbMeta: DbMetadata, me
     : "";
 
   // Set up the codegen artifacts to extend from the base type if necessary
-  const baseEntity = meta.baseClassName ? dbMeta.entities.find((e) => e.name === meta.baseClassName)! : undefined;
+  const baseEntity = dbMeta.entities.find((e) => e.name === meta.baseClassName);
   const subEntities = dbMeta.entities.filter((e) => e.baseClassName === meta.name);
   const base = baseEntity?.entity.type ?? code`${BaseEntity}<${EntityManager}, ${idType}>`;
   const maybeBaseFields = baseEntity ? code`extends ${imp(baseEntity.name + "Fields@./entities")}` : "";

packages/codegen/src/generateMetadataFile.ts

@@ -26,12 +26,16 @@ export function generateMetadataFile(config: Config, dbMeta: DbMetadata, meta: E
   Object.values(fields).forEach((code) => code.asOneline());
 
   const maybeBaseType = meta.baseClassName ? `"${meta.baseClassName}"` : undefined;
+  // We want to put inheritanceType: sti/cti onto base classes as well
+  const maybeInheritanceType = meta.inheritanceType ? `inheritanceType: "${meta.inheritanceType}",` : "";
+  const maybeStiColumn = meta.stiDiscriminatorField ? `stiDiscriminatorField: "${meta.stiDiscriminatorField}",` : "";
+  const maybeStiValue = meta.stiDiscriminatorValue ? `stiDiscriminatorValue: ${meta.stiDiscriminatorValue},` : "";
 
   return code`
     export const ${entity.metaName}: ${EntityMetadata}<${entity.name}> = {
       cstr: ${entity.type},
       type: "${entity.name}",
-      baseType: ${maybeBaseType},
+      baseType: ${maybeBaseType}, ${maybeInheritanceType} ${maybeStiColumn} ${maybeStiValue}
       idType: "${config.idType ?? "tagged-string"}",
       idDbType: "${meta.primaryKey.columnType}",
       tagName: "${meta.tagName}",

packages/codegen/src/index.ts

@@ -2,14 +2,14 @@ import { ConnectionConfig, newPgConnectionConfig } from "joist-utils";
 import { Client } from "pg";
 import pgStructure from "pg-structure";
 import { saveFiles } from "ts-poet";
-import { DbMetadata, EntityDbMetadata, failIfOverlappingFieldNames } from "./EntityDbMetadata";
+import { DbMetadata, EntityDbMetadata, failIfOverlappingFieldNames, makeEntity } from "./EntityDbMetadata";
 import { assignTags } from "./assignTags";
 import { maybeRunTransforms } from "./codemods";
 import { Config, loadConfig, warnInvalidConfigEntries, writeConfig } from "./config";
 import { generateFiles } from "./generate";
 import { createFlushFunction } from "./generateFlushFunction";
 import { loadEnumMetadata, loadPgEnumMetadata } from "./loadMetadata";
-import { isEntityTable, isJoinTable, mapSimpleDbTypeToTypescriptType } from "./utils";
+import { fail, isEntityTable, isJoinTable, mapSimpleDbTypeToTypescriptType } from "./utils";
 
 export {
   DbMetadata,
@@ -107,9 +107,71 @@ async function loadSchemaMetadata(config: Config, client: Client): Promise<DbMet
     .map((table) => new EntityDbMetadata(config, table, enums));
   const totalTables = db.tables.length;
   const joinTables = db.tables.filter((t) => isJoinTable(config, t)).map((t) => t.name);
+  expandSingleTableInheritance(config, entities);
   return { entities, enums, pgEnums, totalTables, joinTables };
 }
 
+// Expand STI tables into multiple entities
+function expandSingleTableInheritance(config: Config, entities: EntityDbMetadata[]): void {
+  for (const entity of entities) {
+    const [fieldName, stiField] =
+      Object.entries(config.entities[entity.name]?.fields || {}).find(([, f]) => !!f.singleTableInheritance) ?? [];
+    if (fieldName && stiField && stiField.singleTableInheritance) {
+      entity.inheritanceType = "sti";
+      // Ensure we have an enum field so that we can bake the STI discriminators into the metadata.ts file
+      const enumField =
+        entity.enums.find((e) => e.fieldName === fieldName) ??
+        fail(`No enum column found for ${entity.name}.${fieldName}, which is required to use singleTableInheritance`);
+      entity.stiDiscriminatorField = enumField.fieldName;
+      for (const [enumCode, config] of Object.entries(stiField.singleTableInheritance)) {
+        // Synthesize an entity for this STI sub-entity
+        const subEntity: EntityDbMetadata = {
+          name: config.entityName,
+          entity: makeEntity(config.entityName),
+          tableName: entity.tableName,
+          primaryKey: entity.primaryKey,
+          primitives: entity.primitives.filter((f) => config.fields.includes(f.fieldName)),
+          enums: entity.enums.filter((f) => config.fields.includes(f.fieldName)),
+          pgEnums: entity.pgEnums.filter((f) => config.fields.includes(f.fieldName)),
+          manyToOnes: entity.manyToOnes.filter((f) => config.fields.includes(f.fieldName)),
+          oneToManys: entity.oneToManys.filter((f) => config.fields.includes(f.fieldName)),
+          largeOneToManys: entity.largeOneToManys.filter((f) => config.fields.includes(f.fieldName)),
+          oneToOnes: entity.oneToOnes.filter((f) => config.fields.includes(f.fieldName)),
+          manyToManys: entity.manyToManys.filter((f) => config.fields.includes(f.fieldName)),
+          largeManyToManys: entity.largeManyToManys.filter((f) => config.fields.includes(f.fieldName)),
+          polymorphics: entity.polymorphics.filter((f) => config.fields.includes(f.fieldName)),
+          tagName: entity.tagName,
+          createdAt: undefined,
+          updatedAt: undefined,
+          deletedAt: undefined,
+          baseClassName: entity.name,
+          inheritanceType: "sti",
+          stiDiscriminatorValue: (
+            enumField.enumRows.find((r) => r.code === enumCode) ??
+            fail(`No enum row found for ${entity.name}.${fieldName}.${enumCode}`)
+          ).id,
+          abstract: false,
+          invalidDeferredFK: false,
+        };
+
+        // Now strip all the subclass fields from the base class
+        entity.primitives = entity.primitives.filter((f) => !config.fields.includes(f.fieldName));
+        entity.enums = entity.enums.filter((f) => !config.fields.includes(f.fieldName));
+        entity.pgEnums = entity.pgEnums.filter((f) => !config.fields.includes(f.fieldName));
+        entity.manyToOnes = entity.manyToOnes.filter((f) => !config.fields.includes(f.fieldName));
+        entity.oneToManys = entity.oneToManys.filter((f) => !config.fields.includes(f.fieldName));
+        entity.largeOneToManys = entity.largeOneToManys.filter((f) => !config.fields.includes(f.fieldName));
+        entity.oneToOnes = entity.oneToOnes.filter((f) => !config.fields.includes(f.fieldName));
+        entity.manyToManys = entity.manyToManys.filter((f) => !config.fields.includes(f.fieldName));
+        entity.largeManyToManys = entity.largeManyToManys.filter((f) => !config.fields.includes(f.fieldName));
+        entity.polymorphics = entity.polymorphics.filter((f) => !config.fields.includes(f.fieldName));
+
+        entities.push(subEntity);
+      }
+    }
+  }
+}
+
 function maybeSetDatabaseUrl(config: Config): void {
   if (!process.env.DATABASE_URL && config.databaseUrl) {
     process.env.DATABASE_URL = config.databaseUrl;

packages/graphql-codegen/src/testUtils.ts

@@ -56,6 +56,7 @@ export function newEntityMetadata(name: string, opts: Partial<EntityDbMetadata>
     createdAt: undefined,
     deletedAt: undefined,
     baseClassName: undefined,
+    inheritanceType: undefined,
     abstract: false,
     invalidDeferredFK: false,
     ...opts,

packages/orm/src/EntityManager.ts

@@ -26,6 +26,7 @@ import {
   DeepPartialOrNull,
   EntityHook,
   EntityMetadata,
+  EnumField,
   ExpressionFilter,
   FilterWithAlias,
   GraphQLFilterWithAlias,
@@ -1028,10 +1029,23 @@ export class EntityManager<C = unknown, Entity extends EntityW = EntityW> {
 
     // Set a default createdAt/updatedAt that we'll keep if this is a new entity, or over-write if we're loaded an existing row
     if (entity.isNewEntity) {
-      const { createdAt, updatedAt } = getBaseMeta(getMetadata(entity)).timestampFields;
+      const baseMeta = getBaseMeta(getMetadata(entity));
+      const { createdAt, updatedAt } = baseMeta.timestampFields;
       const { data } = getOrmField(entity);
       if (createdAt) data[createdAt] = new Date();
       if (updatedAt) data[updatedAt] = new Date();
+      // Set the discriminator for STI
+      if (baseMeta.inheritanceType === "sti") {
+        const typeName = entity.constructor.name;
+        const st = baseMeta.subTypes.find((st) => st.type === typeName);
+        if (st) {
+          const field = baseMeta.fields[baseMeta.stiDiscriminatorField!] as EnumField;
+          const code = (field.enumDetailType.findById(st.stiDiscriminatorValue!) as any).code;
+          (entity as any)[baseMeta.stiDiscriminatorField!] = code;
+        } else {
+          (entity as any)[baseMeta.stiDiscriminatorField!] = undefined;
+        }
+      }
       this.#rm.queueAllDownstreamFields(entity);
     }
   }
@@ -1327,9 +1341,7 @@ export class EntityManager<C = unknown, Entity extends EntityW = EntityW> {
       let entity = this.findExistingInstance(taggedId) as T;
       if (!entity) {
         // Look for __class from the driver telling us which subtype to instantiate
-        const meta = row.__class
-          ? maybeBaseMeta.subTypes.find((st) => st.type === row.__class) ?? maybeBaseMeta
-          : maybeBaseMeta;
+        const meta = findConcreteMeta(maybeBaseMeta, row);
         // Pass id as a hint that we're in hydrate mode
         entity = new (asConcreteCstr(meta.cstr))(this, taggedId) as T;
         getOrmField(entity).row = row;
@@ -1839,3 +1851,23 @@ function getNow(): Date {
   lastNow = now;
   return now;
 }
+
+function findConcreteMeta(maybeBaseMeta: EntityMetadata, row: any): EntityMetadata {
+  if (!row.__class && maybeBaseMeta.inheritanceType !== "sti") {
+    return maybeBaseMeta;
+  }
+  if (row.__class) {
+    // Look for the CTI __class from the driver telling us which subtype to instantiate
+    return maybeBaseMeta.subTypes.find((st) => st.type === row.__class) ?? maybeBaseMeta;
+  } else if (maybeBaseMeta.inheritanceType === "sti") {
+    // Look for the STI discriminator value
+    const baseMeta = getBaseMeta(maybeBaseMeta);
+    const field = baseMeta.fields[baseMeta.stiDiscriminatorField!];
+    if (field.kind !== "enum") throw new Error("Discriminator field must be an enum");
+    const columnName = field.serde.columns[0].columnName;
+    const value = row[columnName];
+    return baseMeta.subTypes.find((st) => st.stiDiscriminatorValue === value) ?? baseMeta;
+  } else {
+    throw new Error("Unknown inheritance type");
+  }
+}

packages/orm/src/EntityMetadata.ts

@@ -20,6 +20,9 @@ export function getMetadata<T extends Entity>(
 /**
  * Runtime metadata about an entity.
  *
+ * The `joist-codegen` step will generate this by reading the database schema at build/codegen
+ * time, along with any customizations in `joist-config.json`.
+ *
  * Note: This has no generic, like `T extends Entity`, because `Entity<IdType>` i.e. with
  * an unknown string/id type, causes issues when we want to generically mix `EntityMetadata`
  * of different types, that even liberally using `EntityMetadata<any>` did not avoid.
@@ -34,6 +37,11 @@ export interface EntityMetadata<T extends Entity = any> {
   tableName: string;
   /** If we're a subtype, our immediate base type's name, e.g. for `SmallPublisher` this would be `Publisher`. */
   baseType: string | undefined;
+  inheritanceType?: "sti" | "cti" | undefined;
+  /** Indicates the field to use to derive which subtype to instantiate; only set on the base meta. */
+  stiDiscriminatorField?: string;
+  /** The discriminator enum value for this subtype; only set on sub metas. */
+  stiDiscriminatorValue?: number;
   tagName: string;
   fields: Record<string, Field>;
   allFields: Record<string, Field & { aliasSuffix: string }>;
@@ -88,7 +96,7 @@ export type EnumField = {
   fieldName: string;
   fieldIdName: undefined;
   required: boolean;
-  enumDetailType: { getValues(): ReadonlyArray<unknown> };
+  enumDetailType: { getValues(): ReadonlyArray<unknown>; findById(id: number): unknown };
   serde: FieldSerde;
   immutable: boolean;
 };

packages/orm/src/QueryParser.ts

@@ -963,7 +963,11 @@ export function joinClauses(joins: ParsedTable[]): string[] {
 }
 
 function needsClassPerTableJoins(meta: EntityMetadata): boolean {
-  return meta.subTypes.length > 0 || meta.baseTypes.length > 0;
+  return meta.inheritanceType === "cti" && (meta.subTypes.length > 0 || meta.baseTypes.length > 0);
+}
+
+function needsStiDiscriminator(meta: EntityMetadata): boolean {
+  return meta.inheritanceType === "sti";
 }
 
 /** Converts a search term like `foo bar` into a SQL `like` pattern like `%foo%bar%`. */

packages/orm/src/drivers/EntityWriter.ts

@@ -2,6 +2,7 @@ import { getOrmField } from "../BaseEntity";
 import { Entity } from "../Entity";
 import {
   EntityMetadata,
+  Field,
   PrimitiveField,
   getBaseAndSelfMetas,
   getBaseSelfAndSubMetas,
@@ -43,10 +44,18 @@ export function generateOps(todos: Record<string, Todo>): Ops {
 
 function addInserts(ops: Ops, todo: Todo): void {
   if (todo.inserts.length > 0) {
+    // If we have subtypes, this todo.metadata will always be the base type
     const meta = todo.metadata;
     if (meta.subTypes.length > 0) {
-      for (const [meta, group] of groupEntitiesByTable(todo.inserts)) {
-        ops.inserts.push(newInsertOp(meta, group));
+      if (meta.inheritanceType === "cti") {
+        // Insert into each of the CTI tables
+        for (const [meta, group] of groupEntitiesByTable(todo.inserts)) {
+          ops.inserts.push(newInsertOp(meta, group));
+        }
+      } else if (meta.inheritanceType === "sti") {
+        ops.inserts.push(newStiInsertOp(meta, todo.inserts));
+      } else {
+        throw new Error(`Found ${meta.tableName} subTypes without a known inheritanceType ${meta.inheritanceType}`);
       }
     } else {
       ops.inserts.push(newInsertOp(meta, todo.inserts));
@@ -62,6 +71,26 @@ function newInsertOp(meta: EntityMetadata, entities: Entity[]): InsertOp {
   return { tableName: meta.tableName, columns, rows };
 }
 
+function newStiInsertOp(root: EntityMetadata, entities: Entity[]): InsertOp {
+  // Get the unique set of subtypes
+  const subTypes = new Set<EntityMetadata>();
+  for (const e of entities) subTypes.add(getMetadata(e));
+  // All the root fields (including id)
+  const fields: Field[] = Object.values(root.fields);
+  // Then the subtype fields that haven't been seen yet (subtypes have the root fields + can share non-root fields)
+  for (const st of subTypes) {
+    for (const f of Object.values(st.fields)) {
+      if (!fields.some((f2) => f2.fieldName === f.fieldName)) {
+        fields.push(f);
+      }
+    }
+  }
+  const columns = fields.filter(hasSerde).flatMap((f) => f.serde.columns);
+  // And then collect the same bindings across each STI
+  const rows = collectBindings(entities, columns);
+  return { tableName: root.tableName, columns, rows };
+}
+
 function addUpdates(ops: Ops, todo: Todo): void {
   if (todo.updates.length > 0) {
     const meta = todo.metadata;

packages/orm/src/index.ts

@@ -124,6 +124,10 @@ export function setOpts<T extends Entity>(
       if (partial && _value === undefined) {
         return;
       }
+      // Ignore the STI discriminator, em.register will set this accordingly
+      if (meta.inheritanceType === "sti" && getBaseMeta(meta).stiDiscriminatorField === key) {
+        return;
+      }
       // We let optional opts fields be `| null` for convenience, and convert to undefined.
       const value = _value === null ? undefined : _value;
       const current = (entity as any)[key];

packages/tests/integration/graphql-codegen-joist.js

@@ -25,6 +25,10 @@ const mappers = {
   PublisherTypeDetail: "src/entities#PublisherType",
   SmallPublisher: "src/entities#SmallPublisher",
   Tag: "src/entities#Tag",
+  Task: "src/entities#Task",
+  TaskNew: "src/entities#TaskNew",
+  TaskOld: "src/entities#TaskOld",
+  TaskTypeDetail: "src/entities#TaskType",
   User: "src/entities#User",
 };
 
@@ -34,6 +38,7 @@ const enumValues = {
   ImageType: "src/entities#ImageType",
   PublisherSize: "src/entities#PublisherSize",
   PublisherType: "src/entities#PublisherType",
+  TaskType: "src/entities#TaskType",
 };
 
 module.exports = { mappers, enumValues };