@aws-amplify/data-schema 0.13.16 → 0.13.18

package/lib-esm/src/ClientSchema.d.ts

@@ -1,34 +1,29 @@
 import { type __modelMeta__ } from '@aws-amplify/data-schema-types';
-import type { ModelSchema } from './ModelSchema';
+import { type GenericModelSchema, type RDSModelSchema } from './ModelSchema';
 import type { ResolveSchema, SchemaTypes } from './MappedTypes/ResolveSchema';
-import type { CreateImplicitModelsFromRelations, ResolveFieldProperties } from './MappedTypes/ResolveFieldProperties';
+import type { CreateImplicitModelsFromRelations, ResolveFieldProperties, ResolveStaticFieldProperties } from './MappedTypes/ResolveFieldProperties';
 import type { ModelIdentifier, ModelSecondaryIndexes, RelationalMetadata } from './MappedTypes/ModelMetadata';
 import type { ExtractNonModelTypes, NonModelTypesShape } from './MappedTypes/ExtractNonModelTypes';
 import { ResolveCustomOperations, CustomOperationHandlerTypes } from './MappedTypes/CustomOperations';
-export type ClientSchema<Schema extends ModelSchema<any, any>> = InternalClientSchema<Schema>;
+export type ClientSchema<Schema extends GenericModelSchema<any>> = InternalClientSchema<Schema>;
 /**
  * Types for unwrapping generic type args into client-consumable types
  *
- * @typeParam Schema - Type Beast schema type
+ * @typeParam Schema - Data schema builder model type
  *
  * The following params are used solely as variables in order to simplify mapped type usage.
  * They should not receive external type args.
  *
  * @internal @typeParam NonModelTypes - Custom Types, Enums, and Custom Operations
- * @internal @typeParam ResolvedSchema - Schema/Models/Fields structure with generic type args extracted
- * @internal @typeParam ResolvedFields - Resovled client-facing types used for CRUDL response shapes
- * @internal @typeParam IdentifierMeta - Map of model primary index metadata
+ * @internal @typeParam ImplicitModels - The implicit models created to represent relationships
+ * @internal @typeParam ResolvedFields - Resolved client-facing types used for CRUDL response shapes
  * @internal @typeParam SecondaryIndexes - Map of model secondary index metadata
- *
- * @internal @typeParam Meta - Stores schema metadata: identifier, relationship metadata;
- * used by `API.generateClient` to craft strongly typed mutation inputs; hidden from customer-facing types behind __modelMeta__ symbol
- *
  */
