feat: allow CustomType and RefType in arguments

.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/src/CustomOperation.ts

@@ -29,7 +29,7 @@ type CustomOperationBrand =
   | typeof subscriptionBrand
   | typeof generationBrand;
 
-type CustomArguments = Record<string, BaseModelField | EnumType>;
+type CustomArguments = Record<string, BaseModelField | EnumType | CustomType<any> | RefType<any, any>>;
 type SubscriptionSource = RefType<any, any>;
 type InternalSubscriptionSource = InternalRef;
 type CustomReturnType = RefType<any> | CustomType<any>;

packages/data-schema/src/runtime/bridge-types.ts

@@ -16,6 +16,8 @@ import { Observable } from 'rxjs';
 import { CustomHeaders, ModelSortDirection } from './client';
 
 import { AiAction, AiCategory } from './internals/ai/getCustomUserAgentDetails';
+import { CustomType } from '../CustomType';
+import { RefType } from '../RefType';
 
 export declare namespace AmplifyServer {
   export interface ContextToken {
@@ -165,7 +167,7 @@ export type CustomOperationArguments = Record<string, CustomOperationArgument>;
 
 export interface CustomOperationArgument {
   name: string;
-  type: InputFieldType;
+  type: InputFieldType | CustomType<any> | RefType<any, any>;
   isArray: boolean;
   isRequired: boolean;
   isArrayNullable?: boolean;
@@ -240,7 +242,7 @@ export type FieldType =
   | ModelFieldType
   | NonModelFieldType;
 
-export type InputFieldType = ScalarType | EnumType | InputType;
+export type InputFieldType = ScalarType | EnumType | InputType | CustomType<any> | RefType<any, any>;
 
 export type FieldAttribute = ModelAttribute;