feat: allow customtypes ref in arguments/main

.changeset/polite-crabs-rush.md

@@ -0,0 +1,5 @@
+---
+"@aws-amplify/data-schema": minor
+---
+
+Allow CustomType and RefType in arguments for custom operations

packages/data-schema/__tests__/CustomOperations.test-d.ts

@@ -16,6 +16,7 @@ import type {
 import { configure } from '../src/ModelSchema';
 import { Nullable } from '../src/ModelField';
 import { defineFunctionStub } from './utils';
+import type { CustomOperation } from '../src/CustomOperation';
 
 describe('custom operations return types', () => {
   describe('when .ref() a basic custom type', () => {
@@ -855,3 +856,153 @@ describe('.for() modifier', () => {
     a.subscription().for(a.ref('Model'));
   });
 });
+
+describe('.arguments() modifier', () => {
+  // Test to verify that CustomType can be used as an argument in custom operations
+  it('accepts CustomType in arguments', () => {
+    type ExpectedType = CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    >;
+
+    const operation = a.query().arguments({
+      customArg: a.customType({
+        field1: a.string(),
+        field2: a.integer(),
+      }),
+    });
+
+    type test = Expect<Equal<ExpectedType, typeof operation>>;
+  });
+
+  // Test to verify that RefType can be used as an argument in custom operations
+  it('accepts RefType in arguments', () => {
+    type ExpectedType = CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    >;
+
+    const operation = a.query().arguments({
+      refArg: a.ref('SomeType'),
+    });
+
+    type test = Expect<Equal<ExpectedType, typeof operation>>;
+  });
+
+  it('handles deeply nested custom types', () => {
+    const schema = a.schema({
+      DeepNested: a.customType({
+        level1: a.customType({
+          level2: a.customType({
+            level3: a.string(),
+          }),
+        }),
+      }),
+      deepQuery: a
+        .query()
+        .arguments({
+          input: a.ref('DeepNested'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        level1?: {
+          level2?: {
+            level3?: string | null;
+          } | null;
+        } | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['deepQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+
+  it('handles mixed custom types and refs', () => {
+    const schema = a.schema({
+      RefType: a.customType({
+        field: a.string(),
+      }),
+      MixedType: a.customType({
+        nested: a.customType({
+          refField: a.ref('RefType'),
+          customField: a.customType({
+            deepField: a.integer(),
+          }),
+        }),
+      }),
+      mixedQuery: a
+        .query()
+        .arguments({
+          input: a.ref('MixedType'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        nested?: {
+          refField?: {
+            field?: string | null;
+          } | null;
+          customField?: {
+            deepField?: number | null;
+          } | null;
+        } | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['mixedQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+
+  it('handles RefType with multi-layered custom types in nested structures', () => {
+    const schema = a.schema({
+      NestedCustomType: a.customType({
+        nestedField: a.string(),
+      }),
+      RefType: a.customType({
+        field: a.string(),
+        nestedCustom: a.ref('NestedCustomType'),
+      }),
+      OuterType: a.customType({
+        refField: a.ref('RefType'),
+        otherField: a.integer(),
+      }),
+      complexQuery: a
+        .query()
+        .arguments({
+          input: a.ref('OuterType'),
+        })
+        .returns(a.string()),
+    });
+
+    type Schema = ClientSchema<typeof schema>;
+
+    type ExpectedArgs = {
+      input?: {
+        refField?: {
+          field?: string | null;
+          nestedCustom?: {
+            nestedField?: string | null;
+          } | null;
+        } | null;
+        otherField?: number | null;
+      } | null;
+    };
+
+    type ActualArgs = Schema['complexQuery']['args'];
+
+    type Test = Expect<Equal<ActualArgs, ExpectedArgs>>;
+  });
+});

packages/data-schema/__tests__/CustomOperations.test.ts

@@ -1089,6 +1089,195 @@ describe('CustomOperation transform', () => {
     });
   });
 
+  describe('custom operations with custom types and refs', () => {
+    test('Schema with custom query using custom type argument', () => {
+      const s = a
+        .schema({
+          CustomArgType: a.customType({
+            field: a.string(),
+          }),
+          queryWithCustomTypeArg: a
+            .query()
+            .arguments({ customArg: a.ref('CustomArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom query using ref argument', () => {
+      const s = a
+        .schema({
+          RefArgType: a.customType({
+            field: a.string(),
+          }),
+          queryWithRefArg: a
+            .query()
+            .arguments({ refArg: a.ref('RefArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using custom type argument', () => {
+      const s = a
+        .schema({
+          CustomArgType: a.customType({
+            field: a.string(),
+          }),
+          mutateWithCustomTypeArg: a
+            .mutation()
+            .arguments({ customArg: a.ref('CustomArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using ref argument', () => {
+      const s = a
+        .schema({
+          RefArgType: a.customType({
+            field: a.string(),
+          }),
+          mutationWithRefArg: a
+            .mutation()
+            .arguments({ refArg: a.ref('RefArgType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom mutation using enum ref in nested custom type', () => {
+      const s = a
+        .schema({
+          NestedCustomType: a.customType({
+            field: a.string(),
+            enumField: a.ref('SomeEnum'),
+          }),
+          SomeEnum: a.enum(['VALUE1', 'VALUE2']),
+          mutateWithNestedEnumRef: a
+            .mutation()
+            .arguments({ nestedArg: a.ref('NestedCustomType') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom query using enum ref as argument', () => {
+      const s = a
+        .schema({
+          SomeEnum: a.enum(['OPTION1', 'OPTION2']),
+          queryWithEnumRefArg: a
+            .query()
+            .arguments({ enumArg: a.ref('SomeEnum') })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+    });
+
+    test('Schema with custom operation using enum directly in arguments', () => {
+      const s = a
+        .schema({
+          directEnumQuery: a
+            .query()
+            .arguments({
+              status: a.enum(['PENDING', 'APPROVED', 'REJECTED']),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum DirectEnumQueryStatus {');
+      expect(result).toContain(
+        'directEnumQuery(status: DirectEnumQueryStatus): String',
+      );
+    });
+
+    test('Schema with custom operation using enum in inline custom type arguments', () => {
+      const s = a
+        .schema({
+          inlineEnumQuery: a
+            .query()
+            .arguments({
+              input: a.customType({
+                name: a.string(),
+                status: a.enum(['ACTIVE', 'INACTIVE']),
+              }),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum InlineEnumQueryInputStatus {');
+      expect(result).toContain('input InlineEnumQueryInput {');
+      expect(result).toContain(
+        'inlineEnumQuery(input: InlineEnumQueryInput): String',
+      );
+    });
+
+    test('Schema with custom operation using enum in referenced custom type arguments', () => {
+      const s = a
+        .schema({
+          StatusInput: a.customType({
+            name: a.string(),
+            status: a.enum(['NEW', 'IN_PROGRESS', 'COMPLETED']),
+          }),
+          refEnumQuery: a
+            .query()
+            .arguments({
+              input: a.ref('StatusInput'),
+            })
+            .returns(a.string())
+            .handler(a.handler.function('myFunc')),
+        })
+        .authorization((allow) => allow.publicApiKey());
+
+      const result = s.transform().schema;
+
+      expect(result).toMatchSnapshot();
+      expect(result).toContain('enum StatusInputStatus {');
+      expect(result).toContain('input StatusInput {');
+      expect(result).toContain('refEnumQuery(input: StatusInput): String');
+    });
+  });
+
   const fakeSecret = () => ({}) as any;
 
   const datasourceConfigMySQL = {