fix: re-enable IAM auth for custom types (#2961)

packages/amplify-graphql-api-construct/src/__tests__/__functional__/disable-gen1-patterns.test.ts

@@ -161,6 +161,22 @@ describe('Deprecate Gen 1 patterns', () => {
     );
   });
 
+  test('does not allow implicit fields on @hasMany', () => {
+    const stack = verifySchema(/* GraphQL */ `
+      type Post @model {
+        author: Author @belongsTo
+      }
+
+      type Author @model {
+        posts: [Post] @hasMany
+      }
+    `);
+    Annotations.fromStack(stack).hasWarning(
+      '/Default/TestApi/GraphQLAPI',
+      'fields argument on @hasMany is deprecated. Modify Author.posts to use references instead. This functionality will be removed in the next major release.',
+    );
+  });
+
   test('does not print warning when fields is not used on @hasMany', () => {
     const stack = verifySchema(/* GraphQL */ `
       type Post @model {

packages/amplify-graphql-auth-transformer/API.md

@@ -110,8 +110,6 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
     // (undocumented)
     addAutoGeneratedRelationalFields: (ctx: TransformerContextProvider, def: ObjectTypeDefinitionNode, allowedFields: Set<string>, fields: readonly string[]) => void;
     // (undocumented)
-    addCustomOperationFieldsToAuthNonModelConfig: (ctx: TransformerTransformSchemaStepContextProvider) => void;
-    // (undocumented)
     addFieldResolverForDynamicAuth: (ctx: TransformerContextProvider, def: ObjectTypeDefinitionNode, typeName: string, fieldName: string) => void;
     // (undocumented)
     addFieldsToObject: (ctx: TransformerTransformSchemaStepContextProvider, modelName: string, ownerFields: Array<string>) => void;

packages/amplify-graphql-auth-transformer/src/__tests__/iam-custom-operations.test.ts

@@ -367,8 +367,7 @@ describe('Custom operations have @aws_iam directives when enableIamAuthorization
       expect(out.schema).not.toMatch(/onUpdateFooCustom: String.*@aws_iam/);
     });
 
-    // TODO: Enable this test once we fix https://github.com/aws-amplify/amplify-category-api/issues/2929
-    test.skip('Adds @aws_iam to non-model custom types when there is no model', () => {
+    test('Adds @aws_iam to non-model custom types when there is no model', () => {
       const strategy = makeStrategy(strategyType);
       const schema = /* GraphQL */ `
         type Foo {
@@ -403,8 +402,7 @@ describe('Custom operations have @aws_iam directives when enableIamAuthorization
       expect(out.schema).toMatch(/type Foo.*@aws_iam/);
     });
 
-    // TODO: Enable this test once we fix https://github.com/aws-amplify/amplify-category-api/issues/2929
-    test.skip('Adds @aws_iam to non-model custom types when there is a model', () => {
+    test('Adds @aws_iam to non-model custom types when there is a model', () => {
       const strategy = makeStrategy(strategyType);
       const schema = /* GraphQL */ `
         type Todo @model {

packages/amplify-graphql-auth-transformer/src/graphql-auth-transformer.ts

@@ -6,10 +6,12 @@ import {
   getModelDataSourceNameForTypeName,
   getSortKeyFieldNames,
   getSubscriptionFilterInputName,
+  hasDirectiveWithName,
   InvalidDirectiveError,
   isBuiltInGraphqlNode,
   isDynamoDbModel,
   isModelType,
+  isObjectTypeDefinitionNode,
   isSqlModel,
   MappingTemplate,
   TransformerAuthBase,
@@ -20,22 +22,22 @@ import {
   DataSourceProvider,
   MutationFieldType,
   QueryFieldType,
-  TransformerTransformSchemaStepContextProvider,
+  TransformerAuthProvider,
+  TransformerBeforeStepContextProvider,
   TransformerContextProvider,
   TransformerResolverProvider,
   TransformerSchemaVisitStepContextProvider,
-  TransformerAuthProvider,
-  TransformerBeforeStepContextProvider,
+  TransformerTransformSchemaStepContextProvider,
 } from '@aws-amplify/graphql-transformer-interfaces';
 import {
   DirectiveNode,
   FieldDefinitionNode,
-  ObjectTypeDefinitionNode,
   InterfaceTypeDefinitionNode,
   Kind,
-  TypeDefinitionNode,
   ListValueNode,
+  ObjectTypeDefinitionNode,
   StringValueNode,
+  TypeDefinitionNode,
 } from 'graphql';
 import { merge } from 'lodash';
 import {
@@ -103,6 +105,7 @@ import {
   isFieldRoleHavingAccessToBothSide,
   isDynamicAuthOrCustomAuth,
   isIdenticalAuthRole,
+  addDirectivesToObject,
 } from './utils';
 import {
   defaultIdentityClaimWarning,
@@ -345,30 +348,46 @@ export class AuthTransformer extends TransformerAuthBase implements TransformerA
   };
 
   /**
-   * Adds custom Queries, Mutations, and Subscriptions to the authNonModelConfig map to ensure they are included when adding implicit
-   * aws_iam auth directives.
+   * If needed, adds aws_iam auth directive to non-model types
    */
-  addCustomOperationFieldsToAuthNonModelConfig = (ctx: TransformerTransformSchemaStepContextProvider): void => {
+  private addIamAuthDirectiveToNonModelTypes = (ctx: TransformerTransformSchemaStepContextProvider): void => {
     if (!ctx.transformParameters.sandboxModeEnabled && !ctx.synthParameters.enableIamAccess) {
       return;
     }
 
-    const hasAwsIamDirective = (field: FieldDefinitionNode): boolean => {
-      return field.directives?.some((dir) => dir.name.value === 'aws_iam');
-    };
+    const nonModelObjects = ctx.inputDocument.definitions
+      .filter(isObjectTypeDefinitionNode)
+      .filter((objectDef) => !isBuiltInGraphqlNode(objectDef))
+      .filter((objectDef) => !hasDirectiveWithName(objectDef, 'model'))
+      .filter((objectDef) => !hasDirectiveWithName(objectDef, 'aws_iam'));
+
+    nonModelObjects.forEach((object) => {
+      const typeName = object.name.value;
+      addDirectivesToObject(ctx, typeName, [makeDirective('aws_iam', [])]);
+    });
+  };
+
+  /**
+   * If needed, adds aws_iam auth directive to custom operations (Queries, Mutations, Subscriptions)
+   */
+  private addIamAuthDirectiveToCustomOperationFields = (ctx: TransformerTransformSchemaStepContextProvider): void => {
+    if (!ctx.transformParameters.sandboxModeEnabled && !ctx.synthParameters.enableIamAccess) {
+      return;
+    }
 
-    const allObjects = ctx.inputDocument.definitions.filter(isBuiltInGraphqlNode);
-    allObjects.forEach((object) => {
+    const builtInObjects = ctx.inputDocument.definitions.filter(isBuiltInGraphqlNode);
+    builtInObjects.forEach((object) => {
       const typeName = object.name.value;
-      const fieldsWithoutIamDirective = object.fields.filter((field) => !hasAwsIamDirective(field));
+      const fieldsWithoutIamDirective = object.fields.filter((field) => !hasDirectiveWithName(field, 'aws_iam'));
       fieldsWithoutIamDirective.forEach((field) => {
         addDirectivesToField(ctx, typeName, field.name.value, [makeDirective('aws_iam', [])]);
       });
     });
   };
 
   transformSchema = (context: TransformerTransformSchemaStepContextProvider): void => {
-    this.addCustomOperationFieldsToAuthNonModelConfig(context);
+    this.addIamAuthDirectiveToNonModelTypes(context);
+    this.addIamAuthDirectiveToCustomOperationFields(context);
 
     const searchableAggregateServiceDirectives = new Set<AuthProvider>();
 

packages/amplify-graphql-auth-transformer/src/utils/schema.ts

@@ -19,6 +19,7 @@ import { ObjectTypeDefinitionNode, FieldDefinitionNode, DirectiveNode, NamedType
 import {
   blankObjectExtension,
   extendFieldWithDirectives,
+  extendObjectWithDirectives,
   extensionWithDirectives,
   graphqlName,
   isListType,
@@ -213,7 +214,7 @@ export const addDirectivesToField = (
   typeName: string,
   fieldName: string,
   directives: Array<DirectiveNode>,
-) => {
+): void => {
   const type = ctx.output.getType(typeName) as ObjectTypeDefinitionNode;
   if (type) {
     const field = type.fields?.find((f) => f.name.value === fieldName);
@@ -230,6 +231,17 @@ export const addDirectivesToField = (
   }
 };
 
+export const addDirectivesToObject = (
+  ctx: TransformerTransformSchemaStepContextProvider,
+  typeName: string,
+  directives: Array<DirectiveNode>,
+): void => {
+  const type = ctx.output.getType(typeName) as ObjectTypeDefinitionNode;
+  if (type) {
+    ctx.output.putType(extendObjectWithDirectives(type, directives));
+  }
+};
+
 /**
  * addSubscriptionArguments
  */

packages/amplify-graphql-transformer-core/API.md

@@ -362,6 +362,9 @@ export interface GraphQLTransformOptions {
     readonly userDefinedSlots?: Record<string, UserDefinedSlot[]>;
 }
 
+// @public (undocumented)
+export const hasDirectiveWithName: (node: FieldDefinitionNode | InterfaceTypeDefinitionNode | ObjectTypeDefinitionNode, name: string) => boolean;
+
 // @public (undocumented)
 export type ImportAppSyncAPIInputs = {
     apiName: string;

packages/amplify-graphql-transformer-core/src/index.ts

@@ -56,6 +56,7 @@ export {
   getSubscriptionFilterInputName,
   getTable,
   getType,
+  hasDirectiveWithName,
   isAmplifyDynamoDbModelDataSourceStrategy,
   isBuiltInGraphqlNode,
   isDefaultDynamoDbModelDataSourceStrategy,

packages/amplify-graphql-transformer-core/src/utils/graphql-utils.ts

@@ -48,3 +48,13 @@ export const getField = (obj: ObjectTypeDefinitionNode, fieldName: string): Fiel
 
 export const getType = (schema: DocumentNode, typeName: string): ObjectTypeDefinitionNode | undefined =>
   schema.definitions.find((def) => isObjectTypeDefinitionNode(def) && def.name.value === typeName) as ObjectTypeDefinitionNode | undefined;
+
+/**
+ * Returns true if the node has a directive named `name`
+ */
+export const hasDirectiveWithName = (
+  node: FieldDefinitionNode | InterfaceTypeDefinitionNode | ObjectTypeDefinitionNode,
+  name: string,
+): boolean => {
+  return node.directives?.some((d) => d.name.value === name) ?? false;
+};

packages/graphql-transformer-common/API.md

@@ -63,6 +63,9 @@ export const directiveExists: (definition: ObjectTypeDefinitionNode, name: strin
 // @public (undocumented)
 export function extendFieldWithDirectives(field: FieldDefinitionNode, directives: DirectiveNode[]): FieldDefinitionNode;
 
+// @public (undocumented)
+export function extendObjectWithDirectives(object: ObjectTypeDefinitionNode, directives: DirectiveNode[]): ObjectTypeDefinitionNode;
+
 // @public (undocumented)
 export function extensionWithDirectives(object: ObjectTypeExtensionNode, directives: DirectiveNode[]): ObjectTypeExtensionNode;
 

packages/graphql-transformer-common/src/__tests__/definition.test.ts

@@ -0,0 +1,150 @@
+import { isObjectTypeDefinitionNode } from '@aws-amplify/graphql-transformer-core';
+import { extendObjectWithDirectives, getNonModelTypes, makeDirective } from '../definition';
+import { Kind, ObjectTypeDefinitionNode, parse } from 'graphql';
+
+const testModel = `
+    type Post @model {
+        id: ID!
+        title: String
+    }
+`;
+
+describe('gets Non-Model Types', () => {
+  it('should return empty list if no non-model types', () => {
+    const document = parse(testModel);
+    const nonModelTypes = getNonModelTypes(document);
+    expect(nonModelTypes).toEqual([]);
+  });
+
+  it('should return non-model types', () => {
+    const testDocument = parse(`
+            type Post @model {
+                id: ID!
+                nonModel: NonModel
+            }
+
+            type NonModel {
+                id: ID!
+            }
+        `);
+    const nonModelTypes = getNonModelTypes(testDocument);
+    expect(nonModelTypes).toHaveLength(1);
+    expect((nonModelTypes[0] as ObjectTypeDefinitionNode).name?.value).toEqual('NonModel');
+  });
+});
+
+describe('extendObjectWithDirectives', () => {
+  it('should extend an object with one directive if the object does not already have directives', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(0);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('dir0', [])]);
+
+    expect(updatedObject.directives?.length).toEqual(1);
+    expect(updatedObject?.directives?.[0].name.value).toEqual('dir0');
+  });
+
+  it('should extend an object with one directive if the object already has directives', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo @existingDir0 @existingDir1 {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(2);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('dir0', [])]);
+
+    expect(updatedObject.directives?.length).toEqual(3);
+    expect(updatedObject?.directives?.[2].name.value).toEqual('dir0');
+  });
+
+  it('should extend an object with multiple directives if the object does not already have directives', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(0);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('dir0', []), makeDirective('dir1', [])]);
+
+    expect(updatedObject.directives?.length).toEqual(2);
+    expect(updatedObject?.directives?.[0].name.value).toEqual('dir0');
+    expect(updatedObject?.directives?.[1].name.value).toEqual('dir1');
+  });
+
+  it('should extend an object with multiple directives if the object already has directives', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo @existingDir0 @existingDir1 {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(2);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('dir0', []), makeDirective('dir1', [])]);
+
+    expect(updatedObject.directives?.length).toEqual(4);
+    expect(updatedObject?.directives?.[2].name.value).toEqual('dir0');
+    expect(updatedObject?.directives?.[3].name.value).toEqual('dir1');
+  });
+
+  it('should not add duplicate directives', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo @existingDir0 @existingDir1 {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(2);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('existingDir0', []), makeDirective('dir1', [])]);
+
+    expect(updatedObject.directives?.length).toEqual(3);
+    expect(updatedObject?.directives?.[0].name.value).toEqual('existingDir0');
+    expect(updatedObject?.directives?.[1].name.value).toEqual('existingDir1');
+    expect(updatedObject?.directives?.[2].name.value).toEqual('dir1');
+  });
+
+  it('should return the object unchanged if directives is empty', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo @existingDir0 @existingDir1 {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(2);
+
+    const updatedObject = extendObjectWithDirectives(object, []);
+
+    expect(updatedObject).toBe(object);
+  });
+
+  it('should return the object unchanged if multiple directives are specified but all already exist', () => {
+    const doc = parse(/* GraphQL */ `
+      type Foo @existingDir0 @existingDir1 {
+        id: ID!
+      }
+    `);
+
+    const object = doc.definitions.find(isObjectTypeDefinitionNode)!;
+    expect(object?.directives?.length).toEqual(2);
+
+    const updatedObject = extendObjectWithDirectives(object, [makeDirective('existingDir0', []), makeDirective('existingDir1', [])]);
+
+    expect(updatedObject).toBe(object);
+  });
+});