@aws-amplify/data-schema 1.5.0 → 1.6.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aws-amplify/data-schema",
-  "version": "1.5.0",
+  "version": "1.6.0",
   "license": "Apache-2.0",
   "repository": {
     "type": "git",

package/src/ClientSchema/Core/ClientCustomOperations.ts

@@ -1,5 +1,8 @@
-import type { CustomOperationParamShape } from '../../CustomOperation';
-import type { BaseModelField } from '../../ModelField';
+import type {
+  CustomOperationParamShape,
+  UltimateFunctionHandlerAsyncType,
+} from '../../CustomOperation';
+import type { BaseModelField, ModelFieldType } from '../../ModelField';
 import type { RefType } from '../../RefType';
 import type { ResolveFieldRequirements } from '../../MappedTypes/ResolveFieldProperties';
 import type { AppSyncResolverHandler } from 'aws-lambda';
@@ -36,7 +39,13 @@ export interface ClientCustomOperation<
    */
   functionHandler: AppSyncResolverHandler<
     CustomOpArguments<Op>,
-    LambdaReturnType<CustomOpReturnType<Op, RefBag>>
+    // If the final handler is an async function, the Schema['fieldname']['functionhandler']
+    // should have a return type of `void`. This only applies to `functionHandler` and not
+    // `returnType` because `returnType` determines the type returned by the mutation / query
+    // in the data client.
+    Op['handlers'] extends UltimateFunctionHandlerAsyncType
+      ? void
+      : LambdaReturnType<CustomOpReturnType<Op, RefBag>>
   >;
 
   /**
@@ -100,6 +109,26 @@ type Normalize<
   ? Exclude<RT, null | undefined>
   : RT;
 
+type _EventInvocationResponseBag = {
+  EventInvocationResponse: {
+    __entityType: 'customType';
+    type: '';
+    data: {
+      fields: {
+        success: {
+          data: {
+            fieldType: ModelFieldType.Boolean;
+            required: true;
+            array: false;
+            arrayRequired: false;
+          };
+        };
+      };
+      type: 'customType';
+    };
+  };
+};
+
 /**
  * Computes the return type from the `returnType` of a custom operation shape.
  *

package/src/CustomOperation.ts

@@ -1,6 +1,6 @@
 import { SetTypeSubArg } from '@aws-amplify/data-schema-types';
 import { Brand, brand } from './util';
-import { InternalField, type BaseModelField } from './ModelField';
+import { InternalField, ModelField, type BaseModelField } from './ModelField';
 import {
   AllowModifierForCustomOperation,
   Authorization,
@@ -10,10 +10,12 @@ import { RefType, InternalRef } from './RefType';
 import { EnumType } from './EnumType';
 import { CustomType } from './CustomType';
 import type {
+  AsyncFunctionHandler,
   CustomHandler,
   FunctionHandler,
   HandlerType as Handler,
 } from './Handler';
+import { brandSymbol } from './util/Brand';
 import { AiModel, InferenceConfiguration } from './ai/ModelType';
 
 const queryBrand = 'queryCustomOperation';
@@ -28,14 +30,31 @@ type CustomOperationBrand =
   | typeof generationBrand;
 
 type CustomArguments = Record<string, BaseModelField | EnumType>;
-
 type SubscriptionSource = RefType<any, any>;
 type InternalSubscriptionSource = InternalRef;
-
 type CustomReturnType = RefType<any> | CustomType<any>;
 type InternalCustomArguments = Record<string, InternalField>;
 type InternalCustomReturnType = InternalRef;
-type HandlerInputType = FunctionHandler[] | CustomHandler[] | Handler;
+type HandlerInputType =
+  | FunctionHandler[]
+  | CustomHandler[]
+  | AsyncFunctionHandler[]
+  | HeterogeneousFunctionHandlerWithLastAsync
+  | HeterogeneousFunctionHandlerType
+  | Handler;
+type HeterogeneousFunctionHandlerType = (
+  | FunctionHandler
+  | AsyncFunctionHandler
+)[];
+type HeterogeneousFunctionHandlerWithLastAsync = [
+  ...HeterogeneousFunctionHandlerType,
+  AsyncFunctionHandler,
+];
+
+export type UltimateFunctionHandlerAsyncType =
+  | AsyncFunctionHandler
+  | AsyncFunctionHandler[]
+  | HeterogeneousFunctionHandlerWithLastAsync;
 
 export const CustomOperationNames = [
   'Query',
@@ -104,7 +123,14 @@ export type CustomOperation<
     >;
     handler<H extends HandlerInputType>(
       handlers: H,
-    ): CustomOperation<T, K | 'handler', B>;
+    ): [H] extends [UltimateFunctionHandlerAsyncType]
+      ? CustomOperation<
+          AsyncFunctionCustomOperation<T>,
+          K | 'handler' | 'returns',
+          B
+        >
+      : CustomOperation<T, K | 'handler', B>;
+
     for<Source extends SubscriptionSource>(
       source: Source | Source[],
     ): CustomOperation<
@@ -183,6 +209,10 @@ function _custom<
           ? handlers
           : ([handlers] as Handler[]);
 
+        if (lastHandlerIsAsyncFunction(handlers)) {
+          data.returnType = eventInvocationResponse;
+        }
+
         return this;
       },
       for(source: SubscriptionSource | SubscriptionSource[]) {
@@ -307,6 +337,39 @@ export function subscription(): CustomOperation<
   return _custom('Subscription', subscriptionBrand);
 }
 
+// #region async Lambda function related types
+type AsyncFunctionCustomOperation<T extends CustomOperationParamShape> =
+  SetTypeSubArg<
+    SetTypeSubArg<T, 'returnType', EventInvocationResponseCustomType>,
+    'handlers',
+    AsyncFunctionHandler
+  >;
+
+type EventInvocationResponseCustomType = CustomType<{
+  fields: {
+    success: ModelField<boolean, 'required', undefined>;
+  };
+}>;
+
+const eventInvocationResponse = {
+  data: {
+    type: 'ref',
+    link: 'EventInvocationResponse',
+    valueRequired: false,
+    array: false,
+    arrayRequired: false,
+    mutationOperations: [],
+    authorization: [],
+  },
+};
+
+function lastHandlerIsAsyncFunction(handlers: HandlerInputType): boolean {
+  const lastHandlerBrandSymbol = Array.isArray(handlers)
+    ? handlers[handlers.length - 1][brandSymbol]
+    : handlers[brandSymbol];
+  return lastHandlerBrandSymbol === 'asyncFunctionHandler';
+}
+// #endregion async Lambda function related types
 export interface GenerationInput {
   aiModel: AiModel;
   systemPrompt: string;

package/src/Handler.ts

@@ -1,12 +1,14 @@
 import type { DefineFunction } from '@aws-amplify/data-schema-types';
 import { Brand, brand } from './util';
+import type { brandSymbol } from './util/Brand.js';
 import { RefType } from './RefType';
 
 export type HandlerType =
   | InlineSqlHandler
   | SqlReferenceHandler
   | CustomHandler
-  | FunctionHandler;
+  | FunctionHandler
+  | AsyncFunctionHandler;
 
 const dataSymbol = Symbol('Data');
 
@@ -18,7 +20,8 @@ type AllHandlers =
   | InlineSqlHandler
   | SqlReferenceHandler
   | CustomHandler
-  | FunctionHandler;
+  | FunctionHandler
+  | AsyncFunctionHandler;
 
 export function getHandlerData<H extends AllHandlers>(
   handler: H,
@@ -30,9 +33,10 @@ export function getHandlerData<H extends AllHandlers>(
 
 const inlineSqlBrand = 'inlineSql';
 
-export type InlineSqlHandler = { [dataSymbol]: string } & Brand<
-  typeof inlineSqlBrand
->;
+export type InlineSqlHandler = {
+  [dataSymbol]: string;
+  [brandSymbol]: typeof inlineSqlBrand;
+};
 
 function inlineSql(sql: string): InlineSqlHandler {
   return { [dataSymbol]: sql, ...buildHandler(inlineSqlBrand) };
@@ -51,7 +55,8 @@ export type SqlReferenceHandlerData = {
 
 export type SqlReferenceHandler = {
   [dataSymbol]: SqlReferenceHandlerData;
-} & Brand<typeof sqlReferenceBrand>;
+  [brandSymbol]: typeof sqlReferenceBrand;
+};
 
 function sqlReference(sqlFilePath: string): SqlReferenceHandler {
   // used to determine caller directory in order to resolve relative path downstream
@@ -88,12 +93,13 @@ export type CustomHandlerData = CustomHandlerInput & {
 
 const customHandlerBrand = 'customHandler';
 
-export type CustomHandler = { [dataSymbol]: CustomHandlerData } & Brand<
-  typeof customHandlerBrand
->;
+export type CustomHandler = {
+  [dataSymbol]: CustomHandlerData;
+  [brandSymbol]: typeof customHandlerBrand;
+};
 
 /**
- * Use a custom JavaScript resolver to handle a query, mutation, or subscription. 
+ * Use a custom JavaScript resolver to handle a query, mutation, or subscription.
  * @see {@link https://docs.amplify.aws/react/build-a-backend/data/custom-business-logic/#step-2---configure-custom-business-logic-handler-code}
  * @param customHandler `{ entry: "path-to-javascript-resolver-file.js", dataSource: "Data Source name added via "backend.data.add*DataSoruce(...)"}`
  * @returns A JavaScript resolver attached to the query, mutation, or subscription.
@@ -107,7 +113,7 @@ export type CustomHandler = { [dataSymbol]: CustomHandlerData } & Brand<
  *     allow.owner(),
  *     allow.authenticated().to(['read'])
  *   ]),
- * 
+ *
  *   likePost: a
  *     .mutation()
  *     .arguments({ postId: a.id() })
@@ -131,17 +137,63 @@ function custom(customHandler: CustomHandlerInput): CustomHandler {
 //#endregion
 
 //#region handler.function
-
-export type FunctionHandlerData = DefineFunction | string;
+export type FunctionHandlerInput = DefineFunction | string;
+export type FunctionHandlerData = {
+  handler: FunctionHandlerInput;
+  invocationType: 'RequestResponse' | 'Event';
+};
 
 const functionHandlerBrand = 'functionHandler';
 
 export type FunctionHandler = {
   [dataSymbol]: FunctionHandlerData;
-} & Brand<typeof functionHandlerBrand>;
+  /**
+   * Use `async()` to have this function handler invoked asynchronously.
+   * If an `async()` function is only handler or the final handler in a pipeline, the return type of this
+   * custom query / mutation is `{ success: boolean }`.
+   * @example
+   * import {
+   *   type ClientSchema,
+   *   a,
+   *   defineData,
+   *   defineFunction // 1.Import "defineFunction" to create new functions
+   * } from '@aws-amplify/backend';
+   *
+   * // 2. define a function
+   * const echoHandler = defineFunction({
+   *   entry: './echo-handler/handler.ts'
+   * })
+   *
+   * const schema = a.schema({
+   *   EchoResponse: a.customType({
+   *     content: a.string(),
+   *     executionDuration: a.float()
+   *   }),
+   *
+   *   echo: a
+   *     .query()
+   *     .arguments({ content: a.string() })
+   *     .returns(a.ref('EchoResponse'))
+   *     .authorization(allow => [allow.publicApiKey()])
+   *     // 3. set the function has the handler
+   *     .handler(a.handler.function(echoHandler).async())
+   *
+   * @see {@link https://docs.amplify.aws/react/build-a-backend/data/INSERT_PATH_WHEN_READY}
+   * @returns A function handler for query / mutation that is asynchronously invoked.
+   */
+  async(): AsyncFunctionHandler;
+  [brandSymbol]: typeof functionHandlerBrand;
+};
+
+const asyncFunctionHandlerBrand = 'asyncFunctionHandler';
+
+export type AsyncFunctionHandler = {
+  [dataSymbol]: FunctionHandlerData;
+  [brandSymbol]: typeof asyncFunctionHandlerBrand;
+};
 
 /**
- * Use a function created via `defineFunction` to handle the custom query/mutation/subscription. In your function handler, 
+ * Use a function created via `defineFunction` to handle the custom query/mutation/subscription. In your function handler,
  * you can use the `Schema["YOUR_QUERY_OR_MUTATION_NAME"]["functionHandler"]` utility type to type the handler function.
  * @example
  * import {
@@ -150,18 +202,18 @@ export type FunctionHandler = {
  *   defineData,
  *   defineFunction // 1.Import "defineFunction" to create new functions
  * } from '@aws-amplify/backend';
- * 
+ *
  * // 2. define a function
  * const echoHandler = defineFunction({
  *   entry: './echo-handler/handler.ts'
  * })
- * 
+ *
  * const schema = a.schema({
  *   EchoResponse: a.customType({
  *     content: a.string(),
  *     executionDuration: a.float()
  *   }),
- * 
+ *
  *   echo: a
  *     .query()
  *     .arguments({ content: a.string() })
@@ -171,12 +223,31 @@ export type FunctionHandler = {
  *     .handler(a.handler.function(echoHandler))
  * });
  * @see {@link https://docs.amplify.aws/react/build-a-backend/data/custom-business-logic/}
- * @param fn A function created via `defineFunction`. Alternatively, you can pass in a "string" of the function name and pass 
+ * @param fn A function created via `defineFunction`. Alternatively, you can pass in a "string" of the function name and pass
  * in a corresponding value into the `functionMap` property of defineData.
  * @returns A handler for the query / mutation / subscription
  */
