From adcd859a05fc70fae810f7108c5edc6d3f7dac46 Mon Sep 17 00:00:00 2001
From: Pavi <13897936+pavinthan@users.noreply.github.com>
Date: Fri, 21 Feb 2025 19:06:38 +0000
Subject: [PATCH 1/3] feat(data-schema): add authorization to custom type

---
 .../__tests__/CustomType.test-d.ts            |   4 +
 .../data-schema/__tests__/CustomType.test.ts  | 371 +++++++++++++-
 .../__snapshots__/ClientSchema.test.ts.snap   |   4 +-
 .../CustomOperations.test.ts.snap             |  40 +-
 .../__snapshots__/CustomType.test.ts.snap     | 468 +++++++++++++++++-
 packages/data-schema/src/CustomOperation.ts   |   1 +
 packages/data-schema/src/CustomType.ts        |  45 +-
 packages/data-schema/src/SchemaProcessor.ts   |   6 +-
 8 files changed, 894 insertions(+), 45 deletions(-)

diff --git a/packages/data-schema/__tests__/CustomType.test-d.ts b/packages/data-schema/__tests__/CustomType.test-d.ts
index fea6598bd..fc35c3316 100644
--- a/packages/data-schema/__tests__/CustomType.test-d.ts
+++ b/packages/data-schema/__tests__/CustomType.test-d.ts
@@ -13,6 +13,7 @@ describe('CustomType', () => {
     });
 
     type Expected = CustomType<{
+      authorization: [];
       fields: {
         lat: ModelField<Nullable<number>, never, undefined>;
         long: ModelField<Nullable<number>, never, undefined>;
@@ -36,12 +37,15 @@ describe('CustomType', () => {
     });
 
     type Expected = CustomType<{
+      authorization: [];
       fields: {
         content: ModelField<string, 'required', undefined>;
         meta: CustomType<{
+          authorization: [];
           fields: {
             enumField: EnumType<readonly ['value1', 'value2']>;
             deepMeta: CustomType<{
+              authorization: [];
               fields: {
                 description: ModelField<string, 'required', undefined>;
               };
diff --git a/packages/data-schema/__tests__/CustomType.test.ts b/packages/data-schema/__tests__/CustomType.test.ts
index a20e333ff..d697125cc 100644
--- a/packages/data-schema/__tests__/CustomType.test.ts
+++ b/packages/data-schema/__tests__/CustomType.test.ts
@@ -1,5 +1,7 @@
 import { expectTypeTestsToPassAsync } from 'jest-tsd';
-import { a } from '../src/index';
+
+import { PrivateProviders, Operations, Operation } from '../src/Authorization';
+import { a, ClientSchema } from '../src/index';
 
 // evaluates type defs in corresponding test-d.ts file
 it('should not produce static type errors', async () => {
@@ -259,3 +261,370 @@ describe('CustomType transform', () => {
     expect(result).toMatchSnapshot();
   });
 });
+
+// TODO: Remove owner based tests once we have a better way to test auth rules
+describe('CustomType auth rules', () => {
+  it('can define public auth with no provider', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.publicApiKey()),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it('fails loudly on invalid provider', () => {
+    expect(() => {
+      a.schema({
+        Post: a
+          .customType({
+            title: a.string().required(),
+          })
+          // @ts-expect-error
+          .authorization((allow) => allow.public('bad-provider')),
+      });
+    }).toThrow();
+  });
+
+  it('can define private auth with no provider', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.authenticated()),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can specify operations `, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.publicApiKey().to(['create', 'read'])),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a static Admins group rule`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.group('Admins')),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a static [Admins, Moderators] groups rule`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.groups(['Admins', 'Moderators'])),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a dynamic singular groups rule`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.groupDefinedIn('businessUnitOwner')),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a dynamic multi groups rule`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) =>
+          allow.groupsDefinedIn('sharedWithGroups').to(['read']),
+        ),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a dynamic singular groups rule with withClaimIn`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) =>
+          allow
+            .groupDefinedIn('businessUnitOwner')
+            .withClaimIn('someClaimsField'),
+        ),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  it(`can create a dynamic multi groups rule with withClaimIn`, () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+        })
+        .authorization((allow) =>
+          allow
+            .groupsDefinedIn('sharedWithGroups')
+            .to(['read'])
+            .withClaimIn('someClaimsField'),
+        ),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  for (const provider of PrivateProviders) {
+    it(`can define private with with provider ${provider}`, () => {
+      const schema = a.schema({
+        Post: a
+          .customType({
+            title: a.string().required(),
+          })
+          .authorization((allow) => allow.authenticated(provider)),
+      });
+
+      const graphql = schema.transform().schema;
+      expect(graphql).toMatchSnapshot();
+    });
+
+    const TestOperations: Operation[][] = [
+      // each individual operation
+      ...Operations.map((op) => [op]),
+
+      // a couple sanity checks to support a combinations
+      ['create', 'read', 'update', 'delete'],
+      ['create', 'read', 'listen'],
+    ];
+
+    for (const operations of TestOperations) {
+      it(`can define private with with provider ${provider} for operations ${operations}`, () => {
+        const schema = a.schema({
+          widget: a
+            .model({
+              title: a.string().required(),
+            })
+            .authorization((allow) =>
+              allow.authenticated(provider).to(operations),
+            ),
+        });
+
+        const graphql = schema.transform().schema;
+        expect(graphql).toMatchSnapshot();
+      });
+    }
+  }
+
+  it(`can define a custom authorization rule`, () => {
+    const schema = a.schema({
+      Widget: a
+        .model({
+          title: a.string().required(),
+        })
+        .authorization((allow) => allow.custom()),
+    });
+
+    const graphql = schema.transform().schema;
+    expect(graphql).toMatchSnapshot();
+  });
+
+  const TestOperations: Operation[][] = [
+    // each individual operation
+    ...Operations.map((op) => [op]),
+
+    // a couple sanity checks to support a combinations
+    ['create', 'read', 'update', 'delete'],
+    ['create', 'read', 'listen'],
+  ];
+
+  for (const operations of TestOperations) {
+    it(`can define custom auth rule for operations ${operations}`, () => {
+      const schema = a.schema({
+        widget: a
+          .model({
+            title: a.string().required(),
+          })
+          .authorization((allow) => allow.custom().to(operations)),
+      });
+
+      const graphql = schema.transform().schema;
+      expect(graphql).toMatchSnapshot();
+    });
+  }
+
+  it('can define define field-level owner and model-level public simultaneously', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a
+            .string()
+            .required()
+            .authorization((allow) => allow.owner()),
+        })
+        .authorization((allow) => allow.publicApiKey()),
+    });
+  });
+
+  it('can define define field-level owner and model-level private simultaneously', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a
+            .string()
+            .required()
+            .authorization((allow) => allow.owner()),
+        })
+        .authorization((allow) => allow.authenticated()),
+    });
+  });
+
+  it('gives a runtime error if field-level owner and model-level owner conflict', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a
+            .string()
+            .required()
+            .authorization((allow) => allow.ownerDefinedIn('someOwnerField')),
+        })
+        .authorization((allow) => allow.ownersDefinedIn('someOwnerField')),
+    });
+    expect(() => schema.transform().schema).toThrow();
+  });
+
+  it('gives a runtime error if model field and model-level owner auth rule conflicts', () => {
+    const schema = a.schema({
+      Post: a
+        .customType({
+          title: a.string().required(),
+          someOwnerField: a.string(),
+        })
+        .authorization((allow) => allow.ownersDefinedIn('someOwnerField')),
+    });
+    expect(() => schema.transform().schema).toThrow();
+  });
+
+  it('gives a runtime error if field-level owner and schema-level owner conflict', () => {
+    const schema = a
+      .schema({
+        Post: a.customType({
+          title: a
+            .string()
+            .required()
+            .authorization((allow) => allow.ownerDefinedIn('someOwnerField')),
+        }),
+      })
+      .authorization((allow) => allow.ownersDefinedIn('someOwnerField'));
+    expect(() => schema.transform().schema).toThrow();
+  });
+
+  it('gives a runtime error if model field and schema-level owner auth rule conflicts', () => {
+    const schema = a
+      .schema({
+        Post: a.customType({
+          title: a.string().required(),
+          someOwnerField: a.string(),
+        }),
+      })
+      .authorization((allow) => allow.ownersDefinedIn('someOwnerField'));
+    expect(() => schema.transform().schema).toThrow();
+  });
+
+  describe('duplicate field validation does not issue errors for `identifier()`', () => {
+    // because `identifier()` doesn't actually need to match on type. it only
+    // needs the field to exist. The field types specified by `identifier()` are
+    // fallback in case the field isn't defined explicitly.
+    it('explicit `id: string` field', () => {
+      const schema = a
+        .schema({
+          Post: a
+            .customType({
+              id: a.string().required(),
+            }),
+        })
+        .authorization((allow) => allow.publicApiKey());
+      expect(schema.transform().schema).toMatchSnapshot();
+    });
+
+    it('explicit `id: ID` field', () => {
+      const schema = a
+        .schema({
+          Post: a
+            .customType({
+              id: a.id().required(),
+            }),
+        })
+        .authorization((allow) => allow.publicApiKey());
+      expect(schema.transform().schema).toMatchSnapshot();
+    });
+
+    it('explicit `customId: string` field', () => {
+      const schema = a
+        .schema({
+          Post: a
+            .customType({
+              customId: a.string().required(),
+            }),
+        })
+        .authorization((allow) => allow.publicApiKey());
+      expect(schema.transform().schema).toMatchSnapshot();
+    });
+
+    it('explicit multi-field string type PK', () => {
+      const schema = a
+        .schema({
+          Post: a
+            .customType({
+              idFieldA: a.string().required(),
+              idFieldB: a.string().required(),
+            }),
+        })
+        .authorization((allow) => allow.publicApiKey());
+      expect(schema.transform().schema).toMatchSnapshot();
+    });
+
+    it('explicit multi-field mixed types PK', () => {
+      const schema = a
+        .schema({
+          Post: a
+            .customType({
+              idFieldA: a.string().required(),
+              idFieldB: a.integer().required(),
+            })
+        })
+        .authorization((allow) => allow.publicApiKey());
+      expect(schema.transform().schema).toMatchSnapshot();
+    });
+  });
+});
diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
index 97e4d41a1..e389d27a9 100644
--- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
@@ -26,7 +26,7 @@ type Query {
 `;
 
 exports[`ai routes conversations 1`] = `
-"type Profile @aws_api_key
+"type Profile @aws_api_key@aws_api_key
 {
   value: Int
   unit: String
@@ -257,7 +257,7 @@ enum PostStatus {
   published
 }
 
-type PostMeta @aws_cognito_user_pools
+type PostMeta @aws_cognito_user_pools@aws_cognito_user_pools
 {
   viewCount: Int
   approvedOn: AWSDate
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
index 366eb95d6..4851c890e 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
@@ -11,7 +11,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation mutation has ref to custom type argument 1`] = `
-"type post 
+"type post @aws_api_key
 {
   field: String
 }
@@ -26,7 +26,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation mutation with multiple references to the same custom type 1`] = `
-"type post 
+"type post @aws_api_key
 {
   name: String
   number: Int
@@ -48,7 +48,7 @@ exports[`CustomOperation transform custom operations with custom types and refs
   VALUE2
 }
 
-type post 
+type post @aws_api_key
 {
   field: String
   enumField: values
@@ -75,7 +75,7 @@ type Query {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation query has ref to custom type argument 1`] = `
-"type post 
+"type post @aws_api_key
 {
   field: String
 }
@@ -106,7 +106,7 @@ type Query {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation query with enum in referenced custom type argument 1`] = `
-"type post 
+"type post @aws_api_key
 {
   name: String
   status: PostStatus
@@ -152,12 +152,12 @@ type Query {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation query with ref to nested custom type argument 1`] = `
-"type post 
+"type post @aws_api_key
 {
   inner: PostInner
 }
 
-type PostInner 
+type PostInner @aws_api_key
 {
   filter: String!
   e1: PostInnerE1
@@ -198,7 +198,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation subscription has ref to a custom type argument 1`] = `
-"type post 
+"type post @aws_api_key
 {
   title: String
 }
@@ -232,7 +232,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation subscription with multiple references to the same custom type 1`] = `
-"type post 
+"type post @aws_api_key
 {
   name: String
 }
@@ -251,12 +251,12 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation subscription with nested custom type argument 1`] = `
-"type nestedType 
+"type nestedType @aws_api_key
 {
   post: NestedTypePost
 }
 
-type NestedTypePost 
+type NestedTypePost @aws_api_key
 {
   inner: String
 }
@@ -279,13 +279,13 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation with circular references between custom types 1`] = `
-"type post 
+"type post @aws_api_key
 {
   value: String
   child: update
 }
 
-type update 
+type update @aws_api_key
 {
   title: String
   field: post
@@ -311,7 +311,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation with inline custom type containing reference to another custom type 1`] = `
-"type post 
+"type post @aws_api_key
 {
   name: String
 }
@@ -331,7 +331,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation with inline custom type referencing self-referential type 1`] = `
-"type post 
+"type post @aws_api_key
 {
   value: String
   child: post
@@ -352,7 +352,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs custom operation with ref to custom type with self-referencing field 1`] = `
-"type post 
+"type post @aws_api_key
 {
   value: String
   child: post
@@ -369,7 +369,7 @@ type Mutation {
 `;
 
 exports[`CustomOperation transform custom operations with custom types and refs multiple custom operations referencing the same custom type 1`] = `
-"type post 
+"type post @aws_api_key
 {
   name: String
   number: Int
@@ -606,7 +606,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema Schema w model, custom query, mutation, and subscription and ref of custom type with array modifier 1`] = `
-"type PostCustomType @aws_api_key
+"type PostCustomType @aws_api_key@aws_api_key
 {
   title: String
 }
@@ -644,7 +644,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema custom subscriptions Custom subscription where .for() resource has a CustomType return type 1`] = `
-"type CreateCustomTypePostReturnType @aws_api_key
+"type CreateCustomTypePostReturnType @aws_api_key@aws_api_key
 {
   title: String!
 }
@@ -726,7 +726,7 @@ exports[`CustomOperation transform dynamo schema handlers a.handler.custom a.han
   title: String
 }
 