-type InternalClientSchema<Schema extends ModelSchema<any, any>, NonModelTypes extends NonModelTypesShape = ExtractNonModelTypes<Schema>, ResolvedSchema = ResolveSchema<Schema>, ImplicitModels = CreateImplicitModelsFromRelations<ResolvedSchema>, ImplicitModelsIdentifierMeta = {
+type InternalClientSchema<Schema extends GenericModelSchema<any>, NonModelTypes extends NonModelTypesShape = ExtractNonModelTypes<Schema>, ResolvedSchema = ResolveSchema<Schema>, ImplicitModels = Schema extends RDSModelSchema<any, any> ? object : CreateImplicitModelsFromRelations<ResolvedSchema>, ImplicitModelsIdentifierMeta = {
     [ImplicitModel in keyof ImplicitModels]: {
         identifier: 'id';
     };
-}, ResolvedFields extends Record<string, unknown> = ResolveFieldProperties<Schema, NonModelTypes, ImplicitModels>, IdentifierMeta extends Record<string, any> = ModelIdentifier<SchemaTypes<Schema>>, SecondaryIndexes extends Record<string, any> = ModelSecondaryIndexes<SchemaTypes<Schema>>> = ResolvedFields & CustomOperationHandlerTypes<ResolveCustomOperations<Schema, ResolvedFields, NonModelTypes>['customOperations']> & ResolvedFields & {
+}, ResolvedFields extends Record<string, unknown> = Schema extends RDSModelSchema<any, any> ? ResolveStaticFieldProperties<Schema, NonModelTypes, object> : ResolveFieldProperties<Schema, NonModelTypes, ImplicitModels>, IdentifierMeta extends Record<string, any> = ModelIdentifier<SchemaTypes<Schema>>, SecondaryIndexes extends Record<string, any> = Schema extends RDSModelSchema<any, any> ? object : ModelSecondaryIndexes<SchemaTypes<Schema>>> = CustomOperationHandlerTypes<ResolveCustomOperations<Schema, ResolvedFields, NonModelTypes>['customOperations']> & ResolvedFields & {
     [__modelMeta__]: IdentifierMeta & ImplicitModelsIdentifierMeta & SecondaryIndexes & RelationalMetadata<ResolvedSchema, ResolvedFields, IdentifierMeta> & NonModelTypes & ResolveCustomOperations<Schema, ResolvedFields, NonModelTypes>;
 };
 export {};

package/lib-esm/src/CustomOperation.d.ts

@@ -11,6 +11,8 @@ declare const mutationBrand = "mutationCustomOperation";
 declare const subscriptionBrand = "subscriptionCustomOperation";
 type CustomOperationBrand = typeof queryBrand | typeof mutationBrand | typeof subscriptionBrand;
 type CustomArguments = Record<string, ModelField<any, any> | EnumType<EnumTypeParamShape>>;
+type SubscriptionSource = RefType<any, any>;
+type InternalSubscriptionSource = InternalRef;
 type CustomReturnType = RefType<any> | CustomType<any>;
 type CustomFunctionRefType = string;
 type InternalCustomArguments = Record<string, InternalField>;
@@ -25,11 +27,13 @@ type CustomData = {
     authorization: Authorization<any, any, any>[];
     typeName: CustomOperationName;
     handlers: Handler[] | null;
+    subscriptionSource: SubscriptionSource[];
 };
 type InternalCustomData = CustomData & {
     arguments: InternalCustomArguments;
     returnType: InternalCustomReturnType;
     functionRef: string | null;
+    subscriptionSource: InternalSubscriptionSource[];
     authorization: Authorization<any, any, any>[];
 };
 export type CustomOperationParamShape = {
@@ -54,6 +58,7 @@ export type CustomOperation<T extends CustomOperationParamShape, K extends keyof
     function<FunctionRef extends CustomFunctionRefType>(functionRefOrName: FunctionRef): CustomOperation<SetTypeSubArg<T, 'functionRef', FunctionRef>, K | 'function', B>;
     authorization<AuthRuleType extends Authorization<any, any, any>>(rules: AuthRuleType[]): CustomOperation<SetTypeSubArg<T, 'authorization', AuthRuleType[]>, K | 'authorization', B>;
     handler<H extends HandlerInputType>(handlers: H): CustomOperation<T, K | 'handler', B>;
+    for(source: SubscriptionSource | SubscriptionSource[]): CustomOperation<T, K | 'for', B>;
 }, K> & Brand<B>;
 /**
  * Internal representation of Custom Type that exposes the `data` property.

package/lib-esm/src/CustomOperation.js

@@ -21,6 +21,7 @@ function _custom(typeName, brand) {
         authorization: [],
         typeName: typeName,
         handlers: null,
+        subscriptionSource: [],
     };
     const builder = brandedBuilder({
         arguments(args) {
@@ -45,6 +46,10 @@ function _custom(typeName, brand) {
                 : [handlers];
             return this;
         },
+        for(source) {
+            data.subscriptionSource = Array.isArray(source) ? source : [source];
+            return this;
+        },
     }, brand);
     return { ...builder, data };
 }

package/lib-esm/src/Handler.d.ts

@@ -23,7 +23,7 @@ type CustomHandlerInput = {
      * Defaults to 'NONE_DS'
      *
      */
-    dataSource?: string | RefType<any, any, any>;
+    dataSource?: string | RefType<any>;
     /**
      * The path to the file that contains the function entry point.
      * If this is a relative path, it is computed relative to the file where this handler is defined

package/lib-esm/src/MappedTypes/CustomOperations.d.ts

@@ -1,4 +1,4 @@
-import type { ModelSchema } from '../ModelSchema';
+import type { GenericModelSchema } from '../ModelSchema';
 import type { NonModelTypesShape } from './ExtractNonModelTypes';
 import type { CustomOperation, CustomOperationParamShape } from '../CustomOperation';
 import type { ModelField } from '../ModelField';
@@ -8,7 +8,7 @@ import type { AppSyncResolverHandler } from 'aws-lambda';
 /**
  * Creates meta types for custom operations from a schema.
  */
-export type ResolveCustomOperations<Schema extends ModelSchema<any, any>, FullyResolvedSchema extends Record<string, unknown>, NonModelTypes extends NonModelTypesShape> = {
+export type ResolveCustomOperations<Schema extends GenericModelSchema<any>, FullyResolvedSchema extends Record<string, unknown>, NonModelTypes extends NonModelTypesShape> = {
     customOperations: {
         [OpName in keyof CustomOpShapes<Schema>]: {
             arguments: CustomOpArguments<CustomOpShapes<Schema>[OpName]>;
@@ -22,7 +22,7 @@ export type ResolveCustomOperations<Schema extends ModelSchema<any, any>, FullyR
 /**
  * Filtered, mapped list of custom operations shapes from a schema.
  */
-export type CustomOpShapes<Schema extends ModelSchema<any, any>> = {
+export type CustomOpShapes<Schema extends GenericModelSchema<any>> = {
     [K in keyof Schema['data']['types'] as Schema['data']['types'][K] extends CustomOperation<any, any> ? K : never]: Schema['data']['types'][K] extends CustomOperation<infer Shape, any> ? Shape : never;
 };
 /**

package/lib-esm/src/MappedTypes/ResolveFieldProperties.d.ts

@@ -2,7 +2,7 @@ import type { UnionToIntersection, LazyLoader, ExcludeEmpty } from '@aws-amplify
 import type { Authorization, ImpliedAuthFields } from '../Authorization';
 import type { ModelField } from '../ModelField';
 import type { ModelType, ModelTypeParamShape } from '../ModelType';
-import type { ModelSchema } from '../ModelSchema';
+import type { GenericModelSchema } from '../ModelSchema';
 import type { ModelRelationalField, ModelRelationalFieldParamShape, ModelRelationalTypeArgFactory, ModelRelationshipTypes } from '../ModelRelationalField';
 import type { AllImpliedFKs } from './ForeignKeys';
 import type { ResolveSchema, SchemaTypes } from './ResolveSchema';
@@ -13,9 +13,10 @@ import type { NonModelTypesShape } from './ExtractNonModelTypes';
 import type { CustomType, CustomTypeParamShape } from '../CustomType';
 import type { EnumType, EnumTypeParamShape } from '../EnumType';
 import type { CustomOperation, CustomOperationParamShape } from '../CustomOperation';
-export type ResolveFieldProperties<Schema extends ModelSchema<any, any>, NonModelTypes extends NonModelTypesShape, ImplicitModelsSchema, ResolvedSchema = ResolveSchema<Schema>, IdentifierMeta extends Record<string, {
+export type ResolveFieldProperties<Schema extends GenericModelSchema<any>, NonModelTypes extends NonModelTypesShape, ImplicitModelsSchema, ResolvedSchema = ResolveSchema<Schema>, IdentifierMeta extends Record<string, {
     identifier: string;
 }> = ModelIdentifier<SchemaTypes<Schema>>, FieldsWithInjectedImplicitFields = InjectImplicitModelFields<ResolvedSchema & ImplicitModelsSchema, IdentifierMeta>, FieldsWithRelationships = ResolveRelationships<FieldsWithInjectedImplicitFields, NonModelTypes>> = Intersection<FilterFieldTypes<MarkModelsNonNullableFieldsRequired<FieldsWithRelationships>>, FilterFieldTypes<MarkModelsNullableFieldsOptional<FieldsWithRelationships>>, FilterFieldTypes<ModelImpliedAuthFields<Schema>>, AllImpliedFKs<ResolvedSchema, IdentifierMeta>>;
+export type ResolveStaticFieldProperties<Schema extends GenericModelSchema<any>, NonModelTypes extends NonModelTypesShape, ImplicitModelsSchema, ResolvedSchema = ResolveSchema<Schema>, FieldsWithInjectedImplicitFields = InjectImplicitModelFields<ResolvedSchema & ImplicitModelsSchema, object>, FieldsWithRelationships = ResolveRelationships<FieldsWithInjectedImplicitFields, NonModelTypes>> = Intersection<FilterFieldTypes<MarkModelsNonNullableFieldsRequired<FieldsWithRelationships>>, FilterFieldTypes<MarkModelsNullableFieldsOptional<FieldsWithRelationships>>>;
 export type CreateImplicitModelsFromRelations<Schema> = UnionToIntersection<ExcludeEmpty<{
     [ModelProp in keyof Schema]: {
         [FieldProp in keyof Schema[ModelProp] as Schema[ModelProp][FieldProp] extends ModelRelationalFieldParamShape ? Schema[ModelProp][FieldProp]['relationName'] extends string ? Schema[ModelProp][FieldProp]['relationName'] extends keyof Schema ? never : Schema[ModelProp][FieldProp]['relationName'] : never : never]: Record<`${Lowercase<ModelProp & string>}`, ModelRelationalTypeArgFactory<ModelProp & string, ModelRelationshipTypes.hasMany, false>>;
@@ -59,9 +60,9 @@ type MarkModelsNonNullableFieldsRequired<Schema> = {
 type Intersection<A = Record<never, never>, B = Record<never, never>, C = Record<never, never>, D = Record<never, never>> = A & B & C & D extends infer U ? {
     [P in keyof U]: U[P];
 } : never;
-export type ModelImpliedAuthFields<Schema extends ModelSchema<any, any>> = {
+export type ModelImpliedAuthFields<Schema extends GenericModelSchema<any>> = {
     [ModelKey in keyof Schema['data']['types'] as Schema['data']['types'][ModelKey] extends EnumType<EnumTypeParamShape> ? never : Schema['data']['types'][ModelKey] extends CustomType<CustomTypeParamShape> ? never : Schema['data']['types'][ModelKey] extends CustomOperation<CustomOperationParamShape, any> ? never : ModelKey]: Schema['data']['types'][ModelKey] extends ModelType<infer Model, any> ? AllAuthFieldsForModel<Schema, Model> : object;
 };
-type AllAuthFieldsForModel<Schema extends ModelSchema<any, any>, Model extends Schema['data']['types'][keyof Schema['data']['types']]> = (Model['authorization'][number] extends never ? Schema['data']['authorization'][number] extends never ? object : ImpliedAuthFields<Schema['data']['authorization'][number]> : ImpliedAuthFields<Model['authorization'][number]>) & ImpliedAuthFieldsFromFields<Model>;
+type AllAuthFieldsForModel<Schema extends GenericModelSchema<any>, Model extends Schema['data']['types'][keyof Schema['data']['types']]> = (Model['authorization'][number] extends never ? Schema['data']['authorization'][number] extends never ? object : ImpliedAuthFields<Schema['data']['authorization'][number]> : ImpliedAuthFields<Model['authorization'][number]>) & ImpliedAuthFieldsFromFields<Model>;
 type ImpliedAuthFieldsFromFields<T> = UnionToIntersection<T extends ModelTypeParamShape ? T['fields'][keyof T['fields']] extends ModelField<any, any, infer Auth> | ModelRelationalField<any, any, any, infer Auth> | RefType<any, any, infer Auth> ? Auth extends Authorization<any, any, any> ? ImpliedAuthFields<Auth> : object : object : object>;
 export {};

package/lib-esm/src/MappedTypes/ResolveSchema.d.ts

@@ -1,5 +1,5 @@
 import type { ModelType } from '../ModelType';
-import type { ModelSchema } from '../ModelSchema';
+import type { GenericModelSchema } from '../ModelSchema';
 import type { ModelRelationalField, ModelRelationshipTypes, RelationTypeFunctionOmitMapping } from '../ModelRelationalField';
 import type { ModelField } from '../ModelField';
 import type { CustomType, CustomTypeParamShape } from '../CustomType';
@@ -7,7 +7,7 @@ import type { EnumType, EnumTypeParamShape } from '../EnumType';
 import type { RefType, RefTypeParamShape } from '../RefType';
 import type { CustomOperation, CustomOperationParamShape } from '../CustomOperation';
 export type ResolveSchema<Schema> = FieldTypes<ModelTypes<SchemaTypes<Schema>>>;
-export type SchemaTypes<T> = T extends ModelSchema<any, any> ? T['data']['types'] : never;
+export type SchemaTypes<T> = T extends GenericModelSchema<any> ? T['data']['types'] : never;
 /**
  * Resolves model types
  *

package/lib-esm/src/ModelSchema.d.ts

@@ -4,6 +4,12 @@ import type { EnumType, EnumTypeParamShape } from './EnumType';
 import type { CustomType, CustomTypeParamShape } from './CustomType';
 import type { CustomOperation, CustomOperationParamShape, InternalCustom, MutationCustomOperation, QueryCustomOperation, SubscriptionCustomOperation } from './CustomOperation';
 import { SchemaAuthorization } from './Authorization';
+import { Brand } from './util';
+export declare const rdsSchemaBrandName = "RDSSchema";
+export declare const rdsSchemaBrand: Brand<"RDSSchema">;
+export type RDSSchemaBrand = Brand<typeof rdsSchemaBrandName>;
+export declare const ddbSchemaBrandName = "DDBSchema";
+export type DDBSchemaBrand = Brand<typeof ddbSchemaBrandName>;
 type SchemaContent = ModelType<ModelTypeParamShape, any> | CustomType<CustomTypeParamShape> | EnumType<EnumTypeParamShape> | CustomOperation<CustomOperationParamShape, any>;
 type ModelSchemaContents = Record<string, SchemaContent>;
 type InternalSchemaModels = Record<string, InternalModel | EnumType<any> | CustomType<any> | InternalCustom>;
@@ -38,7 +44,7 @@ export type ModelSchemaParamShape = {
     authorization: SchemaAuthorization<any, any, any>[];
     configuration: SchemaConfig<any, any>;
 };
-export type SQLModelSchemaParamShape = ModelSchemaParamShape & {
+export type RDSModelSchemaParamShape = ModelSchemaParamShape & {
     setSqlStatementFolderPath?: string;
 };
 export type InternalSchema = {
@@ -48,22 +54,25 @@ export type InternalSchema = {
         configuration: SchemaConfig<any, any>;
     };
 };
-export type ModelSchema<T extends ModelSchemaParamShape, UsedMethods extends 'authorization' = never> = Omit<{
-    authorization: <AuthRules extends SchemaAuthorization<any, any, any>>(auth: AuthRules[]) => ModelSchema<SetTypeSubArg<T, 'authorization', AuthRules[]>, UsedMethods | 'authorization'>;
-}, UsedMethods> & {
+type BaseSchema<T extends ModelSchemaParamShape> = {
     data: T;
     models: {
         [TypeKey in keyof T['types']]: T['types'][TypeKey] extends ModelType<ModelTypeParamShape> ? SchemaModelType<T['types'][TypeKey]> : never;
     };
     transform: () => DerivedApiDefinition;
 };
-type SQLModelSchemaFunctions = 'setSqlStatementFolderPath' | 'addQueries' | 'addMutations' | 'addSubscriptions';
-export type SQLModelSchema<T extends SQLModelSchemaParamShape, UsedMethods extends 'authorization' | SQLModelSchemaFunctions = never> = ModelSchema<T, Exclude<UsedMethods, SQLModelSchemaFunctions>> & Omit<{
-    setSqlStatementFolderPath: (path: string) => SQLModelSchema<T, UsedMethods | 'setSqlStatementFolderPath'>;
-    addQueries: (types: Record<string, QueryCustomOperation>) => SQLModelSchema<T, UsedMethods | 'addQueries'>;
-    addMutations: (types: Record<string, MutationCustomOperation>) => SQLModelSchema<T, UsedMethods | 'addMutations'>;
-    addSubscriptions: (types: Record<string, SubscriptionCustomOperation>) => SQLModelSchema<T, UsedMethods | 'addSubscriptions'>;
-}, UsedMethods>;
+export type GenericModelSchema<T extends ModelSchemaParamShape> = BaseSchema<T> & Brand<string>;
+export type ModelSchema<T extends ModelSchemaParamShape, UsedMethods extends 'authorization' = never> = Omit<{
+    authorization: <AuthRules extends SchemaAuthorization<any, any, any>>(auth: AuthRules[]) => ModelSchema<SetTypeSubArg<T, 'authorization', AuthRules[]>, UsedMethods | 'authorization'>;
+}, UsedMethods> & BaseSchema<T> & DDBSchemaBrand;
+type RDSModelSchemaFunctions = 'setSqlStatementFolderPath' | 'addQueries' | 'addMutations' | 'addSubscriptions' | 'authorization';
+export type RDSModelSchema<T extends RDSModelSchemaParamShape, UsedMethods extends RDSModelSchemaFunctions = never> = Omit<{
+    setSqlStatementFolderPath: (path: string) => RDSModelSchema<T, UsedMethods | 'setSqlStatementFolderPath'>;
+    addQueries: (types: Record<string, QueryCustomOperation>) => RDSModelSchema<T, UsedMethods | 'addQueries'>;
+    addMutations: (types: Record<string, MutationCustomOperation>) => RDSModelSchema<T, UsedMethods | 'addMutations'>;
+    addSubscriptions: (types: Record<string, SubscriptionCustomOperation>) => RDSModelSchema<T, UsedMethods | 'addSubscriptions'>;
+    authorization: <AuthRules extends SchemaAuthorization<any, any, any>>(auth: AuthRules[]) => RDSModelSchema<SetTypeSubArg<T, 'authorization', AuthRules[]>, UsedMethods | 'authorization'> & RDSSchemaBrand;
+}, UsedMethods> & BaseSchema<T> & RDSSchemaBrand;
 /**
  * Amplify API Next Model Schema shape
  */
@@ -78,7 +87,7 @@ type SchemaReturnType<DE extends DatasourceEngine, Types extends ModelSchemaCont
     types: Types;
     authorization: [];
     configuration: any;
-}> : SQLModelSchema<{
+}> : RDSModelSchema<{
     types: Types;
     authorization: [];
     configuration: any;

package/lib-esm/src/ModelSchema.js

@@ -1,8 +1,13 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.configure = exports.schema = exports.isModelSchema = void 0;
+exports.configure = exports.schema = exports.isModelSchema = exports.ddbSchemaBrandName = exports.rdsSchemaBrand = exports.rdsSchemaBrandName = void 0;
 const ModelType_1 = require("./ModelType");
 const SchemaProcessor_1 = require("./SchemaProcessor");
+const util_1 = require("./util");
+exports.rdsSchemaBrandName = 'RDSSchema';
+exports.rdsSchemaBrand = (0, util_1.brand)(exports.rdsSchemaBrandName);
+exports.ddbSchemaBrandName = 'DDBSchema';
+const ddbSchemaBrand = (0, util_1.brand)(exports.ddbSchemaBrandName);
 /**
  * Filter the schema types down to only include the ModelTypes as SchemaModelType
  *
@@ -29,9 +34,23 @@ const isModelSchema = (schema) => {
     return typeof schema === 'object' && schema.data !== undefined;
 };
 exports.isModelSchema = isModelSchema;
-function _sqlSchemaExtension(schema) {
+function _rdsSchema(types, config) {
+    const data = {
+        types,
+        authorization: [],
+        configuration: config,
+    };
     return {
-        ...schema,
+        data,
+        transform() {
+            const internalSchema = { data };
+            return (0, SchemaProcessor_1.processSchema)({ schema: internalSchema });
+        },
+        authorization(rules) {
+            this.data.authorization = rules;
+            const { authorization: _, ...rest } = this;
+            return rest;
+        },
         setSqlStatementFolderPath(path) {
             this.data.setSqlStatementFolderPath = path;
             const { setSqlStatementFolderPath: _, ...rest } = this;
@@ -52,9 +71,10 @@ function _sqlSchemaExtension(schema) {
             const { addSubscriptions: _, ...rest } = this;
             return rest;
         },
+        ...exports.rdsSchemaBrand,
     };
 }
-function _baseSchema(types, config) {
+function _ddbSchema(types, config) {
     const data = {
         types,
         authorization: [],
@@ -72,13 +92,14 @@ function _baseSchema(types, config) {
             return rest;
         },
         models: filterSchemaModelTypes(data.types),
+        ...ddbSchemaBrand,
     };
 }
 function bindConfigToSchema(config) {
     return (types) => {
         return (config.database.engine === 'dynamodb'
-            ? _baseSchema(types, config)
-            : _sqlSchemaExtension(_baseSchema(types, config)));
+            ? _ddbSchema(types, config)
+            : _rdsSchema(types, config));
     };
 }
 /**

package/lib-esm/src/RefType.d.ts

@@ -9,6 +9,7 @@ type RefTypeData = {
     valueRequired: boolean;
     array: boolean;
     arrayRequired: boolean;
+    mutationOperations: MutationOperations[];
     authorization: Authorization<any, any, any>[];
 };
 export type RefTypeParamShape = {
@@ -19,6 +20,7 @@ export type RefTypeParamShape = {
     arrayRequired: boolean;
     authorization: Authorization<any, any, any>[];
 };
+type MutationOperations = 'create' | 'update' | 'delete';
 export type RefType<T extends RefTypeParamShape, K extends keyof RefType<T> = never, Auth = undefined> = Omit<{
     /**
      * Marks a field as required.
@@ -33,6 +35,7 @@ export type RefType<T extends RefTypeParamShape, K extends keyof RefType<T> = ne
      * multiple authorization rules for this field.
      */
     authorization<AuthRuleType extends Authorization<any, any, any>>(rules: AuthRuleType[]): RefType<T, K | 'authorization', AuthRuleType>;
+    mutations(operations: MutationOperations[]): RefType<T, K | 'mutations'>;
 }, K> & {
     [__auth]?: Auth;
 } & Brand<typeof brandName>;

package/lib-esm/src/RefType.js

@@ -13,6 +13,7 @@ function _ref(link) {
         valueRequired: false,
         array: false,
         arrayRequired: false,
+        mutationOperations: [],
         authorization: [],
     };
     const builder = brandedBuilder({
@@ -33,6 +34,10 @@ function _ref(link) {
             data.authorization = rules;
             return this;
         },
+        mutations(operations) {
+            data.mutationOperations = operations;
+            return this;
+        },
     });
     return { ...builder, data };
 }

package/lib-esm/src/SchemaProcessor.js

@@ -176,8 +176,8 @@ function transformFunctionHandler(handlers, callSignature) {
     });
     return { gqlHandlerContent, lambdaFunctionDefinition };
 }
-function customOperationToGql(typeName, typeDef, authorization, isCustom = false, databaseType) {
-    const { arguments: fieldArgs, returnType, functionRef, handlers, } = typeDef.data;
+function customOperationToGql(typeName, typeDef, authorization, isCustom = false, databaseType, getRefType) {
+    const { arguments: fieldArgs, typeName: opType, returnType, functionRef, handlers, subscriptionSource, } = typeDef.data;
     let callSignature = typeName;
     const implicitModels = [];
     const { authString } = isCustom
@@ -218,6 +218,21 @@ function customOperationToGql(typeName, typeDef, authorization, isCustom = false
     else if (databaseType === 'sql' && handler && brand === 'sqlReference') {
         gqlHandlerContent = `@sql(reference: "${(0, Handler_1.getHandlerData)(handler)}") `;
     }
+    if (opType === 'Subscription') {
+        const subscriptionSources = subscriptionSource
+            .flatMap((source) => {
+            const refTarget = source.data.link;
+            const { type } = getRefType(refTarget, typeName);
+            if (type === 'CustomOperation') {
+                return refTarget;
+            }
+            if (type === 'Model') {
+                return source.data.mutationOperations.map((op) => `${op}${refTarget}`);
+            }
+        })
+            .join('", "');
+        gqlHandlerContent += `@aws_subscribe(mutations: ["${subscriptionSources}"]) `;
+    }
     const gqlField = `${callSignature}: ${returnTypeName} ${gqlHandlerContent}${authString}`;
     return { gqlField, models: implicitModels, lambdaFunctionDefinition };
 }
@@ -732,6 +747,31 @@ const extractFunctionSchemaAccess = (authRules) => {
     }
     return { schemaAuth, functionSchemaAccess };
 };
+/**
+ * Returns a closure for retrieving reference type and definition from schema
+ */
+const getRefTypeForSchema = (schema) => {
+    const getRefType = (source, target) => {
+        const typeDef = schema.data.types[source];
+        if (typeDef === undefined) {
+            throw new Error(`Invalid ref. ${target} is referencing ${source} which is not defined in the schema`);
+        }
+        if (isInternalModel(typeDef)) {
+            return { type: 'Model', def: typeDef };
+        }
+        if (isCustomOperation(typeDef)) {
+            return { type: 'CustomOperation', def: typeDef };
+        }
+        if (isCustomType(typeDef)) {
+            return { type: 'CustomType', def: typeDef };
+        }
+        if (isEnumType(typeDef)) {
+            return { type: 'Enum', def: typeDef };
+        }
+        throw new Error(`Invalid ref. ${target} is referencing ${source} which is neither a Model, Custom Operation, Custom Type, or Enum`);
+    };
+    return getRefType;
+};
 const schemaPreprocessor = (schema) => {
     const gqlModels = [];
     const customQueries = [];
@@ -746,6 +786,7 @@ const schemaPreprocessor = (schema) => {
     const fkFields = allImpliedFKs(schema);
     const topLevelTypes = Object.entries(schema.data.types);
     const { schemaAuth, functionSchemaAccess } = extractFunctionSchemaAccess(schema.data.authorization);
+    const getRefType = getRefTypeForSchema(schema);
     for (const [typeName, typeDef] of topLevelTypes) {
         validateAuth(typeDef.data?.authorization);
         const mostRelevantAuthRules = typeDef.data?.authorization?.length > 0
@@ -773,7 +814,7 @@ const schemaPreprocessor = (schema) => {
             }
             else if (isCustomOperation(typeDef)) {
                 const { typeName: opType } = typeDef.data;
-                const { gqlField, models, jsFunctionForField, lambdaFunctionDefinition, } = transformCustomOperations(typeDef, typeName, mostRelevantAuthRules, databaseType);
+                const { gqlField, models, jsFunctionForField, lambdaFunctionDefinition, } = transformCustomOperations(typeDef, typeName, mostRelevantAuthRules, databaseType, getRefType);
                 lambdaFunctions = lambdaFunctionDefinition;
                 topLevelTypes.push(...models);
                 if (jsFunctionForField) {
@@ -849,8 +890,8 @@ const schemaPreprocessor = (schema) => {
         lambdaFunctions,
     };
 };
-function validateCustomOperations(typeDef, typeName, authRules) {
-    const { functionRef, handlers } = typeDef.data;
+function validateCustomOperations(typeDef, typeName, authRules, getRefType) {
+    const { functionRef, handlers, typeName: opType, subscriptionSource, returnType, } = typeDef.data;
     // TODO: remove `functionRef` after deprecating
     const handlerConfigured = functionRef !== null || handlers?.length;
     const authConfigured = authRules.length > 0;
@@ -873,6 +914,57 @@ function validateCustomOperations(typeDef, typeName, authRules) {
             throw new Error(`Field handlers must be of the same type. ${typeName} has been configured with ${configuredHandlersStr}`);
         }
     }
+    if (opType === 'Subscription') {
+        if (subscriptionSource.length < 1) {
+            throw new Error(`${typeName} is missing a mutation source. Custom subscriptions must reference a mutation source via subscription().for(a.ref('ModelOrOperationName')) `);
+        }
+        subscriptionSource.forEach((source) => {
+            const sourceName = source.data.link;
+            const { type, def } = getRefType(sourceName, typeName);
+            if (type !== 'Model' && source.data.mutationOperations.length > 0) {
+                throw new Error(`Invalid subscription definition. .mutations() modifier can only be used with a Model ref. ${typeName} is referencing ${type}`);
+            }
+            if (type === 'Model' && source.data.mutationOperations.length === 0) {
+                throw new Error(`Invalid subscription definition. .mutations() modifier must be used with a Model ref subscription source. ${typeName} is referencing ${sourceName} without specifying a mutation`);
+            }
+            if (type === 'CustomOperation' && def.data.typeName !== 'Mutation') {
+                throw new Error(`Invalid subscription definition. .for() can only reference a mutation. ${typeName} is referencing ${sourceName} which is a ${def.data.typeName}`);
+            }
+            // Ensure subscription return type matches the return type of triggering mutation(s)
+            // TODO: when we remove .returns() for custom subscriptions, minor changes will be needed here. Instead of comparing subscriptionSource return val
+            //  to a root returnType, we'll need to ensure that each subscriptionSource has the same return type
+            if (returnType.data.type === 'ref') {
+                const returnTypeName = returnType.data.link;
+                if (type === 'Model') {
+                    if (returnTypeName !== sourceName ||
+                        returnType.data.array !== source.data.array) {
+                        throw new Error(`Invalid subscription definition. Subscription return type must match the return type of the mutation triggering it. ${typeName} is referencing ${sourceName} which has a different return type`);
+                    }
+                }
+                if (type === 'CustomOperation') {
+                    const customOperationReturnType = def.data.returnType.data.link;
+                    const customOperationReturnTypeArray = def.data.returnType.data.array;
+                    if (returnTypeName !== customOperationReturnType ||
+                        returnType.data.array !== customOperationReturnTypeArray) {
+                        throw new Error(`Invalid subscription definition. Subscription return type must match the return type of the mutation triggering it. ${typeName} is referencing ${sourceName} which has a different return type`);
+                    }
+                }
+            }
+            else if (returnType.data.fieldType !== undefined) {
+                if (type === 'Model') {
+                    throw new Error(`Invalid subscription definition. Subscription return type must match the return type of the mutation triggering it. ${typeName} is referencing ${sourceName} which has a different return type`);
+                }
+                if (type === 'CustomOperation') {
+                    const customOperationReturnType = def.data.returnType.data.fieldType;
+                    const customOperationReturnTypeArray = def.data.returnType.data.array;
+                    if (returnType.data.fieldType !== customOperationReturnType ||
+                        returnType.data.array !== customOperationReturnTypeArray) {
+                        throw new Error(`Invalid subscription definition. Subscription return type must match the return type of the mutation triggering it. ${typeName} is referencing ${sourceName} which has a different return type`);
+                    }
+                }
+            }
+        });
+    }
 }
 const isCustomHandler = (handler) => {
     return Array.isArray(handler) && (0, util_1.getBrand)(handler[0]) === 'customHandler';
@@ -933,15 +1025,15 @@ const handleCustom = (handlers, opType, typeName) => {
     };
     return jsFn;
 };
-function transformCustomOperations(typeDef, typeName, authRules, databaseType) {
+function transformCustomOperations(typeDef, typeName, authRules, databaseType, getRefType) {
     const { typeName: opType, handlers } = typeDef.data;
     let jsFunctionForField = undefined;
-    validateCustomOperations(typeDef, typeName, authRules);
+    validateCustomOperations(typeDef, typeName, authRules, getRefType);
     if (isCustomHandler(handlers)) {
         jsFunctionForField = handleCustom(handlers, opType, typeName);
     }
     const isCustom = Boolean(jsFunctionForField);
-    const { gqlField, models, lambdaFunctionDefinition } = customOperationToGql(typeName, typeDef, authRules, isCustom, databaseType);
+    const { gqlField, models, lambdaFunctionDefinition } = customOperationToGql(typeName, typeDef, authRules, isCustom, databaseType, getRefType);
     return { gqlField, models, jsFunctionForField, lambdaFunctionDefinition };
 }
 function generateCustomOperationTypes({ queries, mutations, subscriptions, }) {