-function fcn(fn: FunctionHandlerData): FunctionHandler {
-  return { [dataSymbol]: fn, ...buildHandler(functionHandlerBrand) };
+function fcn(fn: FunctionHandlerInput): FunctionHandler {
+  return {
+    [dataSymbol]: {
+      handler: fn,
+      invocationType: 'RequestResponse',
+    },
+    async() {
+      return _async(this);
+    },
+    ...buildHandler(functionHandlerBrand),
+  };
+}
+
+function _async(fnHandler: FunctionHandler): AsyncFunctionHandler {
+  return {
+    [dataSymbol]: {
+      handler: fnHandler[dataSymbol].handler,
+      invocationType: 'Event',
+    },
+    ...buildHandler(asyncFunctionHandlerBrand),
+  };
 }
 
 //#endregion

package/src/SchemaProcessor.ts

@@ -42,10 +42,12 @@ import {
   type HandlerType,
   type CustomHandler,
   type SqlReferenceHandler,
-  type FunctionHandler,
+  FunctionHandler,
+  AsyncFunctionHandler
 } from './Handler';
 import * as os from 'os';
 import * as path from 'path';
+import { brandSymbol } from './util/Brand';
 import {
   brandName as conversationBrandName,
   type InternalConversationType,
@@ -279,7 +281,7 @@ function enumFieldToGql(enumName: string, secondaryIndexes: string[] = []) {
 }
 
 function transformFunctionHandler(
-  handlers: FunctionHandler[],
+  handlers: (FunctionHandler | AsyncFunctionHandler)[],
   functionFieldName: string,
 ): {
   gqlHandlerContent: string;
@@ -291,21 +293,25 @@ function transformFunctionHandler(
   handlers.forEach((handler, idx) => {
     const handlerData = getHandlerData(handler);
 
-    if (typeof handlerData === 'string') {
-      gqlHandlerContent += `@function(name: "${handlerData}") `;
-    } else if (typeof handlerData.getInstance === 'function') {
+    if (typeof handlerData.handler === 'string') {
+      gqlHandlerContent += `@function(name: "${handlerData.handler}") `;
+    } else if (typeof handlerData.handler.getInstance === 'function') {
       const fnName = `Fn${capitalize(functionFieldName)}${idx === 0 ? '' : `${idx + 1}`}`;
-
-      lambdaFunctionDefinition[fnName] = handlerData;
-      gqlHandlerContent += `@function(name: "${fnName}") `;
-    } else {
+      lambdaFunctionDefinition[fnName] = handlerData.handler;
+      const invocationTypeArg = handlerData.invocationType === 'Event' ? ', invocationType: Event)' : ')'
+      gqlHandlerContent += `@function(name: "${fnName}"${invocationTypeArg} `;
+    }
+    else {
       throw new Error(
         `Invalid value specified for ${functionFieldName} handler.function(). Expected: defineFunction or string.`,
       );
     }
   });
 
-  return { gqlHandlerContent, lambdaFunctionDefinition };
+  return {
+    gqlHandlerContent,
+    lambdaFunctionDefinition,
+  };
 }
 
 type CustomTypeAuthRules =
@@ -1275,6 +1281,20 @@ const schemaPreprocessor = (
 
   const staticSchema = databaseType === 'sql';
 
+  // If the schema contains a custom operation with an async lambda handler,
+  // we need to add the EventInvocationResponse custom type to the schema.
+  // This is done here so that:
+  // - it only happens once per schema
+  // - downstream validation based on `getRefTypeForSchema` finds the EventInvocationResponse type
+  const containsAsyncLambdaCustomOperation = Object.entries(schema.data.types).find(([_, typeDef]) => {
+    return isCustomOperation(typeDef)
+    && finalHandlerIsAsyncFunctionHandler(typeDef.data.handlers);
+  });
+
+  if (containsAsyncLambdaCustomOperation) {
+    schema.data.types['EventInvocationResponse'] = eventInvocationResponseCustomType;
+  }
+
   const topLevelTypes = sortTopLevelTypes(Object.entries(schema.data.types));
 
   const { schemaAuth, functionSchemaAccess } = extractFunctionSchemaAccess(
@@ -1575,12 +1595,16 @@ function validateCustomOperations(
     }
 
     if (configuredHandlers.size > 1) {
-      const configuredHandlersStr = JSON.stringify(
-        Array.from(configuredHandlers),
-      );
-      throw new Error(
-        `Field handlers must be of the same type. ${typeName} has been configured with ${configuredHandlersStr}`,
-      );
+      configuredHandlers.delete('asyncFunctionHandler');
+      configuredHandlers.delete('functionHandler');
+      if (configuredHandlers.size > 0) {
+        const configuredHandlersStr = JSON.stringify(
+          Array.from(configuredHandlers),
+        );
+        throw new Error(
+          `Field handlers must be of the same type. ${typeName} has been configured with ${configuredHandlersStr}`,
+        );
+      }
     }
   }
 
@@ -1588,11 +1612,15 @@ function validateCustomOperations(
     typeDef.data.returnType === null &&
     (opType === 'Query' || opType === 'Mutation' || opType === 'Generation')
   ) {
-    const typeDescription =
-      opType === 'Generation' ? 'Generation Route' : `Custom ${opType}`;
-    throw new Error(
-      `Invalid ${typeDescription} definition. A ${typeDescription} must include a return type. ${typeName} has no return type specified.`,
-    );
+    // TODO: There should be a more elegant and readable way to handle this check.
+    // Maybe it's not even necessary anymore since we're the setting returnType in the handler() method.
+    if (!handlers || handlers.length === 0 || handlers[handlers.length - 1][brandSymbol] !== 'asyncFunctionHandler') {
+      const typeDescription =
+        opType === 'Generation' ? 'Generation Route' : `Custom ${opType}`;
+      throw new Error(
+        `Invalid ${typeDescription} definition. A ${typeDescription} must include a return type. ${typeName} has no return type specified.`,
+      );
+    }
   }
 
   if (opType !== 'Subscription' && subscriptionSource.length > 0) {
@@ -1677,10 +1705,16 @@ const isCustomHandler = (
 
 const isFunctionHandler = (
   handler: HandlerType[] | null,
-): handler is FunctionHandler[] => {
-  return Array.isArray(handler) && getBrand(handler[0]) === 'functionHandler';
+): handler is (FunctionHandler | AsyncFunctionHandler)[] => {
+  return Array.isArray(handler) && ['functionHandler', 'asyncFunctionHandler'].includes(getBrand(handler[0]));
 };
 
+const finalHandlerIsAsyncFunctionHandler = (
+  handler: HandlerType[] | null,
+): handler is AsyncFunctionHandler[] => {
+  return Array.isArray(handler) && getBrand(handler[handler.length - 1]) === 'asyncFunctionHandler';
+}
+
 const normalizeDataSourceName = (
   dataSource: undefined | string | RefType<any, any, any>,
 ): string => {
@@ -1762,6 +1796,22 @@ const handleCustom = (
   return jsFn;
 };
 
+const eventInvocationResponseCustomType = {
+  data: {
+    fields: {
+      success: {
+        data: {
+          fieldType: ModelFieldType.Boolean,
+          required: true,
+          array: false,
+          arrayRequired: false,
+        }
+      }
+    },
+    type: 'customType'
+  }
+};
+
 function transformCustomOperations(
   typeDef: InternalCustom,
   typeName: string,

package/src/runtime/internals/operations/subscription.ts

@@ -17,6 +17,7 @@ import {
   generateGraphQLDocument,
   getCustomHeaders,
   initializeModel,
+  flattenItems,
 } from '../APIClient';
 
 export function subscriptionFactory(
@@ -59,18 +60,25 @@ export function subscriptionFactory(
     return observable.pipe(
       map((value) => {
         const [key] = Object.keys(value.data);
-        // Will need flattening here if/when custom selection set support is added:
         const data = (value.data as any)[key];
-        const [initialized] = initializeModel(
-          client,
-          name,
-          [data],
-          modelIntrospection,
-          auth.authMode,
-          auth.authToken,
-        );
+        const flattenedResult = flattenItems(modelIntrospection, name, data);
 
-        return initialized;
+        if (flattenedResult === null) {
+          return null;
+        } else if (args?.selectionSet) {
+          return flattenedResult;
+        } else {
+          const [initialized] = initializeModel(
+            client,
+            name,
+            [flattenedResult],
+            modelIntrospection,
+            auth.authMode,
+            auth.authToken,
+          );
+
+          return initialized;
+        }
       }),
     );
   };