@aws-amplify/graphql-model-transformer 0.9.3-beta.0 → 0.10.0-apiext2.0

package/src/__tests__/__snapshots__/model-transformer.test.ts.snap

@@ -13303,7 +13303,7 @@ $util.toJson($UpdateItem)
 ## [End] Mutation Update resolver. **"
 `;
 
-exports[`ModelTransformer:  should support timestamp parameters when generating pipelineFunctions and output schema 1`] = `
+exports[`ModelTransformer:  should support timestamp parameters when generating resolvers and output schema 1`] = `
 "
 type Post {
   id: ID!
@@ -13457,7 +13457,7 @@ type Subscription {
 "
 `;
 
-exports[`ModelTransformer:  should support timestamp parameters when generating pipelineFunctions and output schema 2`] = `
+exports[`ModelTransformer:  should support timestamp parameters when generating resolvers and output schema 2`] = `
 "## [Start] Create Request template. **
 ## Set the default values to put request **
 #set( $mergedValues = $util.defaultIfNull($ctx.stash.defaultValues, {}) )
@@ -13525,7 +13525,7 @@ $util.toJson($PutObject)
 ## [End] Create Request template. **"
 `;
 
-exports[`ModelTransformer:  should support timestamp parameters when generating pipelineFunctions and output schema 3`] = `
+exports[`ModelTransformer:  should support timestamp parameters when generating resolvers and output schema 3`] = `
 "## [Start] Mutation Update resolver. **
 ## Set the default values to put request **
 #set( $mergedValues = $util.defaultIfNull($ctx.stash.defaultValues, {}) )

package/src/__tests__/model-transformer-override.test.ts

@@ -0,0 +1,41 @@
+import { ModelTransformer } from '@aws-amplify/graphql-model-transformer';
+import { GraphQLTransform } from '@aws-amplify/graphql-transformer-core';
+import path from 'path';
+
+const featureFlags = {
+  getBoolean: jest.fn(),
+  getNumber: jest.fn(),
+  getObject: jest.fn(),
+  getString: jest.fn(),
+};
+
+describe('ModelTransformer: ', () => {
+  it('should override  model objects when given override config', () => {
+    const validSchema = `
+      type Post @model {
+        id: ID!
+        comments: [Comment]
+      }
+      type Comment @model{
+        id: String!
+        text: String!
+      }
+    `;
+    const transformer = new GraphQLTransform({
+      transformers: [new ModelTransformer()],
+      overrideConfig: {
+        overrideDir: path.join(__dirname, 'overrides'),
+        overrideFlag: true,
+        resourceName: 'myResource',
+      },
+      featureFlags,
+    });
+    const out = transformer.transform(validSchema);
+    expect(out).toBeDefined();
+    const postStack = out.stacks.Post;
+    const commentStack = out.stacks.Comment;
+
+    expect(postStack).toMatchSnapshot();
+    expect(commentStack).toMatchSnapshot();
+  });
+});

package/src/__tests__/model-transformer.test.ts

@@ -411,7 +411,7 @@ describe('ModelTransformer: ', () => {
     expect(defaultIdField).toBeDefined();
     expect(getBaseType(defaultIdField.type)).toEqual('Int');
     // It should not add default value for ctx.arg.id as id is of type Int
-    expect(result.pipelineFunctions['Mutation.createPost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.createPost.req.vtl']).toMatchSnapshot();
   });
 
   it('should generate only create mutation', () => {
@@ -670,7 +670,7 @@ describe('ModelTransformer: ', () => {
     validateModelSchema(schema);
   });
 
-  it('should support timestamp parameters when generating pipelineFunctions and output schema', () => {
+  it('should support timestamp parameters when generating resolvers and output schema', () => {
     const validSchema = `
     type Post @model(timestamps: { createdAt: "createdOn", updatedAt: "updatedOn"}) {
       id: ID!
@@ -688,8 +688,8 @@ describe('ModelTransformer: ', () => {
     const schema = parse(result.schema);
     validateModelSchema(schema);
 
-    expect(result.pipelineFunctions['Mutation.createPost.req.vtl']).toMatchSnapshot();
-    expect(result.pipelineFunctions['Mutation.updatePost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.createPost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.updatePost.req.vtl']).toMatchSnapshot();
   });
 
   it('should not to auto generate createdAt and updatedAt when the type in schema is not AWSDateTime', () => {
@@ -712,8 +712,8 @@ describe('ModelTransformer: ', () => {
     const schema = parse(result.schema);
     validateModelSchema(schema);
 
-    expect(result.pipelineFunctions['Mutation.createPost.req.vtl']).toMatchSnapshot();
-    expect(result.pipelineFunctions['Mutation.updatePost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.createPost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.updatePost.req.vtl']).toMatchSnapshot();
   });
 
   it('should have timestamps as nullable fields when the type makes it non-nullable', () => {
@@ -737,8 +737,8 @@ describe('ModelTransformer: ', () => {
     const schema = parse(result.schema);
     validateModelSchema(schema);
 
-    expect(result.pipelineFunctions['Mutation.createPost.req.vtl']).toMatchSnapshot();
-    expect(result.pipelineFunctions['Mutation.updatePost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.createPost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.updatePost.req.vtl']).toMatchSnapshot();
   });
 
   it('should not to include createdAt and updatedAt field when timestamps is set to null', () => {
@@ -759,8 +759,8 @@ describe('ModelTransformer: ', () => {
     const schema = parse(result.schema);
     validateModelSchema(schema);
 
-    expect(result.pipelineFunctions['Mutation.createPost.req.vtl']).toMatchSnapshot();
-    expect(result.pipelineFunctions['Mutation.updatePost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.createPost.req.vtl']).toMatchSnapshot();
+    expect(result.resolvers['Mutation.updatePost.req.vtl']).toMatchSnapshot();
   });
 
   it('should filter known input types from create and update input fields', () => {
@@ -897,10 +897,8 @@ describe('ModelTransformer: ', () => {
     const transformer = new GraphQLTransform({
       transformers: [new ModelTransformer()],
       featureFlags,
-      transformConfig: {
-        ResolverConfig: {
-          project: config,
-        },
+      resolverConfig: {
+        project: config,
       },
     });
     const out = transformer.transform(validSchema);
@@ -908,7 +906,7 @@ describe('ModelTransformer: ', () => {
 
     const definition = out.schema;
     expect(definition).toBeDefined();
-    expect(out.pipelineFunctions).toMatchSnapshot();
+    expect(out.resolvers).toMatchSnapshot();
 
     validateModelSchema(parse(definition));
   });
@@ -934,10 +932,8 @@ describe('ModelTransformer: ', () => {
     const transformer = new GraphQLTransform({
       transformers: [new ModelTransformer()],
       featureFlags,
-      transformConfig: {
-        ResolverConfig: {
-          project: config,
-        },
+      resolverConfig: {
+        project: config,
       },
     });
     const out = transformer.transform(validSchema);
@@ -945,7 +941,7 @@ describe('ModelTransformer: ', () => {
 
     const definition = out.schema;
     expect(definition).toBeDefined();
-    expect(out.pipelineFunctions).toMatchSnapshot();
+    expect(out.resolvers).toMatchSnapshot();
 
     validateModelSchema(parse(definition));
   });
@@ -968,10 +964,8 @@ describe('ModelTransformer: ', () => {
     const transformer = new GraphQLTransform({
       transformers: [new ModelTransformer()],
       featureFlags,
-      transformConfig: {
-        ResolverConfig: {
-          project: config,
-        },
+      resolverConfig: {
+        project: config,
       },
     });
     const out = transformer.transform(validSchema);
@@ -979,7 +973,7 @@ describe('ModelTransformer: ', () => {
 
     const definition = out.schema;
     expect(definition).toBeDefined();
-    expect(out.pipelineFunctions).toMatchSnapshot();
+    expect(out.resolvers).toMatchSnapshot();
 
     validateModelSchema(parse(definition));
   });
@@ -1000,10 +994,8 @@ describe('ModelTransformer: ', () => {
     const transformer = new GraphQLTransform({
       transformers: [new ModelTransformer()],
       featureFlags,
-      transformConfig: {
-        ResolverConfig: {
-          project: config,
-        },
+      resolverConfig: {
+        project: config,
       },
     });
     const out = transformer.transform(validSchema);
@@ -1096,12 +1088,11 @@ describe('ModelTransformer: ', () => {
     }`;
 
     const transformer = new GraphQLTransform({
-      transformConfig: {
-        ResolverConfig: {
-          project: {
-            ConflictDetection: 'VERSION',
-            ConflictHandler: ConflictHandlerType.AUTOMERGE,
-          },
+      transformConfig: {},
+      resolverConfig: {
+        project: {
+          ConflictDetection: 'VERSION',
+          ConflictHandler: ConflictHandlerType.AUTOMERGE,
         },
       },
       sandboxModeEnabled: true,
@@ -1115,8 +1106,8 @@ describe('ModelTransformer: ', () => {
     const queryObject = getObjectType(schema, 'Query');
     expectFields(queryObject!, ['syncTodos']);
     // sync resolvers
-    expect(out.pipelineFunctions['Query.syncTodos.req.vtl']).toMatchSnapshot();
-    expect(out.pipelineFunctions['Query.syncTodos.res.vtl']).toMatchSnapshot();
+    expect(out.resolvers['Query.syncTodos.req.vtl']).toMatchSnapshot();
+    expect(out.resolvers['Query.syncTodos.res.vtl']).toMatchSnapshot();
     // ds table
     cdkExpect(out.rootStack).to(
       haveResource('AWS::DynamoDB::Table', {
@@ -1166,4 +1157,108 @@ describe('ModelTransformer: ', () => {
       }),
     );
   });
+
+  it('should add the model parameters at the root sack', () => {
+    const modelParams = {
+      DynamoDBModelTableReadIOPS: expect.objectContaining({
+        Type: 'Number',
+        Default: 5,
+        Description: 'The number of read IOPS the table should support.',
+      }),
+      DynamoDBModelTableWriteIOPS: expect.objectContaining({
+        Type: 'Number',
+        Default: 5,
+        Description: 'The number of write IOPS the table should support.',
+      }),
+      DynamoDBBillingMode: expect.objectContaining({
+        Type: 'String',
+        Default: 'PAY_PER_REQUEST',
+        AllowedValues: ['PAY_PER_REQUEST', 'PROVISIONED'],
+        Description: 'Configure @model types to create DynamoDB tables with PAY_PER_REQUEST or PROVISIONED billing modes.',
+      }),
+      DynamoDBEnablePointInTimeRecovery: expect.objectContaining({
+        Type: 'String',
+        Default: 'false',
+        AllowedValues: ['true', 'false'],
+        Description: 'Whether to enable Point in Time Recovery on the table.',
+      }),
+      DynamoDBEnableServerSideEncryption: expect.objectContaining({
+        Type: 'String',
+        Default: 'true',
+        AllowedValues: ['true', 'false'],
+        Description: 'Enable server side encryption powered by KMS.',
+      }),
+    };
+    const validSchema = `type Todo @model {
+      name: String
+    }`;
+    const transformer = new GraphQLTransform({
+      sandboxModeEnabled: true,
+      transformers: [new ModelTransformer()],
+    });
+    const out = transformer.transform(validSchema);
+
+    const rootStack = out.rootStack;
+    expect(rootStack).toBeDefined();
+    expect(rootStack.Parameters).toMatchObject(modelParams);
+
+    const todoStack = out.stacks['Todo'];
+    expect(todoStack).toBeDefined();
+    expect(todoStack.Parameters).toMatchObject(modelParams);
+  });
+
+  it('global auth enabled should add apiKey if not default mode of auth', () => {
+    const validSchema = `
+    type Post @model {
+      id: ID!
+      title: String!
+      tags: [Tag]
+    }
+
+    type Tag {
+      id: ID
+      tags: [Tag]
+    }`;
+    const transformer = new GraphQLTransform({
+      authConfig: {
+        defaultAuthentication: {
+          authenticationType: 'AMAZON_COGNITO_USER_POOLS',
+        },
+        additionalAuthenticationProviders: [
+          {
+            authenticationType: 'API_KEY',
+          },
+        ],
+      },
+      sandboxModeEnabled: true,
+      transformers: [new ModelTransformer()],
+    });
+    const out = transformer.transform(validSchema);
+    expect(out).toBeDefined();
+
+    const schema = parse(out.schema);
+    validateModelSchema(schema);
+
+    const postType = getObjectType(schema, 'Post')!;
+    expect(postType).toBeDefined();
+    expect(postType.directives).toBeDefined();
+    expect(postType.directives!.some(dir => dir.name.value === 'aws_api_key')).toEqual(true);
+
+    const tagType = getObjectType(schema, 'Tag')!;
+    expect(tagType).toBeDefined();
+    expect(tagType.directives).toBeDefined();
+    expect(tagType.directives!.some(dir => dir.name.value === 'aws_api_key')).toEqual(true);
+
+    // check operations
+    const queryType = getObjectType(schema, 'Query')!;
+    expect(queryType).toBeDefined();
+    const mutationType = getObjectType(schema, 'Mutation')!;
+    expect(mutationType).toBeDefined();
+    const subscriptionType = getObjectType(schema, 'Subscription')!;
+    expect(subscriptionType).toBeDefined();
+
+    for (const field of [...queryType.fields!, ...mutationType.fields!, ...subscriptionType.fields!]) {
+      expect(field.directives!.some(dir => dir.name.value === 'aws_api_key')).toEqual(true);
+    }
+  });
 });

package/src/__tests__/overrides/build/override.js

@@ -0,0 +1,8 @@
+function override(resource) {
+  resource.api.GraphQLAPI.xrayEnabled = true;
+  resource.models['Post'].modelDDBTable.billingMode = 'PROVISIONED';
+  resource.models['Comment'].modelDDBTable.billingMode = 'PROVISIONED';
+  // override resolver
+  resource.models['Post'].resolvers['subscriptionOnUpdatePostResolver'].requestMappingTemplate = 'mockTemplate';
+}
+exports.override = override;

package/src/definitions.ts

@@ -14,3 +14,5 @@ export const BOOLEAN_FUNCTIONS = new Set<string>(['attributeExists', 'attributeT
 export const ATTRIBUTE_TYPES = ['binary', 'binarySet', 'bool', 'list', 'map', 'number', 'numberSet', 'string', 'stringSet', '_null'];
 
 export const OPERATION_KEY = '__operation';
+
+export const API_KEY_DIRECTIVE = 'aws_api_key';

package/src/graphql-model-transformer.ts

@@ -5,6 +5,7 @@ import {
   SyncConfig,
   SyncUtils,
   TransformerModelBase,
+  TransformerNestedStack,
 } from '@aws-amplify/graphql-transformer-core';
 import {
   AppSyncDataSourceType,
@@ -20,6 +21,7 @@ import {
   TransformerSchemaVisitStepContextProvider,
   TransformerTransformSchemaStepContextProvider,
   TransformerValidationStepContextProvider,
+  TransformerBeforeStepContextProvider,
 } from '@aws-amplify/graphql-transformer-interfaces';
 import { AttributeType, CfnTable, ITable, StreamViewType, Table, TableEncryption } from '@aws-cdk/aws-dynamodb';
 import * as iam from '@aws-cdk/aws-iam';
@@ -51,8 +53,10 @@ import {
   toPascalCase,
 } from 'graphql-transformer-common';
 import {
+  addDirectivesToOperation,
   addModelConditionInputs,
   createEnumModelFilters,
+  extendTypeWithDirectives,
   makeCreateInputField,
   makeDeleteInputField,
   makeListQueryFilterInput,
@@ -60,6 +64,7 @@ import {
   makeModelSortDirectionEnumObject,
   makeMutationConditionInput,
   makeUpdateInputField,
+  propagateApiKeyToNestedTypes,
 } from './graphql-types';
 import {
   generateAuthExpressionForSandboxMode,
@@ -82,6 +87,7 @@ import {
 import { FieldWrapper, InputObjectDefinitionWrapper, ObjectDefinitionWrapper } from './wrappers/object-definition-wrapper';
 import { CfnRole } from '@aws-cdk/aws-iam';
 import md5 from 'md5';
+import { API_KEY_DIRECTIVE } from './definitions';
 
 export type Nullable<T> = T | null;
 export type OptionalAndNullable<T> = Partial<T>;
@@ -167,6 +173,37 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
     this.options = this.getOptions(options);
   }
 
+  before = (ctx: TransformerBeforeStepContextProvider) => {
+    // add model related-parameters to the root stack
+    ctx.stackManager.addParameter(ResourceConstants.PARAMETERS.DynamoDBModelTableReadIOPS, {
+      description: 'The number of read IOPS the table should support.',
+      type: 'Number',
+      default: 5,
+    });
+    ctx.stackManager.addParameter(ResourceConstants.PARAMETERS.DynamoDBModelTableWriteIOPS, {
+      description: 'The number of write IOPS the table should support.',
+      type: 'Number',
+      default: 5,
+    });
+    ctx.stackManager.addParameter(ResourceConstants.PARAMETERS.DynamoDBBillingMode, {
+      description: 'Configure @model types to create DynamoDB tables with PAY_PER_REQUEST or PROVISIONED billing modes.',
+      default: 'PAY_PER_REQUEST',
+      allowedValues: ['PAY_PER_REQUEST', 'PROVISIONED'],
+    });
+    ctx.stackManager.addParameter(ResourceConstants.PARAMETERS.DynamoDBEnablePointInTimeRecovery, {
+      description: 'Whether to enable Point in Time Recovery on the table.',
+      type: 'String',
+      default: 'false',
+      allowedValues: ['true', 'false'],
+    });
+    ctx.stackManager.addParameter(ResourceConstants.PARAMETERS.DynamoDBEnableServerSideEncryption, {
+      description: 'Enable server side encryption powered by KMS.',
+      type: 'String',
+      default: 'true',
+      allowedValues: ['true', 'false'],
+    });
+  };
+
   object = (definition: ObjectTypeDefinitionNode, directive: DirectiveNode, ctx: TransformerSchemaVisitStepContextProvider): void => {
     const isTypeNameReserved =
       definition.name.value === ctx.output.getQueryTypeName() ||
@@ -227,6 +264,7 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
     this.ensureModelSortDirectionEnum(ctx);
     for (const type of this.typesWithModelDirective) {
       const def = ctx.output.getObject(type)!;
+      const hasAuth = def.directives!.some(dir => dir.name.value === 'auth');
 
       // add Non Model type inputs
       this.createNonModelInputs(ctx, def);
@@ -246,6 +284,24 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
       if (ctx.isProjectUsingDataStore()) {
         this.addModelSyncFields(ctx, type);
       }
+      // global auth check
+      if (!hasAuth && ctx.sandboxModeEnabled && ctx.authConfig.defaultAuthentication.authenticationType !== 'API_KEY') {
+        const apiKeyDirArray = [makeDirective(API_KEY_DIRECTIVE, [])];
+        extendTypeWithDirectives(ctx, def.name.value, apiKeyDirArray);
+        propagateApiKeyToNestedTypes(ctx as TransformerContextProvider, def, new Set<string>());
+        for (let operationField of queryFields) {
+          const operationName = operationField.name.value;
+          addDirectivesToOperation(ctx, ctx.output.getQueryTypeName()!, operationName, apiKeyDirArray);
+        }
+        for (let operationField of mutationFields) {
+          const operationName = operationField.name.value;
+          addDirectivesToOperation(ctx, ctx.output.getMutationTypeName()!, operationName, apiKeyDirArray);
+        }
+        for (let operationField of subscriptionsFields) {
+          const operationName = operationField.name.value;
+          addDirectivesToOperation(ctx, ctx.output.getSubscriptionTypeName()!, operationName, apiKeyDirArray);
+        }
+      }
     }
   };
 
@@ -274,10 +330,12 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
           default:
             throw new Error('Unknown query field type');
         }
+        // TODO: add mechanism to add an auth like rule to all non auth @models
+        // this way we can just depend on auth to add the check
         resolver.addToSlot(
           'postAuth',
           MappingTemplate.s3MappingTemplateFromString(
-            generateAuthExpressionForSandboxMode(context),
+            generateAuthExpressionForSandboxMode(context.sandboxModeEnabled),
             `${query.typeName}.${query.fieldName}.{slotName}.{slotIndex}.req.vtl`,
           ),
         );
@@ -304,7 +362,7 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
         resolver.addToSlot(
           'postAuth',
           MappingTemplate.s3MappingTemplateFromString(
-            generateAuthExpressionForSandboxMode(context),
+            generateAuthExpressionForSandboxMode(context.sandboxModeEnabled),
             `${mutation.typeName}.${mutation.fieldName}.{slotName}.{slotIndex}.req.vtl`,
           ),
         );
@@ -352,7 +410,7 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
           resolver.addToSlot(
             'postAuth',
             MappingTemplate.s3MappingTemplateFromString(
-              generateAuthExpressionForSandboxMode(context),
+              generateAuthExpressionForSandboxMode(context.sandboxModeEnabled),
               `${subscription.typeName}.${subscription.fieldName}.{slotName}.{slotIndex}.req.vtl`,
             ),
           );
@@ -1085,30 +1143,42 @@ export class ModelTransformer extends TransformerModelBase implements Transforme
       description: 'The number of read IOPS the table should support.',
       type: 'Number',
       default: 5,
-    }).valueAsString;
+    });
     const writeIops = new cdk.CfnParameter(stack, ResourceConstants.PARAMETERS.DynamoDBModelTableWriteIOPS, {
       description: 'The number of write IOPS the table should support.',
       type: 'Number',
       default: 5,
-    }).valueAsString;
+    });
     const billingMode = new cdk.CfnParameter(stack, ResourceConstants.PARAMETERS.DynamoDBBillingMode, {
       description: 'Configure @model types to create DynamoDB tables with PAY_PER_REQUEST or PROVISIONED billing modes.',
       type: 'String',
       default: 'PAY_PER_REQUEST',
       allowedValues: ['PAY_PER_REQUEST', 'PROVISIONED'],
-    }).valueAsString;
+    });
     const pointInTimeRecovery = new cdk.CfnParameter(stack, ResourceConstants.PARAMETERS.DynamoDBEnablePointInTimeRecovery, {
       description: 'Whether to enable Point in Time Recovery on the table.',
       type: 'String',
       default: 'false',
       allowedValues: ['true', 'false'],
-    }).valueAsString;
+    });
     const enableSSE = new cdk.CfnParameter(stack, ResourceConstants.PARAMETERS.DynamoDBEnableServerSideEncryption, {
       description: 'Enable server side encryption powered by KMS.',
       type: 'String',
       default: 'true',
       allowedValues: ['true', 'false'],
-    }).valueAsString;
+    });
+    // add the connection between the root and nested stack so the values can be passed down
+    (stack as TransformerNestedStack).setParameter(readIops.node.id, cdk.Fn.ref(ResourceConstants.PARAMETERS.DynamoDBModelTableReadIOPS));
+    (stack as TransformerNestedStack).setParameter(writeIops.node.id, cdk.Fn.ref(ResourceConstants.PARAMETERS.DynamoDBModelTableWriteIOPS));
+    (stack as TransformerNestedStack).setParameter(billingMode.node.id, cdk.Fn.ref(ResourceConstants.PARAMETERS.DynamoDBBillingMode));
+    (stack as TransformerNestedStack).setParameter(
+      pointInTimeRecovery.node.id,
+      cdk.Fn.ref(ResourceConstants.PARAMETERS.DynamoDBEnablePointInTimeRecovery),
+    );
+    (stack as TransformerNestedStack).setParameter(
+      enableSSE.node.id,
+      cdk.Fn.ref(ResourceConstants.PARAMETERS.DynamoDBEnableServerSideEncryption),
+    );
 
     // Add conditions.
     // eslint-disable-next-line no-new

package/src/graphql-types/common.ts

@@ -1,14 +1,20 @@
-import { TransformerTransformSchemaStepContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
+import { TransformerContextProvider, TransformerTransformSchemaStepContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
 import {
+  DirectiveNode,
   EnumTypeDefinitionNode,
   FieldDefinitionNode,
   InputObjectTypeDefinitionNode,
   Kind,
+  NamedTypeNode,
   ObjectTypeDefinitionNode,
   TypeDefinitionNode,
 } from 'graphql';
 import {
+  blankObjectExtension,
   DEFAULT_SCALARS,
+  extendFieldWithDirectives,
+  extensionWithDirectives,
+  getBaseType,
   makeArgument,
   makeDirective,
   makeField,
@@ -29,6 +35,7 @@ import {
   SIZE_CONDITIONS,
   STRING_CONDITIONS,
   STRING_FUNCTIONS,
+  API_KEY_DIRECTIVE,
 } from '../definitions';
 import {
   EnumWrapper,
@@ -245,7 +252,96 @@ export function makeEnumFilterInput(fieldWrapper: FieldWrapper): InputObjectType
   return input.serialize();
 }
 
+export const addDirectivesToField = (
+  ctx: TransformerTransformSchemaStepContextProvider,
+  typeName: string,
+  fieldName: string,
+  directives: Array<DirectiveNode>,
+) => {
+  const type = ctx.output.getType(typeName) as ObjectTypeDefinitionNode;
+  if (type) {
+    const field = type.fields?.find(f => f.name.value === fieldName);
+    if (field) {
+      const newFields = [...type.fields!.filter(f => f.name.value !== field.name.value), extendFieldWithDirectives(field, directives)];
+
+      const newType = {
+        ...type,
+        fields: newFields,
+      };
+
+      ctx.output.putType(newType);
+    }
+  }
+};
+
+export const addDirectivesToOperation = (
+  ctx: TransformerTransformSchemaStepContextProvider,
+  typeName: string,
+  operationName: string,
+  directives: Array<DirectiveNode>,
+) => {
+  // add directives to the given operation
+  addDirectivesToField(ctx, typeName, operationName, directives);
+
+  // add the directives to the result type of the operation
+  const type = ctx.output.getType(typeName) as ObjectTypeDefinitionNode;
+  if (type) {
+    const field = type.fields!.find(f => f.name.value === operationName);
+
+    if (field) {
+      const returnFieldType = field.type as NamedTypeNode;
+
+      if (returnFieldType.name) {
+        const returnTypeName = returnFieldType.name.value;
+
+        extendTypeWithDirectives(ctx, returnTypeName, directives);
+      }
+    }
+  }
+};
+
+export const extendTypeWithDirectives = (
+  ctx: TransformerTransformSchemaStepContextProvider,
+  typeName: string,
+  directives: Array<DirectiveNode>,
+): void => {
+  let objectTypeExtension = blankObjectExtension(typeName);
+  objectTypeExtension = extensionWithDirectives(objectTypeExtension, directives);
+  ctx.output.addObjectExtension(objectTypeExtension);
+};
+
 export function makeModelSortDirectionEnumObject(): EnumTypeDefinitionNode {
   const name = 'ModelSortDirection';
   return EnumWrapper.create(name, ['ASC', 'DESC']).serialize();
 }
+// the smaller version of it's @auth equivalent since we only support
+// apikey as the only global auth rule
+export const propagateApiKeyToNestedTypes = (
+  ctx: TransformerContextProvider,
+  def: ObjectTypeDefinitionNode,
+  seenNonModelTypes: Set<string>,
+) => {
+  const nonModelTypePredicate = (fieldType: TypeDefinitionNode): TypeDefinitionNode | undefined => {
+    if (fieldType) {
+      if (fieldType.kind !== 'ObjectTypeDefinition') {
+        return undefined;
+      }
+      const typeModel = fieldType.directives!.find(dir => dir.name.value === 'model');
+      return typeModel !== undefined ? undefined : fieldType;
+    }
+    return fieldType;
+  };
+  const nonModelFieldTypes = def
+    .fields!.map(f => ctx.output.getType(getBaseType(f.type)) as TypeDefinitionNode)
+    .filter(nonModelTypePredicate);
+  for (const nonModelFieldType of nonModelFieldTypes) {
+    const nonModelName = nonModelFieldType.name.value;
+    const hasSeenType = seenNonModelTypes.has(nonModelName);
+    const hasApiKey = nonModelFieldType.directives?.some(dir => dir.name.value === API_KEY_DIRECTIVE) ?? false;
+    if (!hasSeenType && !hasApiKey) {
+      seenNonModelTypes.add(nonModelName);
+      extendTypeWithDirectives(ctx, nonModelName, [makeDirective(API_KEY_DIRECTIVE, [])]);
+      propagateApiKeyToNestedTypes(ctx, nonModelFieldType as ObjectTypeDefinitionNode, seenNonModelTypes);
+    }
+  }
+};