feat(data-schema): add support for GSI projection types

.changeset/red-tips-watch.md

@@ -0,0 +1,5 @@
+---
+'@aws-amplify/data-schema': minor
+---
+
+add support for GSI projection types

packages/data-schema/__tests__/ModelIndex.test.ts

@@ -1,6 +1,5 @@
 import { expectTypeTestsToPassAsync } from 'jest-tsd';
-import { ClientSchema, a } from '../src/index';
-import { ExtractModelMeta, Prettify } from '@aws-amplify/data-schema-types';
+import { a } from '../src/index';
 
 // evaluates type defs in corresponding test-d.ts file
 it('should not produce static type errors', async () => {
@@ -75,6 +74,132 @@ describe('secondary index schema generation', () => {
   });
 });
 
+describe('GSI projection functionality', () => {
+  it('generates correct schema for KEYS_ONLY projection', () => {
+    const schema = a
+      .schema({
+        Product: a
+          .model({
+            id: a.id().required(),
+            name: a.string().required(),
+            category: a.string().required(),
+            price: a.float().required(),
+            inStock: a.boolean().required(),
+          })
+          .secondaryIndexes((index) => [
+            index('category').projection('KEYS_ONLY'),
+          ]),
+      })
+      .authorization((allow) => allow.publicApiKey());
+
+    const transformed = schema.transform().schema;
+
+    expect(transformed).toContain('projection: { type: KEYS_ONLY }');
+    expect(transformed).not.toContain('nonKeyAttributes');
+    expect(transformed).toMatchSnapshot();
+  });
+
+  it('generates correct schema for INCLUDE projection with nonKeyAttributes', () => {
+    const schema = a
+      .schema({
+        Product: a
+          .model({
+            id: a.id().required(),
+            name: a.string().required(),
+            category: a.string().required(),
+            price: a.float().required(),
+            inStock: a.boolean().required(),
+          })
+          .secondaryIndexes((index) => [
+            index('category').projection('INCLUDE', ['name', 'price']),
+          ]),
+      })
+      .authorization((allow) => allow.publicApiKey());
+
+    const transformed = schema.transform().schema;
+
+    expect(transformed).toContain(
+      'projection: { type: INCLUDE, nonKeyAttributes: ["name", "price"] }',
+    );
+    expect(transformed).toMatchSnapshot();
+  });
+
+  it('generates correct schema for ALL projection', () => {
+    const schema = a
+      .schema({
+        Product: a
+          .model({
+            id: a.id().required(),
+            name: a.string().required(),
+            category: a.string().required(),
+            price: a.float().required(),
+            inStock: a.boolean().required(),
+          })
+          .secondaryIndexes((index) => [index('category').projection('ALL')]),
+      })
+      .authorization((allow) => allow.publicApiKey());
+
+    const transformed = schema.transform().schema;
+
+    // When projection is ALL and no explicit projection is set, it may be omitted from output
+    expect(transformed).toContain('@index');
+    expect(transformed).not.toContain('nonKeyAttributes');
+    expect(transformed).toMatchSnapshot();
+  });
+
+  it('generates correct schema for multiple indexes with different projection types', () => {
+    const schema = a
+      .schema({
+        Order: a
+          .model({
+            id: a.id().required(),
+            customerId: a.string().required(),
+            status: a.string().required(),
+            total: a.float().required(),
+            createdAt: a.datetime().required(),
+          })
+          .secondaryIndexes((index) => [
+            index('customerId').projection('ALL'),
+            index('status').projection('INCLUDE', ['customerId', 'total']),
+            index('createdAt').projection('KEYS_ONLY'),
+          ]),
+      })
+      .authorization((allow) => allow.publicApiKey());
+
+    const transformed = schema.transform().schema;
+
+    expect(transformed).toContain(
+      'projection: { type: INCLUDE, nonKeyAttributes: ["customerId", "total"] }',
+    );
+    expect(transformed).toContain('projection: { type: KEYS_ONLY }');
+    expect(transformed).toMatchSnapshot();
+  });
+
+  it('generates correct schema without projection (defaults to ALL)', () => {
+    const schema = a
+      .schema({
+        Product: a
+          .model({
+            id: a.id().required(),
+            name: a.string().required(),
+            category: a.string().required(),
+            price: a.float().required(),
+          })
+          .secondaryIndexes((index) => [
+            index('category'), // No projection specified, should default to ALL
+          ]),
+      })
+      .authorization((allow) => allow.publicApiKey());
+
+    const transformed = schema.transform().schema;
+
+    // When no projection is specified, it defaults to ALL and may be omitted from output
+    expect(transformed).toContain('@index');
+    expect(transformed).not.toContain('nonKeyAttributes');
+    expect(transformed).toMatchSnapshot();
+  });
+});
+
 describe('SchemaProcessor validation against secondary indexes', () => {
   it('throws error when a.ref() used as the index partition key points to a non-existing type', () => {
     const schema = a.schema({
@@ -138,9 +263,7 @@ describe('SchemaProcessor validation against secondary indexes', () => {
             content: a.string(),
             status: a.enum(['open', 'in_progress', 'completed']),
           })
-          .secondaryIndexes((index) => [
-            index('status').sortKeys(['title'])
-          ]),
+          .secondaryIndexes((index) => [index('status').sortKeys(['title'])]),
       })
       .authorization((allow) => allow.publicApiKey());
 
@@ -157,7 +280,9 @@ describe('SchemaProcessor validation against secondary indexes', () => {
             status: a.enum(['open', 'in_progress', 'completed']),
           })
           .secondaryIndexes((index) => [
-            index('status').sortKeys(['title']).queryField('userDefinedQueryField')
+            index('status')
+              .sortKeys(['title'])
+              .queryField('userDefinedQueryField'),
           ]),
       })
       .authorization((allow) => allow.publicApiKey());
@@ -175,7 +300,7 @@ describe('SchemaProcessor validation against secondary indexes', () => {
             status: a.enum(['open', 'in_progress', 'completed']),
           })
           .secondaryIndexes((index) => [
-            index('status').sortKeys(['title']).queryField(null)
+            index('status').sortKeys(['title']).queryField(null),
           ]),
       })
       .authorization((allow) => allow.publicApiKey());