-type GetPostDetailsReturnType @aws_api_key
+type GetPostDetailsReturnType @aws_api_key@aws_api_key
 {
   
 }
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
index 3b1a10820..834d9fb8c 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomType.test.ts.snap
@@ -1,12 +1,448 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`CustomType auth rules can create a dynamic multi groups rule 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can create a dynamic multi groups rule with withClaimIn 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can create a dynamic singular groups rule 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can create a dynamic singular groups rule with withClaimIn 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can create a static [Admins, Moderators] groups rule 1`] = `
+"type Post @aws_cognito_user_pools(cognito_groups: ["Admins", "Moderators"])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can create a static Admins group rule 1`] = `
+"type Post @aws_cognito_user_pools(cognito_groups: ["Admins"])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define a custom authorization rule 1`] = `
+"type Widget @model @auth(rules: [{allow: custom}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations create 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [create]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations create,read,listen 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [create, read, listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations create,read,update,delete 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [create, read, update, delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations delete 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations get 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [get]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations list 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [list]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations listen 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations read 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [read]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations search 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [search]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations sync 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [sync]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define custom auth rule for operations update 1`] = `
+"type widget @model @auth(rules: [{allow: custom, operations: [update]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private auth with no provider 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool 1`] = `
+"type Post @aws_iam
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations create 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [create]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations create,read,listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [create, read, listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations create,read,update,delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [create, read, update, delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations get 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [get]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations list 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [list]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations read 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [read]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations search 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [search]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations sync 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [sync]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider identityPool for operations update 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: iam, operations: [update]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc 1`] = `
+"type Post @aws_oidc
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations create 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [create]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations create,read,listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [create, read, listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations create,read,update,delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [create, read, update, delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations get 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [get]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations list 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [list]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations read 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [read]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations search 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [search]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations sync 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [sync]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider oidc for operations update 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: oidc, operations: [update]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools 1`] = `
+"type Post @aws_cognito_user_pools
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations create 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [create]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations create,read,listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [create, read, listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations create,read,update,delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [create, read, update, delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations delete 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [delete]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations get 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [get]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations list 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [list]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations listen 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [listen]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations read 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [read]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations search 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [search]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations sync 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [sync]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define private with with provider userPools for operations update 1`] = `
+"type widget @model @auth(rules: [{allow: private, provider: userPools, operations: [update]}])
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can define public auth with no provider 1`] = `
+"type Post @aws_api_key
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules can specify operations  1`] = `
+"type Post @aws_api_key
+{
+  title: String!
+}"
+`;
+
+exports[`CustomType auth rules duplicate field validation does not issue errors for \`identifier()\` explicit \`customId: string\` field 1`] = `
+"type Post @aws_api_key
+{
+  customId: String!
+}"
+`;
+
+exports[`CustomType auth rules duplicate field validation does not issue errors for \`identifier()\` explicit \`id: ID\` field 1`] = `
+"type Post @aws_api_key
+{
+  id: ID!
+}"
+`;
+
+exports[`CustomType auth rules duplicate field validation does not issue errors for \`identifier()\` explicit \`id: string\` field 1`] = `
+"type Post @aws_api_key
+{
+  id: String!
+}"
+`;
+
+exports[`CustomType auth rules duplicate field validation does not issue errors for \`identifier()\` explicit multi-field mixed types PK 1`] = `
+"type Post @aws_api_key
+{
+  idFieldA: String!
+  idFieldB: Int!
+}"
+`;
+
+exports[`CustomType auth rules duplicate field validation does not issue errors for \`identifier()\` explicit multi-field string type PK 1`] = `
+"type Post @aws_api_key
+{
+  idFieldA: String!
+  idFieldB: String!
+}"
+`;
+
 exports[`CustomType transform Explicit CustomType - required 1`] = `
 "type Post @model @auth(rules: [{allow: public, provider: apiKey}])
 {
   location: Location!
 }
 
-type Location 
+type Location @aws_api_key
 {
   lat: Float
   long: Float
@@ -19,7 +455,7 @@ exports[`CustomType transform Explicit CustomType - with auth 1`] = `
   location: Location @auth(rules: [{allow: owner, ownerField: "owner"}])
 }
 
