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', () => {
@@ -84,15 +85,15 @@ describe('custom operations return types', () => {
     });
 
     it('produces async function handler types', () => {
-      const handler = defineFunctionStub({})
+      const handler = defineFunctionStub({});
       const schema = a.schema({
         aQuery: a
           .query()
           .handler(a.handler.function(handler).async())
           .arguments({
             input: a.string(),
             content: a.string().required(),
-          })
+          }),
       });
 
       type Schema = ClientSchema<typeof schema>;
@@ -109,10 +110,7 @@ describe('custom operations return types', () => {
       type ExpectedResult = {
         success: boolean;
       } | null;
-      type ExpectedFunctionHandler = AppSyncResolverHandler<
-        ActualArgs,
-        void
-      >;
+      type ExpectedFunctionHandler = AppSyncResolverHandler<ActualArgs, void>;
       type _T1 = Expect<Equal<ActualArgs, ExpectedArgs>>;
       type _T2 = Expect<Equal<ActualResult, ExpectedResult>>;
       type _T3 = Expect<Equal<ActualHandler, ExpectedFunctionHandler>>;
@@ -858,3 +856,145 @@ 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', () => {
+    const operation: CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    > = a.query().arguments({
+      customArg: a.customType({
+        field1: a.string(),
+        field2: a.integer(),
+      }),
+    });
+  });
+
+  // Test to verify that RefType can be used as an argument in custom operations
+  it('accepts RefType in arguments', () => {
+    const operation: CustomOperation<
+      any,
+      'arguments' | 'for',
+      'queryCustomOperation'
+    > = a.query().arguments({
+      refArg: a.ref('SomeType'),
+    });
+  });
+
+  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

@@ -907,15 +907,15 @@ describe('CustomOperation transform', () => {
               .query()
               .arguments({})
               .handler(a.handler.function(fn1).async())
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
           expect(schema).toMatchSnapshot();
           expect(lambdaFunctions).toMatchObject({
             FnGetPostDetails: fn1,
           });
-        })
+        });
 
         test('defineFunction sync - async', () => {
           const fn1 = defineFunctionStub({});
@@ -927,15 +927,15 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1),
                 a.handler.function(fn1).async(),
               ])
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
           expect(schema).toMatchSnapshot();
           expect(lambdaFunctions).toMatchObject({
             FnGetPostDetails: fn1,
           });
-        })
+        });
 
         test('defineFunction sync - async with returns generates type errors', () => {
           const fn1 = defineFunctionStub({});
@@ -949,9 +949,9 @@ describe('CustomOperation transform', () => {
               ])
               .authorization((allow) => allow.authenticated())
               // @ts-expect-error
-              .returns({ })
+              .returns({}),
           });
-        })
+        });
 
         test('defineFunction async - async', () => {
           const fn1 = defineFunctionStub({});
@@ -965,7 +965,7 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1).async(),
                 a.handler.function(fn2).async(),
               ])
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
@@ -974,7 +974,7 @@ describe('CustomOperation transform', () => {
             FnGetPostDetails: fn1,
             FnGetPostDetails2: fn2,
           });
-        })
+        });
 
         test('defineFunction async - sync', () => {
           const fn1 = defineFunctionStub({});
@@ -987,12 +987,12 @@ describe('CustomOperation transform', () => {
                 a.handler.function(fn1),
               ])
               .returns(a.customType({}))
-              .authorization((allow) => allow.authenticated())
+              .authorization((allow) => allow.authenticated()),
           });
 
           const { schema, lambdaFunctions } = s.transform();
           expect(schema).toMatchSnapshot();
-        })
+        });
 
         test('pipeline / mix', () => {
           const fn1 = defineFunctionStub({});
@@ -1341,15 +1341,14 @@ describe('custom operations + custom type auth inheritance', () => {
 
   test('implicit custom type inherits auth rules from referencing op', () => {
     const s = a.schema({
+      MyQueryReturnType: a.customType({
+        fieldA: a.string(),
+        fieldB: a.integer(),
+      }),
       myQuery: a
         .query()
         .handler(a.handler.function('myFn'))
-        .returns(
-          a.customType({
-            fieldA: a.string(),
-            fieldB: a.integer(),
-          }),
-        )
+        .returns(a.ref('MyQueryReturnType'))
         .authorization((allow) => allow.publicApiKey()),
     });
 
@@ -1363,23 +1362,22 @@ describe('custom operations + custom type auth inheritance', () => {
 
   test('nested custom types inherit auth rules from top-level referencing op', () => {
     const s = a.schema({
+      MyQueryReturnType: a.customType({
+        fieldA: a.string(),
+        fieldB: a.integer(),
+        nestedCustomType: a.customType({
+          nestedA: a.string(),
+          nestedB: a.string(),
+          grandChild: a.customType({
+            grandA: a.string(),
+            grandB: a.string(),
+          }),
+        }),
+      }),
       myQuery: a
         .query()
         .handler(a.handler.function('myFn'))
-        .returns(
-          a.customType({
-            fieldA: a.string(),
-            fieldB: a.integer(),
-            nestedCustomType: a.customType({
-              nestedA: a.string(),
-              nestedB: a.string(),
-              grandChild: a.customType({
-                grandA: a.string(),
-                grandB: a.string(),
-              }),
-            }),
-          }),
-        )
+        .returns(a.ref('MyQueryReturnType'))
         .authorization((allow) => allow.publicApiKey()),
     });
 
@@ -1401,6 +1399,28 @@ describe('custom operations + custom type auth inheritance', () => {
     );
   });
 
+  test('inline custom type inherits auth rules from referencing op', () => {
+    const s = a.schema({
+      myQuery: a
+        .query()
+        .handler(a.handler.function('myFn'))
+        .returns(
+          a.customType({
+            fieldA: a.string(),
+            fieldB: a.integer(),
+          }),
+        )
+        .authorization((allow) => allow.publicApiKey()),
+    });
+
+    const result = s.transform().schema;
+
+    expect(result).toMatchSnapshot();
+    expect(result).toEqual(
+      expect.stringContaining('type MyQueryReturnType @aws_api_key\n{'),
+    );
+  });
+
   test('top-level custom type with nested top-level custom types inherits combined auth rules from referencing ops', () => {
     const s = a.schema({
       myQuery: a