fix(graphql-model-transformer): derive imported DynamoDB table key schema from @PrimaryKey

Imported Amplify-managed DynamoDB tables (Gen 1 -> Gen 2 migration via
migratedAmplifyGen1DynamoDbTableMappings) hardcoded the table partition key to
`id` when synthesizing the Custom::ImportedAmplifyDynamoDBTable resource. For
owned tables the @PrimaryKey transformer corrects the key schema downstream via
a CloudFormation property override, but that correction does not reach the
TableManager import-validation path, which consumes the initial construct
properties directly.

As a result, any model declaring a custom @PrimaryKey (partition key other than
`id`, and/or a sort key) failed import validation with "Imported table
properties did not match the expected table properties", rolling back the whole
branch stack (a failed CREATE that must then be deleted manually before retry).
GSIs were unaffected because they are already derived from the model.

Fix: in AmplifyDynamoModelResourceGenerator.createModelTable, for imported
tables only, derive the partition key (and single/composite sort key) from the
model's @PrimaryKey via getPrimaryKeyFieldNodes, mirroring how GSIs are derived.
Owned tables are unchanged (still `id`, corrected downstream as before).

Adds regression tests: imported table with a custom @PrimaryKey partition key,
with a sort key, and the default-`id` no-regression case.

Fixes #3489
Related: aws-amplify/amplify-backend#3281

Author: sharonyajain

packages/amplify-graphql-model-transformer/src/__tests__/amplify-dynamodb-table-generator.test.ts

@@ -7,6 +7,7 @@ import {
 import { parse } from 'graphql';
 import { ModelTransformer } from '../graphql-model-transformer';
 import { SearchableModelTransformer } from '@aws-amplify/graphql-searchable-transformer';
+import { PrimaryKeyTransformer } from '@aws-amplify/graphql-index-transformer';
 import { CUSTOM_DDB_CFN_TYPE, CUSTOM_IMPORTED_DDB_CFN_TYPE } from '../resources/amplify-dynamodb-table/amplify-dynamodb-table-construct';
 import { ITERATIVE_TABLE_STACK_NAME } from '../resources/amplify-dynamodb-table/amplify-dynamo-model-resource-generator';
 
@@ -171,4 +172,101 @@ describe('ModelTransformer:', () => {
     validateModelSchema(parse(out.schema));
     parse(out.schema);
   });
+
+  it('derives the imported table key schema from a custom @primaryKey (partition key other than id)', async () => {
+    // An imported Amplify-managed table whose model declares a custom @primaryKey must synthesize
+    // its expected keySchema/attributeDefinitions from the model key, not the hardcoded `id`.
+    // Otherwise the TableManager import validation rejects the (correct) live Gen 1 table with
+    // "Imported table properties did not match the expected table properties" and the migration
+    // stack rolls back.
+    const validSchema = `
+      type AuthSessionToken @model {
+        userAuthenticating: ID! @primaryKey
+        sessionToken: String!
+      }
+    `;
+    const out = testTransform({
+      schema: validSchema,
+      transformers: [new ModelTransformer(), new PrimaryKeyTransformer()],
+      dataSourceStrategies: {
+        AuthSessionToken: {
+          dbType: 'DYNAMODB' as const,
+          provisionStrategy: 'IMPORTED_AMPLIFY_TABLE' as const,
+          tableName: 'AuthSessionToken-myApiId-myEnv',
+        },
+      },
+    });
+    expect(out).toBeDefined();
+    const stack = out.stacks['AuthSessionToken'];
+    expect(stack).toBeDefined();
+    const table = stack.Resources?.AuthSessionTokenTable;
+    expect(table).toBeDefined();
+    expect(table.Type).toBe(CUSTOM_IMPORTED_DDB_CFN_TYPE);
+    expect(table.Properties.isImported).toBe(true);
+    // Key schema and attribute definitions must reflect the custom @primaryKey, not `id`.
+    expect(table.Properties.keySchema).toEqual([{ attributeName: 'userAuthenticating', keyType: 'HASH' }]);
+    expect(table.Properties.attributeDefinitions).toEqual([{ attributeName: 'userAuthenticating', attributeType: 'S' }]);
+    validateModelSchema(parse(out.schema));
+  });
+
+  it('derives the imported table key schema from a custom @primaryKey with a sort key', async () => {
+    const validSchema = `
+      type Event @model {
+        tenantId: ID! @primaryKey(sortKeyFields: ["createdAt"])
+        createdAt: String!
+        name: String
+      }
+    `;
+    const out = testTransform({
+      schema: validSchema,
+      transformers: [new ModelTransformer(), new PrimaryKeyTransformer()],
+      dataSourceStrategies: {
+        Event: {
+          dbType: 'DYNAMODB' as const,
+          provisionStrategy: 'IMPORTED_AMPLIFY_TABLE' as const,
+          tableName: 'Event-myApiId-myEnv',
+        },
+      },
+    });
+    const table = out.stacks['Event'].Resources?.EventTable;
+    expect(table).toBeDefined();
+    expect(table.Type).toBe(CUSTOM_IMPORTED_DDB_CFN_TYPE);
+    expect(table.Properties.keySchema).toEqual([
+      { attributeName: 'tenantId', keyType: 'HASH' },
+      { attributeName: 'createdAt', keyType: 'RANGE' },
+    ]);
+    expect(table.Properties.attributeDefinitions).toEqual(
+      expect.arrayContaining([
+        { attributeName: 'tenantId', attributeType: 'S' },
+        { attributeName: 'createdAt', attributeType: 'S' },
+      ]),
+    );
+    validateModelSchema(parse(out.schema));
+  });
+
+  it('keeps the default id key schema for an imported table without a custom @primaryKey', async () => {
+    const validSchema = `
+      type Note @model {
+        id: ID!
+        body: String
+      }
+    `;
+    const out = testTransform({
+      schema: validSchema,
+      transformers: [new ModelTransformer(), new PrimaryKeyTransformer()],
+      dataSourceStrategies: {
+        Note: {
+          dbType: 'DYNAMODB' as const,
+          provisionStrategy: 'IMPORTED_AMPLIFY_TABLE' as const,
+          tableName: 'Note-myApiId-myEnv',
+        },
+      },
+    });
+    const table = out.stacks['Note'].Resources?.NoteTable;
+    expect(table).toBeDefined();
+    expect(table.Type).toBe(CUSTOM_IMPORTED_DDB_CFN_TYPE);
+    expect(table.Properties.keySchema).toEqual([{ attributeName: 'id', keyType: 'HASH' }]);
+    expect(table.Properties.attributeDefinitions).toEqual([{ attributeName: 'id', attributeType: 'S' }]);
+    validateModelSchema(parse(out.schema));
+  });
 });