-type Location 
+type Location @aws_api_key
 {
   lat: Float
   long: Float
@@ -32,7 +468,7 @@ exports[`CustomType transform Explicit CustomType 1`] = `
   location: Location
 }
 
-type Location 
+type Location @aws_api_key
 {
   lat: Float
   long: Float
@@ -50,13 +486,13 @@ enum PostStatus {
   published
 }
 
-type Meta 
+type Meta @aws_api_key
 {
   status: PostStatus
   nestedMeta: AltMeta
 }
 
-type AltMeta 
+type AltMeta @aws_api_key
 {
   field1: String
 }"
@@ -73,7 +509,7 @@ enum PostStatus {
   published
 }
 
-type Meta 
+type Meta @aws_api_key
 {
   status: PostStatus
   publishedDate: AWSDate
@@ -86,7 +522,7 @@ exports[`CustomType transform Explicit CustomType nests implicit CustomType 1`]
   meta: Meta
 }
 
-type Meta 
+type Meta @aws_api_key
 {
   status: MetaStatus
   nestedMeta: MetaNestedMeta
@@ -97,7 +533,7 @@ enum MetaStatus {
   published
 }
 
-type MetaNestedMeta 
+type MetaNestedMeta @aws_api_key
 {
   field1: String
 }"
@@ -109,7 +545,7 @@ exports[`CustomType transform Explicit CustomType nests implicit enum type 1`] =
   meta: Meta
 }
 
-type Meta 
+type Meta @aws_api_key
 {
   status: MetaStatus
   publishedDate: AWSDate
@@ -127,7 +563,7 @@ exports[`CustomType transform Implicit CustomType 1`] = `
   location: PostLocation
 }
 
