custom procedures

packages/cli/test/ts-schema-gen.test.ts

@@ -184,6 +184,30 @@ model Post {
         });
     });
 
+    it('generates correct procedures with array params and returns', async () => {
+        const { schema } = await generateTsSchema(`
+model User {
+    id Int @id
+}
+
+procedure findByIds(ids: Int[]): User[]
+procedure getIds(): Int[]
+        `);
+
+        expect(schema.procedures).toMatchObject({
+            findByIds: {
+                params: { ids: { name: 'ids', type: 'Int', array: true } },
+                returnType: 'User',
+                returnArray: true,
+            },
+            getIds: {
+                params: {},
+                returnType: 'Int',
+                returnArray: true,
+            },
+        });
+    });
+
     it('merges fields and attributes from mixins', async () => {
         const { schema } = await generateTsSchema(`
 type Timestamped {

packages/clients/tanstack-query/src/common/types.ts

@@ -1,6 +1,6 @@
 import type { Logger, OptimisticDataProvider } from '@zenstackhq/client-helpers';
 import type { FetchFn } from '@zenstackhq/client-helpers/fetch';
-import type { OperationsIneligibleForDelegateModels } from '@zenstackhq/orm';
+import type { GetProcedureNames, OperationsIneligibleForDelegateModels, ProcedureFunc } from '@zenstackhq/orm';
 import type { GetModels, IsDelegateModel, SchemaDef } from '@zenstackhq/schema';
 
 /**
@@ -76,3 +76,7 @@ type WithOptimisticFlag<T> = T extends object
     : T;
 
 export type WithOptimistic<T> = T extends Array<infer U> ? Array<WithOptimisticFlag<U>> : WithOptimisticFlag<T>;
+
+export type ProcedureReturn<Schema extends SchemaDef, Name extends GetProcedureNames<Schema>> = Awaited<
+    ReturnType<ProcedureFunc<Schema, Name>>
+>;

packages/clients/tanstack-query/src/react.ts

@@ -36,8 +36,11 @@ import type {
     FindFirstArgs,
     FindManyArgs,
     FindUniqueArgs,
+    GetProcedure,
+    GetProcedureNames,
     GroupByArgs,
     GroupByResult,
+    ProcedureEnvelope,
     QueryOptions,
     SelectSubset,
     SimplifiedPlainResult,
@@ -55,12 +58,23 @@ import { getQueryKey } from './common/query-key';
 import type {
     ExtraMutationOptions,
     ExtraQueryOptions,
+    ProcedureReturn,
     QueryContext,
     TrimDelegateModelOperations,
     WithOptimistic,
 } from './common/types';
 export type { FetchFn } from '@zenstackhq/client-helpers/fetch';
 
+type ProcedureHookFn<
+    Schema extends SchemaDef,
+    ProcName extends GetProcedureNames<Schema>,
+    Options,
+    Result,
+    Input = ProcedureEnvelope<Schema, ProcName>,
+> = { args: undefined } extends Input
+    ? (input?: Input, options?: Options) => Result
+    : (input: Input, options?: Options) => Result;
+
 /**
  * React context for query settings.
  */
@@ -133,8 +147,61 @@ export type ModelMutationModelResult<
 
 export type ClientHooks<Schema extends SchemaDef, Options extends QueryOptions<Schema> = QueryOptions<Schema>> = {
     [Model in GetModels<Schema> as `${Uncapitalize<Model>}`]: ModelQueryHooks<Schema, Model, Options>;
+} & ProcedureHooks<Schema>;
+
+type ProcedureHookGroup<Schema extends SchemaDef> = {
+    [Name in GetProcedureNames<Schema>]: GetProcedure<Schema, Name> extends { mutation: true }
+        ? {
+              useMutation(
+                  options?: Omit<
+                      UseMutationOptions<ProcedureReturn<Schema, Name>, DefaultError, ProcedureEnvelope<Schema, Name>>,
+                      'mutationFn'
+                  > &
+                      QueryContext,
+              ): UseMutationResult<ProcedureReturn<Schema, Name>, DefaultError, ProcedureEnvelope<Schema, Name>>;
+          }
+        : {
+              useQuery: ProcedureHookFn<
+                  Schema,
+                  Name,
+                  Omit<ModelQueryOptions<ProcedureReturn<Schema, Name>>, 'optimisticUpdate'>,
+                  UseQueryResult<ProcedureReturn<Schema, Name>, DefaultError> & { queryKey: QueryKey }
+              >;
+
+              useSuspenseQuery: ProcedureHookFn<
+                  Schema,
+                  Name,
+                  Omit<ModelSuspenseQueryOptions<ProcedureReturn<Schema, Name>>, 'optimisticUpdate'>,
+                  UseSuspenseQueryResult<ProcedureReturn<Schema, Name>, DefaultError> & { queryKey: QueryKey }
+              >;
+
+              //   Infinite queries for procedures are currently disabled, will add back later if needed
+              //
+              //   useInfiniteQuery: ProcedureHookFn<
+              //       Schema,
+              //       Name,
+              //       ModelInfiniteQueryOptions<ProcedureReturn<Schema, Name>>,
+              //       ModelInfiniteQueryResult<InfiniteData<ProcedureReturn<Schema, Name>>>
+              //   >;
+
+              //   useSuspenseInfiniteQuery: ProcedureHookFn<
+              //       Schema,
+              //       Name,
+              //       ModelSuspenseInfiniteQueryOptions<ProcedureReturn<Schema, Name>>,
+              //       ModelSuspenseInfiniteQueryResult<InfiniteData<ProcedureReturn<Schema, Name>>>
+              //   >;
+          };
 };
 
+export type ProcedureHooks<Schema extends SchemaDef> = Schema extends { procedures: Record<string, any> }
+    ? {
+          /**
+           * Custom procedures.
+           */
+          $procs: ProcedureHookGroup<Schema>;
+      }
+    : {};
+
 // Note that we can potentially use TypeScript's mapped type to directly map from ORM contract, but that seems
 // to significantly slow down tsc performance ...
 export type ModelQueryHooks<
@@ -263,7 +330,7 @@ export function useClientQueries<Schema extends SchemaDef, Options extends Query
     schema: Schema,
     options?: QueryContext,
 ): ClientHooks<Schema, Options> {
-    return Object.keys(schema.models).reduce(
+    const result = Object.keys(schema.models).reduce(
         (acc, model) => {
             (acc as any)[lowerCaseFirst(model)] = useModelQueries<Schema, GetModels<Schema>, Options>(
                 schema,
@@ -274,6 +341,46 @@ export function useClientQueries<Schema extends SchemaDef, Options extends Query
         },
         {} as ClientHooks<Schema, Options>,
     );
+
+    const procedures = (schema as any).procedures as Record<string, { mutation?: boolean }> | undefined;
+    if (procedures) {
+        const buildProcedureHooks = (endpointModel: '$procs') => {
+            return Object.keys(procedures).reduce((acc, name) => {
+                const procDef = procedures[name];
+                if (procDef?.mutation) {
+                    acc[name] = {
+                        useMutation: (hookOptions?: any) =>
+                            useInternalMutation(schema, endpointModel, 'POST', name, { ...options, ...hookOptions }),
+                    };
+                } else {
+                    acc[name] = {
+                        useQuery: (args?: any, hookOptions?: any) =>
+                            useInternalQuery(schema, endpointModel, name, args, { ...options, ...hookOptions }),
+                        useSuspenseQuery: (args?: any, hookOptions?: any) =>
+                            useInternalSuspenseQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                        useInfiniteQuery: (args?: any, hookOptions?: any) =>
+                            useInternalInfiniteQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                        useSuspenseInfiniteQuery: (args?: any, hookOptions?: any) =>
+                            useInternalSuspenseInfiniteQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                    };
+                }
+                return acc;
+            }, {} as any);
+        };
+
+        (result as any).$procs = buildProcedureHooks('$procs');
+    }
+
+    return result;
 }
 
 /**

packages/clients/tanstack-query/src/svelte/index.svelte.ts

@@ -37,8 +37,11 @@ import type {
     FindFirstArgs,
     FindManyArgs,
     FindUniqueArgs,
+    GetProcedure,
+    GetProcedureNames,
     GroupByArgs,
     GroupByResult,
+    ProcedureEnvelope,
     QueryOptions,
     SelectSubset,
     SimplifiedPlainResult,
@@ -56,12 +59,23 @@ import { getQueryKey } from '../common/query-key';
 import type {
     ExtraMutationOptions,
     ExtraQueryOptions,
+    ProcedureReturn,
     QueryContext,
     TrimDelegateModelOperations,
     WithOptimistic,
 } from '../common/types';
 export type { FetchFn } from '@zenstackhq/client-helpers/fetch';
 
+type ProcedureHookFn<
+    Schema extends SchemaDef,
+    ProcName extends GetProcedureNames<Schema>,
+    Options,
+    Result,
+    Input = ProcedureEnvelope<Schema, ProcName>,
+> = { args: undefined } extends Input
+    ? (args?: Accessor<Input>, options?: Accessor<Options>) => Result
+    : (args: Accessor<Input>, options?: Accessor<Options>) => Result;
+
 /**
  * Key for setting and getting the global query context.
  */
@@ -88,6 +102,14 @@ function useQuerySettings() {
     return { endpoint: endpoint ?? DEFAULT_QUERY_ENDPOINT, ...rest };
 }
 
+function merge(rootOpt: unknown, opt: unknown): Accessor<any> {
+    return () => {
+        const rootOptVal = typeof rootOpt === 'function' ? (rootOpt as any)() : rootOpt;
+        const optVal = typeof opt === 'function' ? (opt as any)() : opt;
+        return { ...rootOptVal, ...optVal };
+    };
+}
+
 export type ModelQueryOptions<T> = Omit<CreateQueryOptions<T, DefaultError>, 'queryKey'> & ExtraQueryOptions;
 
 export type ModelQueryResult<T> = CreateQueryResult<WithOptimistic<T>, DefaultError> & { queryKey: QueryKey };
@@ -122,8 +144,51 @@ export type ModelMutationModelResult<
 
 export type ClientHooks<Schema extends SchemaDef, Options extends QueryOptions<Schema> = QueryOptions<Schema>> = {
     [Model in GetModels<Schema> as `${Uncapitalize<Model>}`]: ModelQueryHooks<Schema, Model, Options>;
+} & ProcedureHooks<Schema>;
+
+type ProcedureHookGroup<Schema extends SchemaDef> = {
+    [Name in GetProcedureNames<Schema>]: GetProcedure<Schema, Name> extends { mutation: true }
+        ? {
+              useMutation(
+                  options?: Omit<
+                      CreateMutationOptions<
+                          ProcedureReturn<Schema, Name>,
+                          DefaultError,
+                          ProcedureEnvelope<Schema, Name>
+                      >,
+                      'mutationFn'
+                  > &
+                      QueryContext,
+              ): CreateMutationResult<ProcedureReturn<Schema, Name>, DefaultError, ProcedureEnvelope<Schema, Name>>;
+          }
+        : {
+              useQuery: ProcedureHookFn<
+                  Schema,
+                  Name,
+                  Omit<ModelQueryOptions<ProcedureReturn<Schema, Name>>, 'optimisticUpdate'>,
+                  CreateQueryResult<ProcedureReturn<Schema, Name>, DefaultError> & { queryKey: QueryKey }
+              >;
+
+              //   Infinite queries for procedures are currently disabled, will add back later if needed
+              //
+              //   useInfiniteQuery: ProcedureHookFn<
+              //       Schema,
+              //       Name,
+              //       ModelInfiniteQueryOptions<ProcedureReturn<Schema, Name>>,
+              //       ModelInfiniteQueryResult<InfiniteData<ProcedureReturn<Schema, Name>>>
+              //   >;
+          };
 };
 
+export type ProcedureHooks<Schema extends SchemaDef> = Schema extends { procedures: Record<string, any> }
+    ? {
+          /**
+           * Custom procedures.
+           */
+          $procs: ProcedureHookGroup<Schema>;
+      }
+    : {};
+
 // Note that we can potentially use TypeScript's mapped type to directly map from ORM contract, but that seems
 // to significantly slow down tsc performance ...
 export type ModelQueryHooks<
@@ -212,7 +277,7 @@ export function useClientQueries<Schema extends SchemaDef, Options extends Query
     schema: Schema,
     options?: Accessor<QueryContext>,
 ): ClientHooks<Schema, Options> {
-    return Object.keys(schema.models).reduce(
+    const result = Object.keys(schema.models).reduce(
         (acc, model) => {
             (acc as any)[lowerCaseFirst(model)] = useModelQueries<Schema, GetModels<Schema>, Options>(
                 schema,
@@ -223,6 +288,33 @@ export function useClientQueries<Schema extends SchemaDef, Options extends Query
         },
         {} as ClientHooks<Schema, Options>,
     );
+
+    const procedures = (schema as any).procedures as Record<string, { mutation?: boolean }> | undefined;
+    if (procedures) {
+        const buildProcedureHooks = (endpointModel: '$procs') => {
+            return Object.keys(procedures).reduce((acc, name) => {
+                const procDef = procedures[name];
+                if (procDef?.mutation) {
+                    acc[name] = {
+                        useMutation: (hookOptions?: any) =>
+                            useInternalMutation(schema, endpointModel, 'POST', name, merge(options, hookOptions)),
+                    };
+                } else {
+                    acc[name] = {
+                        useQuery: (args?: any, hookOptions?: any) =>
+                            useInternalQuery(schema, endpointModel, name, args, merge(options, hookOptions)),
+                        useInfiniteQuery: (args?: any, hookOptions?: any) =>
+                            useInternalInfiniteQuery(schema, endpointModel, name, args, merge(options, hookOptions)),
+                    };
+                }
+                return acc;
+            }, {} as any);
+        };
+
+        (result as any).$procs = buildProcedureHooks('$procs');
+    }
+
+    return result;
 }
 
 /**
@@ -240,14 +332,6 @@ export function useModelQueries<
 
     const modelName = modelDef.name;
 
-    const merge = (rootOpt: unknown, opt: unknown): Accessor<any> => {
-        return () => {
-            const rootOptVal = typeof rootOpt === 'function' ? rootOpt() : rootOpt;
-            const optVal = typeof opt === 'function' ? opt() : opt;
-            return { ...rootOptVal, ...optVal };
-        };
-    };
-
     return {
         useFindUnique: (args: any, options?: any) => {
             return useInternalQuery(schema, modelName, 'findUnique', args, merge(rootOptions, options));