feat(orm)!: Add typed typed mutation hooks #2365

BREAKINGCHANGES.md

@@ -4,3 +4,4 @@
 1. `@omit` and `@password` attributes have been removed
 1. SWR plugin is removed
 1. `makeModelSchema()` no longer includes relation fields by default — use `include` or `select` options to opt in, mirroring ORM behaviour
+1. ORM entity mutation catch-all hooks now need to be registered under `onEntityMutation.$all`

packages/orm/src/client/executor/zenstack-query-executor.ts

@@ -123,7 +123,12 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
     }
 
     private get hasEntityMutationPluginsWithAfterMutationHooks() {
-        return (this.client.$options.plugins ?? []).some((plugin) => plugin.onEntityMutation?.afterEntityMutation);
+        return (this.client.$options.plugins ?? []).some((plugin) => {
+            if (!plugin.onEntityMutation) return false;
+            if (plugin.onEntityMutation['$all']?.afterEntityMutation) return true;
+            const models = Object.keys(this.client.$schema.models);
+            return models.some((model) => (plugin.onEntityMutation as any)?.[model]?.afterEntityMutation);
+        });
     }
 
     private get hasOnKyselyHooks() {
@@ -351,18 +356,34 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
         if (this.options.plugins) {
             for (const plugin of this.options.plugins) {
                 const onEntityMutation = plugin.onEntityMutation;
-                if (!onEntityMutation?.beforeEntityMutation) {
+                if (!onEntityMutation) {
                     continue;
                 }
 
-                await onEntityMutation.beforeEntityMutation({
-                    model: mutationInfo.model,
-                    action: mutationInfo.action,
-                    queryNode,
-                    loadBeforeMutationEntities,
-                    client,
-                    queryId,
-                });
+                // all-model hook
+                if (onEntityMutation['$all']?.beforeEntityMutation) {
+                    await onEntityMutation['$all'].beforeEntityMutation({
+                        model: mutationInfo.model,
+                        action: mutationInfo.action,
+                        queryNode,
+                        loadBeforeMutationEntities,
+                        client,
+                        queryId,
+                    });
+                }
+
+                // per-model hook
+                const modelHooks = (onEntityMutation as Record<string, any>)[mutationInfo.model];
+                if (modelHooks?.beforeEntityMutation) {
+                    await modelHooks.beforeEntityMutation({
+                        model: mutationInfo.model,
+                        action: mutationInfo.action,
+                        queryNode,
+                        loadBeforeMutationEntities,
+                        client,
+                        queryId,
+                    });
+                }
             }
         }
     }
@@ -373,22 +394,29 @@ export class ZenStackQueryExecutor extends DefaultQueryExecutor {
 
         const hooks: AfterEntityMutationCallback<SchemaDef>[] = [];
 
-        // tsc perf
         for (const plugin of this.options.plugins ?? []) {
             const onEntityMutation = plugin.onEntityMutation;
 
-            if (!onEntityMutation?.afterEntityMutation) {
-                continue;
-            }
-            if (filterFor === 'inTx' && !onEntityMutation.runAfterMutationWithinTransaction) {
+            if (!onEntityMutation) {
                 continue;
             }
 
-            if (filterFor === 'outTx' && onEntityMutation.runAfterMutationWithinTransaction) {
-                continue;
+            // all-model hook
+            if (onEntityMutation['$all']?.afterEntityMutation) {
+                const runInTx = onEntityMutation['$all'].runAfterMutationWithinTransaction ?? false;
+                if (filterFor === 'all' || (filterFor === 'inTx' && runInTx) || (filterFor === 'outTx' && !runInTx)) {
+                    hooks.push(onEntityMutation['$all'].afterEntityMutation.bind(plugin));
+                }
             }
 
-            hooks.push(onEntityMutation.afterEntityMutation.bind(plugin));
+            // per-model hook
+            const modelHooks = (onEntityMutation as Record<string, any>)[mutationInfo.model];
+            if (modelHooks?.afterEntityMutation) {
+                const runInTx = modelHooks.runAfterMutationWithinTransaction ?? false;
+                if (filterFor === 'all' || (filterFor === 'inTx' && runInTx) || (filterFor === 'outTx' && !runInTx)) {
+                    hooks.push(modelHooks.afterEntityMutation.bind(plugin));
+                }
+            }
         }
 
         if (hooks.length === 0) {

packages/orm/src/client/plugin.ts

@@ -5,6 +5,7 @@ import type { GetModelFields, GetModels, NonRelationFields, SchemaDef } from '@z
 import type { MaybePromise } from '../utils/type-utils';
 import type { MapModelFieldType } from './crud-types';
 import type { AllCrudOperations, CoreCrudOperations } from './crud/operations/base';
+import type { DefaultModelResult } from './crud-types';
 
 type AllowedExtQueryArgKeys = CoreCrudOperations | '$create' | '$read' | '$update' | '$delete' | '$all';
 
@@ -287,6 +288,12 @@ type OnQueryHookContext<Schema extends SchemaDef> = {
 // #region OnEntityMutation hooks
 
 export type EntityMutationHooksDef<Schema extends SchemaDef> = {
+    [M in '$all' | GetModels<Schema>]?: M extends '$all'
+        ? AllEntityMutationHooksDef<Schema>
+        : ModelEntityMutationHooksDef<Schema, M & GetModels<Schema>>;
+};
+
+export type AllEntityMutationHooksDef<Schema extends SchemaDef> = {
     /**
      * Called before entities are mutated.
      */
@@ -374,13 +381,38 @@ export type PluginAfterEntityMutationArgs<Schema extends SchemaDef> = MutationHo
 
     /**
      * The ZenStack client you can use to perform additional operations.
-     * See {@link EntityMutationHooksDef.runAfterMutationWithinTransaction} for detailed transaction behavior.
+     * See {@link AllEntityMutationHooksDef.runAfterMutationWithinTransaction} for detailed transaction behavior.
      *
      * Mutations initiated from this client will NOT trigger entity mutation hooks to avoid infinite loops.
      */
     client: ClientContract<Schema>;
 };
 
+export type PluginModelBeforeEntityMutationArgs<Schema extends SchemaDef, Model extends GetModels<Schema>> = Omit<
+    MutationHooksArgs<Schema>,
+    'model'
+> & {
+    model: Model;
+    loadBeforeMutationEntities(): Promise<DefaultModelResult<Schema, Model>[] | undefined>;
+    client: ClientContract<Schema>;
+};
+
+export type PluginModelAfterEntityMutationArgs<Schema extends SchemaDef, Model extends GetModels<Schema>> = Omit<
+    MutationHooksArgs<Schema>,
+    'model'
+> & {
+    model: Model;
+    loadAfterMutationEntities(): Promise<DefaultModelResult<Schema, Model>[] | undefined>;
+    beforeMutationEntities?: DefaultModelResult<Schema, Model>[];
+    client: ClientContract<Schema>;
+};
+
+export type ModelEntityMutationHooksDef<Schema extends SchemaDef, Model extends GetModels<Schema>> = {
+    beforeEntityMutation?: (args: PluginModelBeforeEntityMutationArgs<Schema, Model>) => MaybePromise<void>;
+    afterEntityMutation?: (args: PluginModelAfterEntityMutationArgs<Schema, Model>) => MaybePromise<void>;
+    runAfterMutationWithinTransaction?: boolean;
+};
+
 // #endregion
 
 // #region OnKyselyQuery hooks