-type PostLocation 
+type PostLocation @aws_api_key
 {
   lat: Float
   long: Float
@@ -140,12 +576,12 @@ exports[`CustomType transform Implicit CustomType nests explicit CustomType 1`]
   meta: PostMeta
 }
 
-type AltMeta 
+type AltMeta @aws_api_key
 {
   field1: String
 }
 
-type PostMeta 
+type PostMeta @aws_api_key
 {
   status: PostMetaStatus
   nestedMeta: AltMeta
@@ -168,7 +604,7 @@ enum PostStatus {
   published
 }
 
-type PostMeta 
+type PostMeta @aws_api_key
 {
   status: PostStatus
   publishedDate: AWSDate
@@ -181,7 +617,7 @@ exports[`CustomType transform Implicit CustomType nests implicit CustomType 1`]
   meta: PostMeta
 }
 
-type PostMeta 
+type PostMeta @aws_api_key
 {
   status: PostMetaStatus
   nestedMeta: PostMetaNestedMeta
@@ -192,7 +628,7 @@ enum PostMetaStatus {
   published
 }
 
-type PostMetaNestedMeta 
+type PostMetaNestedMeta @aws_api_key
 {
   field1: String
 }"
@@ -204,7 +640,7 @@ exports[`CustomType transform Implicit CustomType nests implicit enum type 1`] =
   meta: PostMeta
 }
 
-type PostMeta 
+type PostMeta @aws_api_key
 {
   status: PostMetaStatus
   publishedDate: AWSDate
diff --git a/packages/data-schema/src/CustomOperation.ts b/packages/data-schema/src/CustomOperation.ts
index 4f79fe524..807da288f 100644
--- a/packages/data-schema/src/CustomOperation.ts
+++ b/packages/data-schema/src/CustomOperation.ts
@@ -353,6 +353,7 @@ type AsyncFunctionCustomOperation<T extends CustomOperationParamShape> =
   >;
 
 type EventInvocationResponseCustomType = CustomType<{
+  authorization: [];
   fields: {
     success: ModelField<boolean, 'required', undefined>;
   };
diff --git a/packages/data-schema/src/CustomType.ts b/packages/data-schema/src/CustomType.ts
index 27d166984..7f714d5ed 100644
--- a/packages/data-schema/src/CustomType.ts
+++ b/packages/data-schema/src/CustomType.ts
@@ -1,3 +1,9 @@
+import { 
+  type Authorization,
+  type BaseAllowModifier, 
+  type AnyAuthorization,
+  allow,
+} from './Authorization';
 import type { Brand } from './util';
 import type { InternalField, BaseModelField } from './ModelField';
 import type { RefType } from './RefType';
@@ -26,6 +32,7 @@ type InternalModelFields = Record<string, InternalField>;
 type CustomTypeData = {
   fields: CustomTypeFields;
   type: 'customType';
+  authorization: Authorization<any, any, any>[];
 };
 
 type InternalCustomTypeData = CustomTypeData & {
@@ -34,6 +41,7 @@ type InternalCustomTypeData = CustomTypeData & {
 
 export type CustomTypeParamShape = {
   fields: CustomTypeFields;
+  authorization: Authorization<any, any, any>[];
 };
 
 /**
@@ -41,8 +49,26 @@ export type CustomTypeParamShape = {
  *
  * @param T - The shape of the custom type container
  */
-export type CustomType<T extends CustomTypeParamShape> = T &
-  Brand<'customType'>;
+export type CustomType<T extends CustomTypeParamShape> = T & Brand<'customType'> & {
+    /**
+     * Configures authorization rules for public, signed-in user, per user, and per user group data access
+     *
+     * @param callback A function that receives an allow modifier to define authorization rules
+     * @returns A ModelType instance with updated authorization rules
+     *
+     * @example
+     * a.customType({}).authorization((allow) => [
+     *   allow.guest(),
+     *   allow.publicApiKey(),
+     *   allow.authenticated(),
+     * ])
+     */
+    authorization<AuthRuleType extends AnyAuthorization>(
+      callback: (
+        allow: BaseAllowModifier,
+      ) => AuthRuleType | AuthRuleType[],
+    ): CustomType<T>;
+  };
 
 /**
  * Internal representation of CustomType that exposes the `data` property.
@@ -56,9 +82,20 @@ function _customType<T extends CustomTypeParamShape>(fields: T['fields']) {
   const data: CustomTypeData = {
     fields,
     type: 'customType',
+    authorization: [],
   };
 
-  return { data } as InternalCustomType as CustomType<T>;
+  const builder = {
+    authorization<AuthRuleType extends AnyAuthorization>(callback: (allow: BaseAllowModifier) => AuthRuleType | AuthRuleType[]) {
+      const { resource: _, ...rest } = allow;
+      const rules = callback(rest);
+      data.authorization = Array.isArray(rules) ? rules : [rules];
+
+      return this;
+    },
+  }
+
+  return { ...builder, data } as InternalCustomType as CustomType<T>;
 }
 
 /**
@@ -95,6 +132,6 @@ function _customType<T extends CustomTypeParamShape>(fields: T['fields']) {
  */
 export function customType<T extends CustomTypeFields>(
   fields: T,
-): CustomType<{ fields: T }> {
+): CustomType<{ fields: T, authorization: [] }> {
   return _customType(fields);
 }
diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts
index b8b452f29..e1f8c3e66 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -1580,13 +1580,15 @@ const schemaPreprocessor = (
           ),
         );
 
-        let customAuth = '';
+        const { authString } = mapToNativeAppSyncAuthDirectives(mostRelevantAuthRules, false);
+
+        let customAuth = authString;
         if (typeName in customTypeInheritedAuthRules) {
           const { authString } = mapToNativeAppSyncAuthDirectives(
             customTypeInheritedAuthRules[typeName],
             false,
           );
-          customAuth = authString;
+          customAuth += authString;
         }
 
         const authFields = {};

From 47d004454553fb6b8d0d0c5486425656f194c682 Mon Sep 17 00:00:00 2001
From: Pavi <13897936+pavinthan@users.noreply.github.com>
Date: Fri, 21 Feb 2025 19:28:56 +0000
Subject: [PATCH 2/3] feat(data-schema): trim rules

---
 .../__tests__/__snapshots__/ClientSchema.test.ts.snap       | 4 ++--
 .../__tests__/__snapshots__/CustomOperations.test.ts.snap   | 6 +++---
 packages/data-schema/src/SchemaProcessor.ts                 | 4 ++--
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
index e389d27a9..4ffe9c0eb 100644
--- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
@@ -26,7 +26,7 @@ type Query {
 `;
 
 exports[`ai routes conversations 1`] = `
-"type Profile @aws_api_key@aws_api_key
+"type Profile @aws_api_key @aws_api_key
 {
   value: Int
   unit: String
@@ -257,7 +257,7 @@ enum PostStatus {
   published
 }
 
-type PostMeta @aws_cognito_user_pools@aws_cognito_user_pools
+type PostMeta @aws_cognito_user_pools @aws_cognito_user_pools
 {
   viewCount: Int
   approvedOn: AWSDate
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
index 4851c890e..5717d63b6 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
@@ -606,7 +606,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema Schema w model, custom query, mutation, and subscription and ref of custom type with array modifier 1`] = `
-"type PostCustomType @aws_api_key@aws_api_key
+"type PostCustomType @aws_api_key @aws_api_key
 {
   title: String
 }
@@ -644,7 +644,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema custom subscriptions Custom subscription where .for() resource has a CustomType return type 1`] = `
-"type CreateCustomTypePostReturnType @aws_api_key@aws_api_key
+"type CreateCustomTypePostReturnType @aws_api_key @aws_api_key
 {
   title: String!
 }
@@ -726,7 +726,7 @@ exports[`CustomOperation transform dynamo schema handlers a.handler.custom a.han
   title: String
 }
 
-type GetPostDetailsReturnType @aws_api_key@aws_api_key
+type GetPostDetailsReturnType @aws_api_key @aws_api_key
 {
   
 }
diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts
index e1f8c3e66..ca7f12a28 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -1588,7 +1588,7 @@ const schemaPreprocessor = (
             customTypeInheritedAuthRules[typeName],
             false,
           );
-          customAuth += authString;
+          customAuth += ` ${authString}`;
         }
 
         const authFields = {};
@@ -1613,7 +1613,7 @@ const schemaPreprocessor = (
 
         const joined = gqlFields.join('\n  ');
 
-        const model = `type ${typeName} ${customAuth}\n{\n  ${joined}\n}`;
+        const model = `type ${typeName} ${customAuth.trim()}\n{\n  ${joined}\n}`;
         gqlModels.push(model);
       } else if (isCustomOperation(typeDef)) {
         // TODO: add generation route logic.

From 770ba6e6f4c31730bdde984b8b50bc1be5aa2185 Mon Sep 17 00:00:00 2001
From: Pavi <13897936+pavinthan@users.noreply.github.com>
Date: Fri, 21 Feb 2025 19:38:52 +0000
Subject: [PATCH 3/3] feat(data-schema): fix duplicates

---
 .../__tests__/__snapshots__/ClientSchema.test.ts.snap     | 4 ++--
 .../__tests__/__snapshots__/CustomOperations.test.ts.snap | 6 +++---
 packages/data-schema/src/SchemaProcessor.ts               | 8 +++++---
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
index 4ffe9c0eb..97e4d41a1 100644
--- a/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/ClientSchema.test.ts.snap
@@ -26,7 +26,7 @@ type Query {
 `;
 
 exports[`ai routes conversations 1`] = `
-"type Profile @aws_api_key @aws_api_key
+"type Profile @aws_api_key
 {
   value: Int
   unit: String
@@ -257,7 +257,7 @@ enum PostStatus {
   published
 }
 
-type PostMeta @aws_cognito_user_pools @aws_cognito_user_pools
+type PostMeta @aws_cognito_user_pools
 {
   viewCount: Int
   approvedOn: AWSDate
diff --git a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
index 5717d63b6..8d2d1a6ff 100644
--- a/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
+++ b/packages/data-schema/__tests__/__snapshots__/CustomOperations.test.ts.snap
@@ -606,7 +606,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema Schema w model, custom query, mutation, and subscription and ref of custom type with array modifier 1`] = `
-"type PostCustomType @aws_api_key @aws_api_key
+"type PostCustomType @aws_api_key
 {
   title: String
 }
@@ -644,7 +644,7 @@ type Subscription {
 `;
 
 exports[`CustomOperation transform dynamo schema custom subscriptions Custom subscription where .for() resource has a CustomType return type 1`] = `
-"type CreateCustomTypePostReturnType @aws_api_key @aws_api_key
+"type CreateCustomTypePostReturnType @aws_api_key
 {
   title: String!
 }
@@ -726,7 +726,7 @@ exports[`CustomOperation transform dynamo schema handlers a.handler.custom a.han
   title: String
 }
 
-type GetPostDetailsReturnType @aws_api_key @aws_api_key
+type GetPostDetailsReturnType @aws_api_key
 {
   
 }
diff --git a/packages/data-schema/src/SchemaProcessor.ts b/packages/data-schema/src/SchemaProcessor.ts
index ca7f12a28..a796ec71f 100644
--- a/packages/data-schema/src/SchemaProcessor.ts
+++ b/packages/data-schema/src/SchemaProcessor.ts
@@ -1582,13 +1582,13 @@ const schemaPreprocessor = (
 
         const { authString } = mapToNativeAppSyncAuthDirectives(mostRelevantAuthRules, false);
 
-        let customAuth = authString;
+        let customAuth = authString.split(' ');
         if (typeName in customTypeInheritedAuthRules) {
           const { authString } = mapToNativeAppSyncAuthDirectives(
             customTypeInheritedAuthRules[typeName],
             false,
           );
-          customAuth += ` ${authString}`;
+          customAuth = customAuth.concat(authString.split(' '));
         }
 
         const authFields = {};
@@ -1613,7 +1613,9 @@ const schemaPreprocessor = (
 
         const joined = gqlFields.join('\n  ');
 
-        const model = `type ${typeName} ${customAuth.trim()}\n{\n  ${joined}\n}`;
+        const customAuthRules = Array.from(new Set(customAuth)).join(' ').trim()
+
+        const model = `type ${typeName} ${customAuthRules}\n{\n  ${joined}\n}`;
         gqlModels.push(model);
       } else if (isCustomOperation(typeDef)) {
         // TODO: add generation route logic.