packages/amplify-graphql-model-transformer/src/resources/amplify-dynamodb-table/amplify-dynamo-model-resource-generator.ts

@@ -1,8 +1,12 @@
 import * as cdk from 'aws-cdk-lib';
 import { TransformerContextProvider } from '@aws-amplify/graphql-transformer-interfaces';
-import { ModelResourceIDs, ResourceConstants } from 'graphql-transformer-common';
+import { ModelResourceIDs, ResourceConstants, attributeTypeFromScalar } from 'graphql-transformer-common';
 import { ObjectTypeDefinitionNode } from 'graphql';
-import { setResourceName, isImportedAmplifyDynamoDbModelDataSourceStrategy } from '@aws-amplify/graphql-transformer-core';
+import {
+  setResourceName,
+  isImportedAmplifyDynamoDbModelDataSourceStrategy,
+  getPrimaryKeyFieldNodes,
+} from '@aws-amplify/graphql-transformer-core';
 import { AttributeType, StreamViewType, TableEncryption } from 'aws-cdk-lib/aws-dynamodb';
 import { Construct } from 'constructs';
 import { Duration, aws_iam, aws_lambda } from 'aws-cdk-lib';
@@ -172,6 +176,46 @@ export class AmplifyDynamoModelResourceGenerator extends DynamoModelResourceGene
     const isTableImported = isImportedAmplifyDynamoDbModelDataSourceStrategy(strategy);
     const tableName = isTableImported ? strategy.tableName : context.resourceHelper.generateTableName(modelName);
 
+    // Determine the table's key schema (partition key, and sort key when declared).
+    //
+    // For owned (Amplify-managed) tables we default to a `id` partition key; the `@primaryKey`
+    // transformer (see `replaceDdbPrimaryKey` in the index transformer) corrects the key schema
+    // downstream via a CloudFormation property override before the table is created.
+    //
+    // For imported tables that downstream correction does not reach the TableManager import
+    // validation, which consumes the initial construct properties directly. If we left the key
+    // schema hardcoded to `id`, any model declaring a custom `@primaryKey` (partition key other
+    // than `id`, and/or a sort key) would fail import with "Imported table properties did not
+    // match the expected table properties". So for imported tables we derive the key schema from
+    // the model definition up front, mirroring how GSIs are already derived from the model.
+    let partitionKey = {
+      name: 'id',
+      type: AttributeType.STRING,
+    };
+    let sortKey: { name: string; type: AttributeType } | undefined;
+    if (isTableImported) {
+      const [primaryKeyFieldNode, ...sortKeyFieldNodes] = getPrimaryKeyFieldNodes(def);
+      partitionKey = {
+        name: primaryKeyFieldNode.name.value,
+        type: attributeTypeFromScalar(primaryKeyFieldNode.type) === 'N' ? AttributeType.NUMBER : AttributeType.STRING,
+      };
+      if (sortKeyFieldNodes.length === 1) {
+        // A single sort key field maps directly to a sort key attribute of the field's scalar type.
+        sortKey = {
+          name: sortKeyFieldNodes[0].name.value,
+          type: attributeTypeFromScalar(sortKeyFieldNodes[0].type) === 'N' ? AttributeType.NUMBER : AttributeType.STRING,
+        };
+      } else if (sortKeyFieldNodes.length > 1) {
+        // Composite sort keys are stored as a single string attribute whose name is the sort key
+        // field names joined by the model composite key separator (matches `getSortKeyName` in the
+        // index transformer's `replaceDdbPrimaryKey`).
+        sortKey = {
+          name: sortKeyFieldNodes.map((node) => node.name.value).join(ModelResourceIDs.ModelCompositeKeySeparator()),
+          type: AttributeType.STRING,
+        };
+      }
+    }
+
     // Add parameters.
     const { readIops, writeIops, billingMode, pointInTimeRecovery } = this.createDynamoDBParameters(scope, true);
 
@@ -194,10 +238,8 @@ export class AmplifyDynamoModelResourceGenerator extends DynamoModelResourceGene
       allowDestructiveGraphqlSchemaUpdates: context.transformParameters.allowDestructiveGraphqlSchemaUpdates,
       replaceTableUponGsiUpdate: context.transformParameters.replaceTableUponGsiUpdate,
       tableName,
-      partitionKey: {
-        name: 'id',
-        type: AttributeType.STRING,
-      },
+      partitionKey,
+      ...(sortKey ? { sortKey } : undefined),
       stream: StreamViewType.NEW_AND_OLD_IMAGES,
       encryption: TableEncryption.DEFAULT,
       removalPolicy,