fix(orm,server): fix procedure slicing in toJSONSchema and RPC OpenAPI spec (#2596)

Author: ymc9

packages/orm/src/client/zod/factory.ts

@@ -220,9 +220,11 @@ export class ZodSchemaFactory<
             this.makeAggregateSchema(m);
             this.makeGroupBySchema(m);
         }
-        // Eagerly build args schemas for all procedures.
+        // Eagerly build args schemas for allowed procedures only.
         for (const procName of Object.keys(this.schema.procedures ?? {})) {
-            this.makeProcedureArgsSchema(procName);
+            if (this.isProcedureAllowed(procName)) {
+                this.makeProcedureArgsSchema(procName);
+            }
         }
         return z.toJSONSchema(this.schemaRegistry, { unrepresentable: 'any' });
     }
@@ -2495,6 +2497,29 @@ export class ZodSchemaFactory<
         return true;
     }
 
+    private isProcedureAllowed(procName: string): boolean {
+        const slicing = this.options.slicing;
+        if (!slicing) {
+            return true;
+        }
+
+        const { includedProcedures, excludedProcedures } = slicing;
+
+        if (includedProcedures !== undefined) {
+            if (!(includedProcedures as readonly string[]).includes(procName)) {
+                return false;
+            }
+        }
+
+        if (excludedProcedures !== undefined) {
+            if ((excludedProcedures as readonly string[]).includes(procName)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
     // #endregion
 }
 

packages/server/src/api/rpc/openapi.ts

@@ -1,6 +1,6 @@
 import { lowerCaseFirst, upperCaseFirst } from '@zenstackhq/common-helpers';
 import { CoreCrudOperations, createQuerySchemaFactory, type ZodSchemaFactory } from '@zenstackhq/orm';
-import type { BuiltinType, ModelDef, ProcedureDef, SchemaDef, TypeDefDef } from '@zenstackhq/orm/schema';
+import type { BuiltinType, EnumDef, ModelDef, ProcedureDef, SchemaDef, TypeDefDef } from '@zenstackhq/orm/schema';
 import type { OpenAPIV3_1 } from 'openapi-types';
 import type { RPCApiHandlerOptions } from '.';
 import { PROCEDURE_ROUTE_PREFIXES } from '../common/procedures';
@@ -468,6 +468,7 @@ export class RPCApiSpecGenerator<Schema extends SchemaDef = SchemaDef> {
         } else {
             if (hasParams) {
                 op['requestBody'] = {
+                    ...(hasRequiredParams && { required: true }),
                     content: {
                         [JSON_CT]: { schema: envelopeSchema },
                     },
@@ -529,6 +530,12 @@ export class RPCApiSpecGenerator<Schema extends SchemaDef = SchemaDef> {
     }
 
     private generateSharedSchemas(): Record<string, SchemaObject> {
+        // Generate schemas for enums
+        const enumSchemas: Record<string, SchemaObject> = {};
+        for (const [enumName, enumDef] of Object.entries(this.schema.enums ?? {})) {
+            enumSchemas[enumName] = this.buildEnumSchema(enumDef);
+        }
+
         // Generate schemas for typedefs (e.g. `type Address { city String }`)
         const typeDefSchemas: Record<string, SchemaObject> = {};
         for (const [typeName, typeDef] of Object.entries(this.schema.typeDefs ?? {})) {
@@ -542,6 +549,7 @@ export class RPCApiSpecGenerator<Schema extends SchemaDef = SchemaDef> {
         }
 
         return {
+            ...enumSchemas,
             ...typeDefSchemas,
             ...modelEntitySchemas,
             _integer: { type: 'integer', minimum: -9007199254740991, maximum: 9007199254740991 },
@@ -637,6 +645,10 @@ export class RPCApiSpecGenerator<Schema extends SchemaDef = SchemaDef> {
     /**
      * Builds a JSON Schema object describing a custom typedef
      */
+    private buildEnumSchema(enumDef: EnumDef): SchemaObject {
+        return { type: 'string', enum: Object.values(enumDef.values) };
+    }
+
     private buildTypeDefSchema(typeDef: TypeDefDef): SchemaObject {
         const properties: Record<string, SchemaObject | ReferenceObject> = {};
         const required: string[] = [];

packages/server/test/openapi/baseline/rpc.baseline.yaml

@@ -4225,6 +4225,7 @@ paths:
                             schema:
                                 $ref: "#/components/schemas/_rpcErrorResponse"
             requestBody:
+                required: true
                 content:
                     application/json:
                         schema:

packages/server/test/openapi/rpc-openapi.test.ts

@@ -444,6 +444,33 @@ describe('RPC OpenAPI spec generation - response data shapes', () => {
         expect(postSchema.required).toContain('title');
         expect(postSchema.required).toContain('published');
     });
+
+    it('enum schemas are present in components and enum fields reference them', async () => {
+        const enumSchema = `
+enum Role {
+    USER
+    ADMIN
+}
+
+model User {
+    id   Int    @id @default(autoincrement())
+    role Role
+}
+`;
+        const client = await createTestClient(enumSchema);
+        const handler = new RPCApiHandler({ schema: client.$schema });
+        const s = await generateSpec(handler);
+
+        // Enum component must be registered so the $ref resolves
+        expect(s.components?.schemas?.['Role']).toBeDefined();
+        expect(s.components?.schemas?.['Role'].type).toBe('string');
+        expect(s.components?.schemas?.['Role'].enum).toEqual(['USER', 'ADMIN']);
+
+        // The entity schema must reference the enum via $ref
+        const userSchema = s.components?.schemas?.['User'] as any;
+        const roleField = userSchema?.properties?.role;
+        expect(roleField?.['$ref']).toBe('#/components/schemas/Role');
+    });
 });
 
 describe('RPC OpenAPI spec generation - operationIds', () => {
@@ -753,6 +780,7 @@ procedure optionalSearch(query: String?): User[]
 
         const operation = spec?.paths?.['/$procs/createUser']?.post;
         expect(operation?.requestBody).toBeDefined();
+        expect((operation?.requestBody as any)?.required).toBe(true);
         const bodySchema = (operation?.requestBody as any)?.content?.['application/json']?.schema;
         // args is a $ref to the registered ProcArgs component schema
         const argsRef = bodySchema?.properties?.args?.$ref;
@@ -763,6 +791,23 @@ procedure optionalSearch(query: String?): User[]
         expect(argsSchema?.required).toContain('name');
     });
 
+    it('mutation with only optional params does not set requestBody.required', async () => {
+        const optionalMutationSchema = `
+model User {
+    id Int @id @default(autoincrement())
+    name String
+}
+mutation procedure softDelete(id: Int?): User
+`;
+        const client = await createTestClient(optionalMutationSchema);
+        const handler = new RPCApiHandler({ schema: client.$schema });
+        const spec = await generateSpec(handler);
+
+        const operation = spec?.paths?.['/$procs/softDelete']?.post;
+        expect(operation?.requestBody).toBeDefined();
+        expect((operation?.requestBody as any)?.required).toBeUndefined();
+    });
+
     it('optional procedure params are not in required array', async () => {
         const client = await createTestClient(procSchema);
         const handler = new RPCApiHandler({ schema: client.$schema });
@@ -799,6 +844,33 @@ procedure optionalSearch(query: String?): User[]
         expect(spec.paths?.['/$procs/getUser']).toBeUndefined();
         expect(spec.paths?.['/$procs/createUser']).toBeDefined();
     });
+
+    it('slicing excludedProcedures removes procedure args from components schemas', async () => {
+        const client = await createTestClient(procSchema);
+        const handler = new RPCApiHandler({
+            schema: client.$schema,
+            queryOptions: { slicing: { excludedProcedures: ['getUser'] as any } },
+        });
+        const spec = await generateSpec(handler);
+
+        const schemaKeys = Object.keys(spec.components?.schemas ?? {});
+        expect(schemaKeys.some((k) => k.startsWith('getUser'))).toBe(false);
+        expect(schemaKeys.some((k) => k.startsWith('createUser'))).toBe(true);
+    });
+
+    it('slicing includedProcedures removes non-listed procedure args from components schemas', async () => {
+        const client = await createTestClient(procSchema);
+        const handler = new RPCApiHandler({
+            schema: client.$schema,
+            queryOptions: { slicing: { includedProcedures: ['createUser'] as any } },
+        });
+        const spec = await generateSpec(handler);
+
+        const schemaKeys = Object.keys(spec.components?.schemas ?? {});
+        expect(schemaKeys.some((k) => k.startsWith('createUser'))).toBe(true);
+        expect(schemaKeys.some((k) => k.startsWith('getUser'))).toBe(false);
+        expect(schemaKeys.some((k) => k.startsWith('optionalSearch'))).toBe(false);
+    });
 });
 
 describe('RPC OpenAPI spec generation - respectAccessPolicies', () => {