feat: plugin-extended query args

packages/orm/src/client/client-impl.ts

@@ -13,6 +13,7 @@ import {
 import type { ProcedureDef, SchemaDef } from '../schema';
 import type { AnyKysely } from '../utils/kysely-utils';
 import type { UnwrapTuplePromises } from '../utils/type-utils';
+import type { ClientOptions, ProceduresOptions } from './options';
 import type {
     AuthType,
     ClientConstructor,
@@ -21,13 +22,13 @@ import type {
     TransactionIsolationLevel,
 } from './contract';
 import { AggregateOperationHandler } from './crud/operations/aggregate';
-import type { AllCrudOperation, CoreCrudOperation } from './crud/operations/base';
+import type { AllCrudOperations, CoreCrudOperations } from './crud/operations/base';
 import { BaseOperationHandler } from './crud/operations/base';
 import { CountOperationHandler } from './crud/operations/count';
 import { CreateOperationHandler } from './crud/operations/create';
 import { DeleteOperationHandler } from './crud/operations/delete';
-import { FindOperationHandler } from './crud/operations/find';
 import { ExistsOperationHandler } from './crud/operations/exists';
+import { FindOperationHandler } from './crud/operations/find';
 import { GroupByOperationHandler } from './crud/operations/group-by';
 import { UpdateOperationHandler } from './crud/operations/update';
 import { InputValidator } from './crud/validator';
@@ -36,7 +37,6 @@ import { ZenStackDriver } from './executor/zenstack-driver';
 import { ZenStackQueryExecutor } from './executor/zenstack-query-executor';
 import * as BuiltinFunctions from './functions';
 import { SchemaDbPusher } from './helpers/schema-db-pusher';
-import type { ClientOptions, ProceduresOptions } from './options';
 import type { RuntimePlugin } from './plugin';
 import { createZenStackPromise, type ZenStackPromise } from './promise';
 import { ResultProcessor } from './result-processor';
@@ -292,9 +292,9 @@ export class ClientImpl {
         await new SchemaDbPusher(this.schema, this.kysely).push();
     }
 
-    $use(plugin: RuntimePlugin<SchemaDef>) {
+    $use(plugin: RuntimePlugin<any>) {
         // tsc perf
-        const newPlugins: RuntimePlugin<SchemaDef>[] = [...(this.$options.plugins ?? []), plugin];
+        const newPlugins: RuntimePlugin<any>[] = [...(this.$options.plugins ?? []), plugin];
         const newOptions: ClientOptions<SchemaDef> = {
             ...this.options,
             plugins: newPlugins,
@@ -304,7 +304,7 @@ export class ClientImpl {
 
     $unuse(pluginId: string) {
         // tsc perf
-        const newPlugins: RuntimePlugin<SchemaDef>[] = [];
+        const newPlugins: RuntimePlugin<any>[] = [];
         for (const plugin of this.options.plugins ?? []) {
             if (plugin.id !== pluginId) {
                 newPlugins.push(plugin);
@@ -321,7 +321,7 @@ export class ClientImpl {
         // tsc perf
         const newOptions: ClientOptions<SchemaDef> = {
             ...this.options,
-            plugins: [] as RuntimePlugin<SchemaDef>[],
+            plugins: [] as RuntimePlugin<any>[],
         };
         return new ClientImpl(this.schema, newOptions, this);
     }
@@ -419,12 +419,13 @@ function createModelCrudHandler(
     resultProcessor: ResultProcessor<any>,
 ): ModelOperations<any, any> {
     const createPromise = (
-        operation: CoreCrudOperation,
-        nominalOperation: AllCrudOperation,
+        operation: CoreCrudOperations,
+        nominalOperation: AllCrudOperations,
         args: unknown,
         handler: BaseOperationHandler<any>,
         postProcess = false,
         throwIfNoResult = false,
+        pluginOptions?: unknown,
     ) => {
         return createZenStackPromise(async (txClient?: ClientContract<any>) => {
             let proceed = async (_args: unknown) => {
@@ -448,16 +449,22 @@ function createModelCrudHandler(
                 const onQuery = plugin.onQuery;
                 if (onQuery) {
                     const _proceed = proceed;
-                    proceed = (_args: unknown) =>
-                        onQuery({
+                    proceed = (_args: unknown) => {
+                        const ctx: any = {
                             client,
                             model,
                             operation: nominalOperation,
                             // reflect the latest override if provided
                             args: _args,
                             // ensure inner overrides are propagated to the previous proceed
                             proceed: (nextArgs: unknown) => _proceed(nextArgs),
-                        }) as Promise<unknown>;
+                        };
+                        // Add plugin options to context if provided
+                        if (pluginOptions !== undefined) {
+                            ctx.options = pluginOptions;
+                        }
+                        return (onQuery as (ctx: any) => Promise<unknown>)(ctx);
+                    };
                 }
             }
 
@@ -509,13 +516,16 @@ function createModelCrudHandler(
             );
         },
 
-        findMany: (args: unknown) => {
+        findMany: (...allArgs: unknown[]) => {
+            const [args, pluginOptions] = allArgs;
             return createPromise(
                 'findMany',
                 'findMany',
                 args,
                 new FindOperationHandler<any>(client, model, inputValidator),
                 true,
+                false,
+                pluginOptions,
             );
         },