fix(redshift-alpha): extract tableName from custom resource functions (#32452)

### Issue # (if applicable)

Closes 
 - #31817
 - #29231

### Reason for this change



This regression was introduced in PR#24308, which added support for executing the `ALT TABLE` SQL statement when the table name is changed in the CDK construct. The root cause is in the modification of the physical ID of the custom resource Lambda function. Previously, this Lambda function was returning the `tableName` after table is created. However, the physical ID was updated to a combined ID in the format `clusterId:databaseName:tableNamePrefix:stackIdSuffix`.

This change breaks the logic that depends on the return value being used as the `tableName`. For example, in the reported issue, the integration test relies on this value to grant permissions to specific users. The new combined ID format causes SQL statement errors, as described in the issue.

### Description of changes



The fix involves adding logic to detect whether the resource's physical ID is in the new format. If it is, the logic reconstructs the `tableName` from the combined ID. This ensures that Lambda functions referencing the `tableName` receive the correct value.

### Description of how you validated changes



- rerun unit test cases
- deployed integration test cases

### Checklist
- [X] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md)

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: 5d

packages/@aws-cdk/aws-redshift-alpha/lib/private/database-query-provider/privileges.ts

@@ -11,17 +11,18 @@ export async function handler(props: UserTablePrivilegesHandlerProps & ClusterPr
   const clusterProps = props;
 
   if (event.RequestType === 'Create') {
-    await grantPrivileges(username, tablePrivileges, clusterProps);
+    await grantPrivileges(username, tablePrivileges, clusterProps, event.StackId);
     return { PhysicalResourceId: makePhysicalId(username, clusterProps, event.RequestId) };
   } else if (event.RequestType === 'Delete') {
-    await revokePrivileges(username, tablePrivileges, clusterProps);
+    await revokePrivileges(username, tablePrivileges, clusterProps, event.StackId);
     return;
   } else if (event.RequestType === 'Update') {
     const { replace } = await updatePrivileges(
       username,
       tablePrivileges,
       clusterProps,
       event.OldResourceProperties as unknown as UserTablePrivilegesHandlerProps & ClusterProps,
+      event.StackId,
     );
     const physicalId = replace ? makePhysicalId(username, clusterProps, event.RequestId) : event.PhysicalResourceId;
     return { PhysicalResourceId: physicalId };
@@ -31,19 +32,35 @@ export async function handler(props: UserTablePrivilegesHandlerProps & ClusterPr
   }
 }
 
-async function revokePrivileges(username: string, tablePrivileges: TablePrivilege[], clusterProps: ClusterProps) {
+async function revokePrivileges(
+  username: string,
+  tablePrivileges: TablePrivilege[],
+  clusterProps: ClusterProps,
+  stackId: string,
+) {
   // Limited by human input
   // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
   await Promise.all(tablePrivileges.map(({ tableName, actions }) => {
-    return executeStatement(`REVOKE ${actions.join(', ')} ON ${tableName} FROM ${username}`, clusterProps);
+    return executeStatement(
+      `REVOKE ${actions.join(', ')} ON ${normalizedTableName(tableName, stackId)} FROM ${username}`,
+      clusterProps,
+    );
   }));
 }
 
-async function grantPrivileges(username: string, tablePrivileges: TablePrivilege[], clusterProps: ClusterProps) {
+async function grantPrivileges(
+  username: string,
+  tablePrivileges: TablePrivilege[],
+  clusterProps: ClusterProps,
+  stackId: string,
+) {
   // Limited by human input
   // eslint-disable-next-line @cdklabs/promiseall-no-unbounded-parallelism
   await Promise.all(tablePrivileges.map(({ tableName, actions }) => {
-    return executeStatement(`GRANT ${actions.join(', ')} ON ${tableName} TO ${username}`, clusterProps);
+    return executeStatement(
+      `GRANT ${actions.join(', ')} ON ${normalizedTableName(tableName, stackId)} TO ${username}`,
+      clusterProps,
+    );
   }));
 }
 
@@ -52,16 +69,17 @@ async function updatePrivileges(
   tablePrivileges: TablePrivilege[],
   clusterProps: ClusterProps,
   oldResourceProperties: UserTablePrivilegesHandlerProps & ClusterProps,
+  stackId: string,
 ): Promise<{ replace: boolean }> {
   const oldClusterProps = oldResourceProperties;
   if (clusterProps.clusterName !== oldClusterProps.clusterName || clusterProps.databaseName !== oldClusterProps.databaseName) {
-    await grantPrivileges(username, tablePrivileges, clusterProps);
+    await grantPrivileges(username, tablePrivileges, clusterProps, stackId);
     return { replace: true };
   }
 
   const oldUsername = oldResourceProperties.username;
   if (oldUsername !== username) {
-    await grantPrivileges(username, tablePrivileges, clusterProps);
+    await grantPrivileges(username, tablePrivileges, clusterProps, stackId);
     return { replace: true };
   }
 
@@ -72,7 +90,7 @@ async function updatePrivileges(
     ))
   ));
   if (tablesToRevoke.length > 0) {
-    await revokePrivileges(username, tablesToRevoke, clusterProps);
+    await revokePrivileges(username, tablesToRevoke, clusterProps, stackId);
   }
 
   const tablesToGrant = tablePrivileges.filter(({ tableId, tableName, actions }) => {
@@ -85,8 +103,21 @@ async function updatePrivileges(
     return tableAdded || actionsAdded;
   });
   if (tablesToGrant.length > 0) {
-    await grantPrivileges(username, tablesToGrant, clusterProps);
+    await grantPrivileges(username, tablesToGrant, clusterProps, stackId);
   }
 
   return { replace: false };
 }
+
+/**
+ * We need this normalization logic because some of the `TableName` values
+ * are physical IDs generated in the {@link makePhysicalId} function.
+ * */
+const normalizedTableName = (tableName: string, stackId: string): string => {
+  const segments = tableName.split(':');
+  const suffix = segments.slice(-1);
+  if (suffix != null && stackId.endsWith(suffix[0])) {
+    return segments.slice(-2)[0] ?? tableName;
+  }
+  return tableName;
+};

packages/@aws-cdk/aws-redshift-alpha/lib/private/database-query-provider/table.ts

@@ -15,7 +15,7 @@ export async function handler(props: TableAndClusterProps, event: AWSLambda.Clou
 
   if (event.RequestType === 'Create') {
     tableName = await createTable(tableNamePrefix, getTableNameSuffix(props.tableName.generateSuffix), tableColumns, tableAndClusterProps);
-    return { PhysicalResourceId: makePhysicalId(tableNamePrefix, tableAndClusterProps, event.StackId.substring(event.StackId.length - 12)) };
+    return { PhysicalResourceId: makePhysicalId(tableName, tableAndClusterProps, event.StackId.substring(event.StackId.length - 12)) };
   } else if (event.RequestType === 'Delete') {
     await dropTable(
       event.PhysicalResourceId.includes(event.StackId.substring(event.StackId.length - 12)) ? tableName : event.PhysicalResourceId,

packages/@aws-cdk/aws-redshift-alpha/lib/private/privileges.ts

@@ -5,7 +5,7 @@ import { ITable, TableAction } from '../table';
 import { IUser } from '../user';
 import { DatabaseQuery } from './database-query';
 import { HandlerName } from './database-query-provider/handler-name';
-import { TablePrivilege as SerializedTablePrivilege, UserTablePrivilegesHandlerProps } from './handler-props';
+import { UserTablePrivilegesHandlerProps } from './handler-props';
 
 /**
  * The Redshift table and action that make up a privilege that can be granted to a Redshift user.
@@ -61,35 +61,14 @@ export class UserTablePrivileges extends Construct {
       properties: {
         username: props.user.username,
         tablePrivileges: cdk.Lazy.any({
-          produce: () => {
-            const reducedPrivileges = this.privileges.reduce((privileges, { table, actions }) => {
-              const tableId = table.node.id;
-              if (!(tableId in privileges)) {
-                privileges[tableId] = {
-                  tableName: table.tableName,
-                  actions: [],
-                };
-              }
-              actions = actions.concat(privileges[tableId].actions);
-              if (actions.includes(TableAction.ALL)) {
-                actions = [TableAction.ALL];
-              }
-              if (actions.includes(TableAction.UPDATE) || actions.includes(TableAction.DELETE)) {
-                actions.push(TableAction.SELECT);
-              }
-              privileges[tableId] = {
-                tableName: table.tableName,
-                actions: Array.from(new Set(actions)),
-              };
-              return privileges;
-            }, {} as { [key: string]: { tableName: string; actions: TableAction[] } });
-            const serializedPrivileges: SerializedTablePrivilege[] = Object.entries(reducedPrivileges).map(([tableId, config]) => ({
-              tableId,
-              tableName: config.tableName,
-              actions: config.actions.map(action => TableAction[action]),
-            }));
-            return serializedPrivileges;
-          },
+          produce: () =>
+            Object.entries(groupPrivilegesByTable(this.privileges))
+              .map(([tableId, tablePrivileges]) => ({
+                tableId,
+                // The first element always exists since the groupBy element is at least a singleton.
+                tableName: tablePrivileges[0]!.table.tableName,
+                actions: unifyTableActions(tablePrivileges).map(action => TableAction[action]),
+              })),
         }) as any,
       },
     });
@@ -102,3 +81,29 @@ export class UserTablePrivileges extends Construct {
     this.privileges.push({ table, actions });
   }
 }
+
+const unifyTableActions = (tablePrivileges: TablePrivilege[]): TableAction[] => {
+  const set = new Set<TableAction>(tablePrivileges.flatMap(x => x.actions));
+
+  if (set.has(TableAction.ALL)) {
+    return [TableAction.ALL];
+  }
+
+  if (set.has(TableAction.UPDATE) || set.has(TableAction.DELETE)) {
+    set.add(TableAction.SELECT);
+  }
+
+  return [...set];
+};
+
+const groupPrivilegesByTable = (privileges: TablePrivilege[]): Record<string, TablePrivilege[]> => {
+  return privileges.reduce((grouped, privilege) => {
+    const { table } = privilege;
+    const tableId = table.node.id;
+    const tablePrivileges = grouped[tableId] ?? [];
+    return {
+      ...grouped,
+      [tableId]: [...tablePrivileges, privilege],
+    };
+  }, {} as Record<string, TablePrivilege[]>);
+};

packages/@aws-cdk/aws-redshift-alpha/lib/table.ts

@@ -272,7 +272,7 @@ export class Table extends TableBase {
       properties: {
         tableName: {
           prefix: props.tableName ?? cdk.Names.uniqueId(this),
-          generateSuffix: !props.tableName ? 'true' : 'false',
+          generateSuffix: (props.tableName == null).toString(),
         },
         tableColumns: this.tableColumns,
         distStyle: props.distStyle,
@@ -282,7 +282,7 @@ export class Table extends TableBase {
       },
     });
 
-    this.tableName = this.resource.ref;
+    this.tableName = props.tableName ?? this.resource.ref;
   }
 
   /**

packages/@aws-cdk/aws-redshift-alpha/test/database-query-provider/privileges.test.ts

@@ -40,6 +40,7 @@ jest.mock('@aws-sdk/client-redshift-data', () => {
 });
 
 import { handler as managePrivileges } from '../../lib/private/database-query-provider/privileges';
+import { makePhysicalId } from '../../lib/private/database-query-provider/util';
 
 beforeEach(() => {
   jest.clearAllMocks();
@@ -61,6 +62,31 @@ describe('create', () => {
       Sql: `GRANT INSERT, SELECT ON ${tableName} TO ${username}`,
     }));
   });
+
+  test('serializes properties in statement when tableName in physical resource ID', async () => {
+    const properties = {
+      ...resourceProperties,
+      tablePrivileges: [{
+        tableId,
+        tableName: `${makePhysicalId(tableName, resourceProperties, requestId)}`,
+        actions,
+      }],
+    };
+
+    const event = {
+      ...baseEvent,
+      ResourceProperties: properties,
+      StackId: 'xxxxx:' + requestId,
+    };
+
+    await expect(managePrivileges(properties, event)).resolves.toEqual({
+      PhysicalResourceId: 'clusterName:databaseName:username:requestId',
+    });
+
+    expect(mockExecuteStatement).toHaveBeenCalledWith(expect.objectContaining({
+      Sql: `GRANT INSERT, SELECT ON ${tableName} TO ${username}`,
+    }));
+  });
 });
 
 describe('delete', () => {
@@ -79,6 +105,29 @@ describe('delete', () => {
       Sql: `REVOKE INSERT, SELECT ON ${tableName} FROM ${username}`,
     }));
   });
+
+  test('serializes properties in statement when tableName in physical resource ID', async () => {
+    const properties = {
+      ...resourceProperties,
+      tablePrivileges: [{
+        tableId,
+        tableName: `${makePhysicalId(tableName, resourceProperties, requestId)}`,
+        actions,
+      }],
+    };
+
+    const event = {
+      ...baseEvent,
+      ResourceProperties: properties,
+      StackId: 'xxxxx:' + requestId,
+    };
+
+    await managePrivileges(properties, event);
+
+    expect(mockExecuteStatement).toHaveBeenCalledWith(expect.objectContaining({
+      Sql: `REVOKE INSERT, SELECT ON ${tableName} FROM ${username}`,
+    }));
+  });
 });
 
 describe('update', () => {
@@ -244,4 +293,50 @@ describe('update', () => {
       Sql: expect.stringMatching(new RegExp(`.+ ON ${tableName} FROM ${username}`)),
     }));
   });
+
+  test('serializes properties in grant statement when tableName in physical resource ID', async () => {
+    const properties = {
+      ...resourceProperties,
+      tablePrivileges: [{
+        tableId,
+        tableName: `${makePhysicalId(tableName, resourceProperties, requestId)}`,
+        actions,
+      }],
+    };
+
+    const newEvent = {
+      ...event,
+      ResourceProperties: properties,
+      StackId: 'xxxxx:' + requestId,
+    };
+
+    await managePrivileges(properties, newEvent);
+
+    expect(mockExecuteStatement).toHaveBeenCalledWith(expect.objectContaining({
+      Sql: `GRANT INSERT, SELECT ON ${tableName} TO ${username}`,
+    }));
+  });
+
+  test('serializes properties in drop statement when tableName in physical resource ID', async () => {
+    const properties = {
+      ...resourceProperties,
+      tablePrivileges: [{
+        tableId,
+        tableName: `${makePhysicalId(tableName, resourceProperties, requestId)}`,
+        actions: ['DROP'],
+      }],
+    };
+
+    const newEvent = {
+      ...event,
+      ResourceProperties: properties,
+      StackId: 'xxxxx:' + requestId,
+    };
+
+    await managePrivileges(properties, newEvent);
+
+    expect(mockExecuteStatement).toHaveBeenCalledWith(expect.objectContaining({
+      Sql: `REVOKE INSERT, SELECT ON ${tableName} FROM ${username}`,
+    }));
+  });
 });

packages/@aws-cdk/aws-redshift-alpha/test/database-query-provider/table.test.ts

@@ -61,7 +61,7 @@ describe('create', () => {
     const event = baseEvent;
 
     await expect(manageTable(resourceProperties, event)).resolves.toEqual({
-      PhysicalResourceId: 'clusterName:databaseName:tableNamePrefix:111111111111',
+      PhysicalResourceId: 'clusterName:databaseName:tableNamePrefix111111111111:111111111111',
     });
     expect(mockExecuteStatement).toHaveBeenCalledWith(expect.objectContaining({
       Sql: `CREATE TABLE ${tableNamePrefix}${stackIdTruncated} (col1 varchar(1))`,