feat(client types): enable custom subscriptions

.changeset/nine-kangaroos-promise.md

@@ -0,0 +1,6 @@
+---
+"@aws-amplify/data-schema-types": patch
+"@aws-amplify/data-schema": patch
+---
+
+client types for custom subscriptions

packages/data-schema-types/src/client/index.ts

@@ -1,6 +1,6 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import {
+import type {
   DeepReadOnlyObject,
   UnwrapArray,
   UnionToIntersection,
@@ -744,6 +744,28 @@ export type CustomMutations<
   ModelMeta extends Record<any, any> = ExtractModelMeta<Schema>,
 > = CustomOperations<Schema, 'Mutation', Context, ModelMeta>;
 
+export type CustomSubscriptions<
+  Schema extends Record<any, any>,
+  Context extends ContextType = 'CLIENT',
+  ModelMeta extends Record<any, any> = ExtractModelMeta<Schema>,
+> = CustomOperations<Schema, 'Subscription', Context, ModelMeta>;
+
+type CustomOperationMethodOptions = {
+  // selectionSet?: SelectionSet;
+  authMode?: AuthMode;
+  authToken?: string;
+  headers?: CustomHeaders;
+};
+
+/**
+ * Generates Custom Operations function params based on whether .arguments() were specified in the schema builder
+ */
+type CustomOperationFnParams<Args extends Record<string, unknown> | never> = [
+  Args,
+] extends [never]
+  ? [CustomOperationMethodOptions?]
+  : [Args, CustomOperationMethodOptions?];
+
 export type CustomOperations<
   Schema extends Record<any, any>,
   OperationType extends 'Query' | 'Mutation' | 'Subscription',
@@ -754,36 +776,27 @@ export type CustomOperations<
     ? OpName
     : never]: {
     CLIENT: (
-      input: ModelMeta['customOperations'][OpName]['arguments'],
-      options?: {
-        // selectionSet?: SelectionSet;
-        authMode?: AuthMode;
-        authToken?: string;
-        headers?: CustomHeaders;
-      },
-    ) => SingularReturnValue<
-      ModelMeta['customOperations'][OpName]['returnType']
-    >;
+      ...params: CustomOperationFnParams<
+        ModelMeta['customOperations'][OpName]['arguments']
+      >
+    ) => // we only generate subscriptions on the clientside; so this isn't applied to COOKIES | REQUEST
+    ModelMeta['customOperations'][OpName]['typeName'] extends 'Subscription'
+      ? ObservedReturnValue<ModelMeta['customOperations'][OpName]['returnType']>
+      : SingularReturnValue<
+          ModelMeta['customOperations'][OpName]['returnType']
+        >;
     COOKIES: (
-      input: ModelMeta['customOperations'][OpName]['arguments'],
-      options?: {
-        // selectionSet?: SelectionSet;
-        authMode?: AuthMode;
-        authToken?: string;
-        headers?: CustomHeaders;
-      },
+      ...params: CustomOperationFnParams<
+        ModelMeta['customOperations'][OpName]['arguments']
+      >
     ) => SingularReturnValue<
       ModelMeta['customOperations'][OpName]['returnType']
     >;
     REQUEST: (
       contextSpec: any,
-      input: ModelMeta['customOperations'][OpName]['arguments'],
-      options?: {
-        // selectionSet?: SelectionSet;
-        authMode?: AuthMode;
-        authToken?: string;
-        headers?: CustomHeaders;
-      },
+      ...params: CustomOperationFnParams<
+        ModelMeta['customOperations'][OpName]['arguments']
+      >
     ) => SingularReturnValue<
       ModelMeta['customOperations'][OpName]['returnType']
     >;

packages/data-schema-types/src/client/index.v3.ts

@@ -38,3 +38,9 @@ export type CustomMutations<
   _Context extends string = 'CLIENT',
   _ModelMeta extends Record<any, any> = ExtractModelMeta<Schema>,
 > = any;
+
+export type CustomSubscriptions<
+  Schema extends Record<any, any>,
+  _Context extends string = 'CLIENT',
+  _ModelMeta extends Record<any, any> = ExtractModelMeta<Schema>,
+> = any;

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

@@ -39,6 +39,19 @@ describe('Custom Operations mapper utils', () => {
     type T = Expect<Equal<Actual, Expected>>;
   });
 
+  test('can resolve omitted arguments type from custom op', () => {
+    const aMutation = a
+      .mutation()
+      .returns(a.string())
+      .authorization([a.allow.public()])
+      .handler(a.handler.function('asdf'));
+
+    type Actual = CustomOpArguments<OpShape<typeof aMutation>>;
+    type Expected = never;
+
+    type T = Expect<Equal<Actual, Expected>>;
+  });
+
   test('can select custom op shapes from a schema', () => {
     const aQuery = a
       .query()

packages/data-schema/src/CustomOperation.ts

@@ -62,7 +62,7 @@ type InternalCustomData = CustomData & {
 };
 
 export type CustomOperationParamShape = {
-  arguments: CustomArguments;
+  arguments: CustomArguments | null;
   returnType: CustomReturnType | null;
   functionRef: string | null;
   authorization: Authorization<any, any, any>[];
@@ -206,7 +206,7 @@ export type QueryCustomOperation = CustomOperation<
 
 export function query(): CustomOperation<
   {
-    arguments: CustomArguments;
+    arguments: null;
     returnType: null;
     functionRef: null;
     authorization: [];
@@ -227,7 +227,7 @@ export type MutationCustomOperation = CustomOperation<
 
 export function mutation(): CustomOperation<
   {
-    arguments: CustomArguments;
+    arguments: null;
     returnType: null;
     functionRef: null;
     authorization: [];
@@ -248,7 +248,7 @@ export type SubscriptionCustomOperation = CustomOperation<
 
 export function subscription(): CustomOperation<
   {
-    arguments: CustomArguments;
+    arguments: null;
     returnType: null;
     functionRef: null;
     authorization: [];

packages/data-schema/src/MappedTypes/CustomOperations.ts

@@ -55,15 +55,18 @@ export type CustomOpShapes<Schema extends GenericModelSchema<any>> = {
 /**
  * Digs out custom operation arguments, mapped to the intended graphql types.
  */
-export type CustomOpArguments<Shape extends CustomOperationParamShape> = {
-  [FieldName in keyof Shape['arguments']]: Shape['arguments'][FieldName] extends ModelField<
-    infer R,
-    any,
-    any
-  >
-    ? R
-    : never;
-};
+export type CustomOpArguments<Shape extends CustomOperationParamShape> =
+  Shape['arguments'] extends null
+    ? never
+    : {
+        [FieldName in keyof Shape['arguments']]: Shape['arguments'][FieldName] extends ModelField<
+          infer R,
+          any,
+          any
+        >
+          ? R
+          : never;
+      };
 
 /**
  * Computes the return type from the `returnType` of a custom operation shape.

packages/integration-tests/__tests__/operation-params.test-d.ts

@@ -47,18 +47,15 @@ describe('Basic operations', () => {
         await client.models.Post.get();
       });
 
-      // TODO: broken because JS lib is using old types package name and we're not able to share the symbol correctly.
-      // Re-enable once we update the JS dep on
-
-      // test('parameter must contain PK', async () => {
-      //   // @ts-expect-error
-      //   await client.models.Post.get({});
-      // });
-
-      // test('parameter must not contain extra fields', async () => {
-      //   // @ts-expect-error
-      //   await client.models.Post.get({ id: 'some-id', title: 'whatever' });
-      // });
+      test('parameter must contain PK', async () => {
+        // @ts-expect-error
+        await client.models.Post.get({});
+      });
+
+      test('parameter must not contain extra fields', async () => {
+        // @ts-expect-error
+        await client.models.Post.get({ id: 'some-id', title: 'whatever' });
+      });
     });
 
     describe('list', () => {
@@ -90,6 +87,7 @@ describe('Basic operations', () => {
           },
         });
       });
+
       test('lazy loaded hasMany returns a non-nullable list of non-nullable elements', async () => {
         const { data } = await client.models.Post.get({ id: 'something' });
         const comments = await data.comments();

packages/integration-tests/__tests__/pseudo-client.test-d.ts

@@ -6,6 +6,7 @@ import {
   type HasKey,
   type CustomQueries,
   type CustomMutations,
+  type CustomSubscriptions,
   type ModelTypes,
   type GraphQLFormattedError,
   __modelMeta__,
@@ -24,6 +25,7 @@ type PartialClient<T extends Record<any, any> = never> = ExcludeNeverFields<{
   models: ModelTypes<T>;
   queries: CustomQueries<T>;
   mutations: CustomMutations<T>;
+  subscriptions: CustomSubscriptions<T>;
 }>;
 
 function generateClient<T extends Record<any, any>>() {
@@ -213,6 +215,98 @@ describe('client', () => {
     type test = Expect<Equal<ResponseType, Expected>>;
   });
 
+  test('subscription - no args', async () => {
+    const schema = a
+      .schema({
+        Post: a.model({
+          title: a.string(),
+          liked: a.boolean(),
+        }),
+
+        onCreateOrUpdatePost: a
+          .subscription()
+          .for(a.ref('Post').mutations(['create', 'update']))
+          .returns(a.ref('Post'))
+          .handler(a.handler.function('onCreateOrUpdatePost')),
+      })
+      .authorization([a.allow.public()]);
+
+    type Schema = ClientSchema<typeof schema>;
+    const client = generateClient<Schema>();
+
+    client.subscriptions.onCreateOrUpdatePost().subscribe({
+      next: (data) => {
+        type ResponseType = Prettify<typeof data>;
+        type ExpectedType = {
+          readonly id: string;
+          readonly createdAt: string;
+          readonly updatedAt: string;
+          title?: string | null;
+          liked?: boolean | null;
+        } | null;
+        type test = Expect<Equal<ResponseType, ExpectedType>>;
+      },
+      error: (error) => {},
+    });
+
+    // Accepts options
+    client.subscriptions.onCreateOrUpdatePost({
+      authMode: 'apiKey',
+      authToken: 'abc123',
+      headers: { someHeader: 'someValue' },
+    });
+  });
+
+  test('subscription - with args', async () => {
+    const schema = a
+      .schema({
+        Post: a.model({
+          title: a.string(),
+          liked: a.boolean(),
+        }),
+
+        onCreateOrUpdatePost: a
+          .subscription()
+          .for(a.ref('Post').mutations(['create', 'update']))
+          .arguments({
+            postId: a.string().required(),
+          })
+          .returns(a.ref('Post'))
+          .handler(a.handler.function('onCreateOrUpdatePost')),
+      })
+      .authorization([a.allow.public()]);
+
+    type Schema = ClientSchema<typeof schema>;
+    const client = generateClient<Schema>();
+
+    const sub = client.subscriptions
+      .onCreateOrUpdatePost({ postId: 'abc' })
+      .subscribe({
+        next: (data) => {
+          type ResponseType = Prettify<typeof data>;
+          type ExpectedType = {
+            readonly id: string;
+            readonly createdAt: string;
+            readonly updatedAt: string;
+            title?: string | null;
+            liked?: boolean | null;
+          } | null;
+          type test = Expect<Equal<ResponseType, ExpectedType>>;
+        },
+        error: (error) => {},
+      });
+
+    // Accepts options
+    client.subscriptions.onCreateOrUpdatePost(
+      { postId: 'abc' },
+      {
+        authMode: 'apiKey',
+        authToken: 'abc123',
+        headers: { someHeader: 'someValue' },
+      },
+    );
+  });
+
   test('custom ops are absent from models property', async () => {
     const schema = a
       .schema({