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,9 @@ 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', () => {
@@ -858,3 +861,22 @@ 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')
+    });
+  });
+});

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

packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap

@@ -242,20 +242,20 @@ type Query {
 `;
 
 exports[`custom operations Add entities to SQL schema add custom type, enum, and custom query to generated SQL schema 1`] = `
-"type post @model(timestamps: null) @auth(rules: [{allow: private}])
-{
-  title: String!
-  description: String
-  author: String
-}
-
-enum PostStatus {
+"enum PostStatus {
   draft
   pending
   approved
   published
 }
 
+type post @model(timestamps: null) @auth(rules: [{allow: private}])
+{
+  title: String!
+  description: String
+  author: String
+}
+
 type PostMeta @aws_cognito_user_pools
 {
   viewCount: Int
@@ -595,7 +595,11 @@ exports[`schema auth rules global public auth - multiple models 1`] = `
   "functionSlots": [],
   "jsFunctions": [],
   "lambdaFunctions": {},
-  "schema": "type A @model @auth(rules: [{allow: public, provider: apiKey}])
+  "schema": "enum DTired {
+  ?
+}
+
+type A @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   field: String
 }
@@ -621,10 +625,6 @@ type D @model @auth(rules: [{allow: public, provider: apiKey}])
   tired: DTired
   cId: ID
   c: C @belongsTo(references: ["cId"])
-}
-
-enum DTired {
-  ?
 }",
 }
 `;

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

@@ -1,16 +1,16 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
 exports[`CustomOperation transform dynamo schema Custom Mutation w required arg and enum 1`] = `
-"type Post @model @auth(rules: [{allow: private}])
-{
-  title: String
-}
-
-enum LikePostReactionType {
+"enum LikePostReactionType {
   :shipit:
   :risitas:
 }
 
+type Post @model @auth(rules: [{allow: private}])
+{
+  title: String
+}
+
 type Mutation {
   likePost(postId: String!, reactionType: LikePostReactionType): Post 
 }"
@@ -45,7 +45,7 @@ exports[`CustomOperation transform dynamo schema Custom mutation w inline boolea
 `;
 
 exports[`CustomOperation transform dynamo schema Custom mutation w inline custom return type 1`] = `
-"type LikePostReturnType 
+"type LikePostReturnType
 {
   stringField: String
   intField: Int
@@ -109,7 +109,7 @@ type Query {
 `;
 
 exports[`CustomOperation transform dynamo schema Custom query w inline custom return type 1`] = `
-"type GetPostDetailsReturnType 
+"type GetPostDetailsReturnType
 {
   stringField: String
   intField: Int
@@ -545,6 +545,18 @@ type Query {
 }"
 `;
 
+exports[`custom operations + custom type auth inheritance inline custom type inherits auth rules from referencing op 1`] = `
+"type MyQueryReturnType @aws_api_key
+{
+  fieldA: String
+  fieldB: Int
+}
+
+type Query {
+  myQuery: MyQueryReturnType @function(name: "myFn") @auth(rules: [{allow: public, provider: apiKey}])
+}"
+`;
+
 exports[`custom operations + custom type auth inheritance nested custom types inherit auth rules from top-level referencing op 1`] = `
 "type MyQueryReturnType @aws_api_key
 {