packages/data-schema/__tests__/__snapshots__/ModelIndex.test.ts.snap

@@ -1,5 +1,59 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`GSI projection functionality generates correct schema for ALL projection 1`] = `
+"type Product @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  id: ID! @primaryKey
+  name: String!
+  category: String! @index(queryField: "listProductByCategory")
+  price: Float!
+  inStock: Boolean!
+}"
+`;
+
+exports[`GSI projection functionality generates correct schema for INCLUDE projection with nonKeyAttributes 1`] = `
+"type Product @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  id: ID! @primaryKey
+  name: String!
+  category: String! @index(queryField: "listProductByCategory", projection: { type: INCLUDE, nonKeyAttributes: ["name", "price"] })
+  price: Float!
+  inStock: Boolean!
+}"
+`;
+
+exports[`GSI projection functionality generates correct schema for KEYS_ONLY projection 1`] = `
+"type Product @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  id: ID! @primaryKey
+  name: String!
+  category: String! @index(queryField: "listProductByCategory", projection: { type: KEYS_ONLY })
+  price: Float!
+  inStock: Boolean!
+}"
+`;
+
+exports[`GSI projection functionality generates correct schema for multiple indexes with different projection types 1`] = `
+"type Order @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  id: ID! @primaryKey
+  customerId: String! @index(queryField: "listOrderByCustomerId")
+  status: String! @index(queryField: "listOrderByStatus", projection: { type: INCLUDE, nonKeyAttributes: ["customerId", "total"] })
+  total: Float!
+  createdAt: AWSDateTime! @index(queryField: "listOrderByCreatedAt", projection: { type: KEYS_ONLY })
+}"
+`;
+
+exports[`GSI projection functionality generates correct schema without projection (defaults to ALL) 1`] = `
+"type Product @model @auth(rules: [{allow: public, provider: apiKey}])
+{
+  id: ID! @primaryKey
+  name: String!
+  category: String! @index(queryField: "listProductByCategory")
+  price: Float!
+}"
+`;
+
 exports[`SchemaProcessor validation against secondary indexes creates a queryField with a default name 1`] = `
 "type Todo @model @auth(rules: [{allow: public, provider: apiKey}])
 {

packages/data-schema/src/ModelIndex.ts

@@ -2,11 +2,26 @@ import { Brand, brand } from './util';
 
 const brandName = 'modelIndexType';
 
+/**
+ * DynamoDB Global Secondary Index (GSI) projection types.
+ * - `KEYS_ONLY`: Only the index and primary keys are projected into the index.
+ * - `INCLUDE`: In addition to the attributes in KEYS_ONLY, includes specified non-key attributes.
+ * - `ALL`: All attributes from the base table are projected into the index.
+ */
+export type GSIProjectionType = 'KEYS_ONLY' | 'INCLUDE' | 'ALL';
+
+/**
+ * Configuration data for a model's secondary index.
+ */
 export type ModelIndexData = {
   partitionKey: string;
   sortKeys: readonly unknown[];
   indexName: string;
   queryField: string | null;
+  /** The projection type for the GSI. Defaults to 'ALL' if not specified. */
+  projectionType?: GSIProjectionType;
+  /** Non-key attributes to include when projectionType is 'INCLUDE'. */
+  nonKeyAttributes?: readonly string[];
 };
 
 export type InternalModelIndexType = ModelIndexType<any, any, any, any> & {
@@ -36,6 +51,13 @@ export type ModelIndexType<
     >(
       field: QF,
     ): ModelIndexType<MF, PK, SK, QF, K | 'queryField'>;
+    projection<
+      PT extends GSIProjectionType,
+      NKA extends PT extends 'INCLUDE' ? readonly string[] : never = never,
+    >(
+      type: PT,
+      ...args: PT extends 'INCLUDE' ? [nonKeyAttributes: NKA] : []
+    ): ModelIndexType<ModelFieldKeys, PK, SK, QueryField, K | 'projection'>;
   },
   K
 > &
@@ -52,6 +74,8 @@ function _modelIndex<
     sortKeys: [],
     indexName: '',
     queryField: '',
+    projectionType: 'ALL',
+    nonKeyAttributes: undefined,
   };
 
   const builder = {
@@ -70,6 +94,14 @@ function _modelIndex<
 
       return this;
     },
+    projection(type, ...args) {
+      data.projectionType = type;
+      if (type === 'INCLUDE') {
+        data.nonKeyAttributes = args[0];
+      }
+
+      return this;
+    },
     ...brand(brandName),
   } as ModelIndexType<ModelFieldKeys, PK, SK, QueryField>;
 

packages/data-schema/src/SchemaProcessor.ts

@@ -1117,6 +1117,8 @@ const transformedSecondaryIndexesForModel = (
     sortKeys: readonly string[],
     indexName: string,
     queryField: string | null,
+    projectionType?: string,
+    nonKeyAttributes?: readonly string[],
   ): string => {
     for (const keyName of [partitionKey, ...sortKeys]) {
       const field = modelFields[keyName];
@@ -1131,7 +1133,7 @@ const transformedSecondaryIndexesForModel = (
       }
     }
 
-    if (!sortKeys.length && !indexName && !queryField && queryField !== null) {
+    if (!sortKeys.length && !indexName && !queryField && queryField !== null && !projectionType) {
       return `@index(queryField: "${secondaryIndexDefaultQueryField(
         modelName,
         partitionKey,
@@ -1165,13 +1167,23 @@ const transformedSecondaryIndexesForModel = (
       );
     }
 
+    // Add projection attributes if specified
+    if (projectionType && projectionType !== 'ALL') {
+      if (projectionType === 'KEYS_ONLY') {
+        attributes.push(`projection: { type: KEYS_ONLY }`);
+      } else if (projectionType === 'INCLUDE' && nonKeyAttributes?.length) {
+        const nonKeyAttrsStr = nonKeyAttributes.map(attr => `"${attr}"`).join(', ');
+        attributes.push(`projection: { type: INCLUDE, nonKeyAttributes: [${nonKeyAttrsStr}] }`);
+      }
+    }
+
     return `@index(${attributes.join(', ')})`;
   };
 
   return secondaryIndexes.reduce(
     (
       acc: TransformedSecondaryIndexes,
-      { data: { partitionKey, sortKeys, indexName, queryField } },
+      { data: { partitionKey, sortKeys, indexName, queryField, projectionType, nonKeyAttributes } },
     ) => {
       acc[partitionKey] = acc[partitionKey] || [];
       acc[partitionKey].push(
@@ -1180,6 +1192,8 @@ const transformedSecondaryIndexesForModel = (
           sortKeys as readonly string[],
           indexName,
           queryField,
+          projectionType,
+          nonKeyAttributes,
         ),
       );