feat: custom procs

Author: mwillbanks

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: [{ 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/react.ts

@@ -38,6 +38,7 @@ import type {
     FindUniqueArgs,
     GroupByArgs,
     GroupByResult,
+    ProcedureFunc,
     QueryOptions,
     SelectSubset,
     SimplifiedPlainResult,
@@ -61,6 +62,25 @@ import type {
 } from './common/types';
 export type { FetchFn } from '@zenstackhq/client-helpers/fetch';
 
+type ExtractProcedures<Schema extends SchemaDef> = Schema extends { procedures: Record<string, any> }
+    ? NonNullable<Schema['procedures']>
+    : never;
+
+type ProcedureArgsTuple<Schema extends SchemaDef, Name extends keyof ExtractProcedures<Schema>> = Parameters<
+    ProcedureFunc<Schema, ExtractProcedures<Schema>[Name]>
+>;
+
+type ProcedureReturn<Schema extends SchemaDef, Name extends keyof ExtractProcedures<Schema>> = Awaited<ReturnType<
+    ProcedureFunc<Schema, ExtractProcedures<Schema>[Name]>
+>>;
+
+type ProcedurePayload<Schema extends SchemaDef, Name extends keyof ExtractProcedures<Schema>> =
+    ProcedureArgsTuple<Schema, Name> extends []
+        ? undefined
+        : ProcedureArgsTuple<Schema, Name> extends [infer A]
+          ? A
+          : ProcedureArgsTuple<Schema, Name>;
+
 /**
  * React context for query settings.
  */
@@ -133,8 +153,56 @@ 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 keyof ExtractProcedures<Schema>]: ExtractProcedures<Schema>[Name] extends { mutation: true }
+        ? {
+              useMutation(
+                  options?: Omit<
+                      UseMutationOptions<ProcedureReturn<Schema, Name>, DefaultError, ProcedurePayload<Schema, Name>>,
+                      'mutationFn'
+                  > &
+                      QueryContext,
+              ): UseMutationResult<ProcedureReturn<Schema, Name>, DefaultError, ProcedurePayload<Schema, Name>>;
+          }
+        : {
+              useQuery(
+                  args?: ProcedurePayload<Schema, Name>,
+                  options?: ModelQueryOptions<ProcedureReturn<Schema, Name>>,
+              ): ModelQueryResult<ProcedureReturn<Schema, Name>>;
+
+              useSuspenseQuery(
+                  args?: ProcedurePayload<Schema, Name>,
+                  options?: ModelSuspenseQueryOptions<ProcedureReturn<Schema, Name>>,
+              ): ModelSuspenseQueryResult<ProcedureReturn<Schema, Name>>;
+
+              useInfiniteQuery(
+                  args?: ProcedurePayload<Schema, Name>,
+                  options?: ModelInfiniteQueryOptions<ProcedureReturn<Schema, Name>>,
+              ): ModelInfiniteQueryResult<InfiniteData<ProcedureReturn<Schema, Name>>>;
+
+              useSuspenseInfiniteQuery(
+                  args?: ProcedurePayload<Schema, Name>,
+                  options?: ModelSuspenseInfiniteQueryOptions<ProcedureReturn<Schema, Name>>,
+              ): ModelSuspenseInfiniteQueryResult<InfiniteData<ProcedureReturn<Schema, Name>>>;
+          };
 };
 
+export type ProcedureHooks<Schema extends SchemaDef> = Schema extends { procedures: Record<string, any> }
+    ? {
+          /**
+           * Preferred procedures API.
+           */
+          $procs: ProcedureHookGroup<Schema>;
+
+          /**
+           * Backward-compatible procedures API.
+           */
+          $procedures: 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 +331,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 +342,137 @@ 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' | '$procedures') => {
+            return Object.keys(procedures).reduce((acc, name) => {
+                const procDef = procedures[name];
+                if (procDef?.mutation) {
+                    acc[name] = {
+                        useMutation: (hookOptions?: any) =>
+                            useInternalProcedureMutation(schema, endpointModel, name, { ...options, ...hookOptions }),
+                    };
+                } else {
+                    acc[name] = {
+                        useQuery: (args?: any, hookOptions?: any) =>
+                            useInternalProcedureQuery(schema, endpointModel, name, args, { ...options, ...hookOptions }),
+                        useSuspenseQuery: (args?: any, hookOptions?: any) =>
+                            useInternalProcedureSuspenseQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                        useInfiniteQuery: (args?: any, hookOptions?: any) =>
+                            useInternalProcedureInfiniteQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                        useSuspenseInfiniteQuery: (args?: any, hookOptions?: any) =>
+                            useInternalProcedureSuspenseInfiniteQuery(schema, endpointModel, name, args, {
+                                ...options,
+                                ...hookOptions,
+                            }),
+                    };
+                }
+                return acc;
+            }, {} as any);
+        };
+
+        (result as any).$procs = buildProcedureHooks('$procs');
+        (result as any).$procedures = buildProcedureHooks('$procedures');
+    }
+
+    return result;
+}
+
+export function useInternalProcedureQuery<TQueryFnData, TData>(
+    _schema: SchemaDef,
+    endpointModel: '$procs' | '$procedures',
+    procedure: string,
+    args?: unknown,
+    options?: Omit<UseQueryOptions<TQueryFnData, DefaultError, TData>, 'queryKey'> & ExtraQueryOptions,
+) {
+    const { endpoint, fetch } = useFetchOptions(options);
+    const reqUrl = makeUrl(endpoint, endpointModel, procedure, args);
+    const queryKey = getQueryKey(endpointModel, procedure, args, {
+        infinite: false,
+        optimisticUpdate: false,
+    });
+    return {
+        queryKey,
+        ...useQuery({
+            queryKey,
+            queryFn: ({ signal }) => fetcher<TQueryFnData>(reqUrl, { signal }, fetch),
+            ...options,
+        }),
+    };
+}
+
+export function useInternalProcedureSuspenseQuery<TQueryFnData, TData>(
+    _schema: SchemaDef,
+    endpointModel: '$procs' | '$procedures',
+    procedure: string,
+    args?: unknown,
+    options?: Omit<UseSuspenseQueryOptions<TQueryFnData, DefaultError, TData>, 'queryKey'> & ExtraQueryOptions,
+) {
+    const { endpoint, fetch } = useFetchOptions(options);
+    const reqUrl = makeUrl(endpoint, endpointModel, procedure, args);
+    const queryKey = getQueryKey(endpointModel, procedure, args, {
+        infinite: false,
+        optimisticUpdate: false,
+    });
+    return {
+        queryKey,
+        ...useSuspenseQuery({
+            queryKey,
+            queryFn: ({ signal }) => fetcher<TQueryFnData>(reqUrl, { signal }, fetch),
+            ...options,
+        }),
+    };
+}
+
+export function useInternalProcedureInfiniteQuery<TQueryFnData, TData>(
+    schema: SchemaDef,
+    endpointModel: '$procs' | '$procedures',
+    procedure: string,
+    args: unknown,
+    options:
+        | (Omit<
+              UseInfiniteQueryOptions<TQueryFnData, DefaultError, InfiniteData<TData>>,
+              'queryKey' | 'initialPageParam'
+          > &
+              QueryContext)
+        | undefined,
+) {
+    return useInternalInfiniteQuery(schema, endpointModel, procedure, args, options);
+}
+
+export function useInternalProcedureSuspenseInfiniteQuery<TQueryFnData, TData>(
+    schema: SchemaDef,
+    endpointModel: '$procs' | '$procedures',
+    procedure: string,
+    args: unknown,
+    options: Omit<
+        UseSuspenseInfiniteQueryOptions<TQueryFnData, DefaultError, InfiniteData<TData>> & QueryContext,
+        'queryKey' | 'initialPageParam'
+    >,
+) {
+    return useInternalSuspenseInfiniteQuery(schema, endpointModel, procedure, args, options);
+}
+
+export function useInternalProcedureMutation<TArgs, R = any>(
+    _schema: SchemaDef,
+    endpointModel: '$procs' | '$procedures',
+    procedure: string,
+    options?: Omit<UseMutationOptions<R, DefaultError, TArgs>, 'mutationFn'> & QueryContext,
+) {
+    const { endpoint, fetch } = useFetchOptions(options);
+    const mutationFn = (data: any) => {
+        const reqUrl = makeUrl(endpoint, endpointModel, procedure, data);
+        return fetcher<R>(reqUrl, { method: 'POST' }, fetch) as Promise<R>;
+    };
+
+    return useMutation({ ...options, mutationFn });
 }
 
 /**