feat(redshift): Add support for distStyle, distKey, sortStyle and sortKey to Table

packages/@aws-cdk/aws-redshift/README.md

@@ -167,6 +167,34 @@ new Table(this, 'Table', {
 });
 ```
 
+The table can be configured to have distStyle attribute and a distKey column:
+
+```ts fixture=cluster
+new Table(this, 'Table', {
+  tableColumns: [
+    { name: 'col1', dataType: 'varchar(4)', distKey: true }, 
+    { name: 'col2', dataType: 'float' },
+  ],
+  cluster: cluster,
+  databaseName: 'databaseName',
+  distStyle: TableDistStyle.KEY,
+});
+```
+
+The table can also be configured to have sortStyle attribute and sortKey columns:
+
+```ts fixture=cluster
+new Table(this, 'Table', {
+  tableColumns: [
+    { name: 'col1', dataType: 'varchar(4)', sortKey: true }, 
+    { name: 'col2', dataType: 'float', sortKey: true },
+  ],
+  cluster: cluster,
+  databaseName: 'databaseName',
+  sortStyle: TableSortStyle.COMPOUND,
+});
+```
+
 ### Granting Privileges
 
 You can give a user privileges to perform certain actions on a table by using the

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

@@ -1,7 +1,8 @@
 /* eslint-disable-next-line import/no-unresolved */
 import * as AWSLambda from 'aws-lambda';
 import { TablePrivilege, UserTablePrivilegesHandlerProps } from '../handler-props';
-import { ClusterProps, executeStatement, makePhysicalId } from './util';
+import { ClusterProps } from './types';
+import { executeStatement, makePhysicalId } from './util';
 
 export async function handler(props: UserTablePrivilegesHandlerProps & ClusterProps, event: AWSLambda.CloudFormationCustomResourceEvent) {
   const username = props.username;

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

@@ -1,29 +1,29 @@
 /* eslint-disable-next-line import/no-unresolved */
 import * as AWSLambda from 'aws-lambda';
 import { Column } from '../../table';
-import { TableHandlerProps } from '../handler-props';
-import { ClusterProps, executeStatement } from './util';
+import { ClusterProps, TableAndClusterProps, TableSortStyle } from './types';
+import { areColumnsEqual, executeStatement, getDistKeyColumn, getSortKeyColumns } from './util';
 
-export async function handler(props: TableHandlerProps & ClusterProps, event: AWSLambda.CloudFormationCustomResourceEvent) {
+export async function handler(props: TableAndClusterProps, event: AWSLambda.CloudFormationCustomResourceEvent) {
   const tableNamePrefix = props.tableName.prefix;
   const tableNameSuffix = props.tableName.generateSuffix === 'true' ? `${event.RequestId.substring(0, 8)}` : '';
   const tableColumns = props.tableColumns;
-  const clusterProps = props;
+  const tableAndClusterProps = props;
 
   if (event.RequestType === 'Create') {
-    const tableName = await createTable(tableNamePrefix, tableNameSuffix, tableColumns, clusterProps);
+    const tableName = await createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
     return { PhysicalResourceId: tableName };
   } else if (event.RequestType === 'Delete') {
-    await dropTable(event.PhysicalResourceId, clusterProps);
+    await dropTable(event.PhysicalResourceId, tableAndClusterProps);
     return;
   } else if (event.RequestType === 'Update') {
     const tableName = await updateTable(
       event.PhysicalResourceId,
       tableNamePrefix,
       tableNameSuffix,
       tableColumns,
-      clusterProps,
-      event.OldResourceProperties as TableHandlerProps & ClusterProps,
+      tableAndClusterProps,
+      event.OldResourceProperties as TableAndClusterProps,
     );
     return { PhysicalResourceId: tableName };
   } else {
@@ -32,10 +32,33 @@ export async function handler(props: TableHandlerProps & ClusterProps, event: AW
   }
 }
 
-async function createTable(tableNamePrefix: string, tableNameSuffix: string, tableColumns: Column[], clusterProps: ClusterProps): Promise<string> {
+async function createTable(
+  tableNamePrefix: string,
+  tableNameSuffix: string,
+  tableColumns: Column[],
+  tableAndClusterProps: TableAndClusterProps,
+): Promise<string> {
   const tableName = tableNamePrefix + tableNameSuffix;
   const tableColumnsString = tableColumns.map(column => `${column.name} ${column.dataType}`).join();
-  await executeStatement(`CREATE TABLE ${tableName} (${tableColumnsString})`, clusterProps);
+
+  let statement = `CREATE TABLE ${tableName} (${tableColumnsString})`;
+
+  if (tableAndClusterProps.distStyle) {
+    statement += ` DISTSTYLE ${tableAndClusterProps.distStyle}`;
+  }
+
+  const distKeyColumn = getDistKeyColumn(tableColumns);
+  if (distKeyColumn) {
+    statement += ` DISTKEY(${distKeyColumn.name})`;
+  }
+
+  const sortKeyColumns = getSortKeyColumns(tableColumns);
+  if (sortKeyColumns.length > 0) {
+    const sortKeyColumnsString = getSortKeyColumnsString(sortKeyColumns);
+    statement += ` ${tableAndClusterProps.sortStyle} SORTKEY(${sortKeyColumnsString})`;
+  }
+
+  await executeStatement(statement, tableAndClusterProps);
   return tableName;
 }
 
@@ -48,28 +71,79 @@ async function updateTable(
   tableNamePrefix: string,
   tableNameSuffix: string,
   tableColumns: Column[],
-  clusterProps: ClusterProps,
-  oldResourceProperties: TableHandlerProps & ClusterProps,
+  tableAndClusterProps: TableAndClusterProps,
+  oldResourceProperties: TableAndClusterProps,
 ): Promise<string> {
+  const alterationStatements: string[] = [];
+
   const oldClusterProps = oldResourceProperties;
-  if (clusterProps.clusterName !== oldClusterProps.clusterName || clusterProps.databaseName !== oldClusterProps.databaseName) {
-    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, clusterProps);
+  if (tableAndClusterProps.clusterName !== oldClusterProps.clusterName || tableAndClusterProps.databaseName !== oldClusterProps.databaseName) {
+    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
   }
 
   const oldTableNamePrefix = oldResourceProperties.tableName.prefix;
   if (tableNamePrefix !== oldTableNamePrefix) {
-    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, clusterProps);
+    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
   }
 
   const oldTableColumns = oldResourceProperties.tableColumns;
   if (!oldTableColumns.every(oldColumn => tableColumns.some(column => column.name === oldColumn.name && column.dataType === oldColumn.dataType))) {
-    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, clusterProps);
+    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
   }
 
-  const additions = tableColumns.filter(column => {
+  const columnAdditions = tableColumns.filter(column => {
     return !oldTableColumns.some(oldColumn => column.name === oldColumn.name && column.dataType === oldColumn.dataType);
   }).map(column => `ADD ${column.name} ${column.dataType}`);
-  await Promise.all(additions.map(addition => executeStatement(`ALTER TABLE ${tableName} ${addition}`, clusterProps)));
+  if (columnAdditions.length > 0) {
+    alterationStatements.push(...columnAdditions.map(addition => `ALTER TABLE ${tableName} ${addition}`));
+  }
+
+  const oldDistStyle = oldResourceProperties.distStyle;
+  if ((!oldDistStyle && tableAndClusterProps.distStyle) ||
+    (oldDistStyle && !tableAndClusterProps.distStyle)) {
+    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
+  } else if (oldDistStyle !== tableAndClusterProps.distStyle) {
+    alterationStatements.push(`ALTER TABLE ${tableName} ALTER DISTSTYLE ${tableAndClusterProps.distStyle}`);
+  }
+
+  const oldDistKey = getDistKeyColumn(oldTableColumns)?.name;
+  const newDistKey = getDistKeyColumn(tableColumns)?.name;
+  if ((!oldDistKey && newDistKey ) || (oldDistKey && !newDistKey)) {
+    return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
+  } else if (oldDistKey !== newDistKey) {
+    alterationStatements.push(`ALTER TABLE ${tableName} ALTER DISTKEY ${newDistKey}`);
+  }
+
+  const oldSortKeyColumns = getSortKeyColumns(oldTableColumns);
+  const newSortKeyColumns = getSortKeyColumns(tableColumns);
+  const oldSortStyle = oldResourceProperties.sortStyle;
+  const newSortStyle = tableAndClusterProps.sortStyle;
+  if ((oldSortStyle === newSortStyle && !areColumnsEqual(oldSortKeyColumns, newSortKeyColumns))
+    || (oldSortStyle !== newSortStyle)) {
+    switch (newSortStyle) {
+      case TableSortStyle.INTERLEAVED:
+        // INTERLEAVED sort key addition requires replacement.
+        // https://docs.aws.amazon.com/redshift/latest/dg/r_ALTER_TABLE.html
+        return createTable(tableNamePrefix, tableNameSuffix, tableColumns, tableAndClusterProps);
+
+      case TableSortStyle.COMPOUND: {
+        const sortKeyColumnsString = getSortKeyColumnsString(newSortKeyColumns);
+        alterationStatements.push(`ALTER TABLE ${tableName} ALTER ${newSortStyle} SORTKEY(${sortKeyColumnsString})`);
+        break;
+      }
+
+      case TableSortStyle.AUTO: {
+        alterationStatements.push(`ALTER TABLE ${tableName} ALTER SORTKEY ${newSortStyle}`);
+        break;
+      }
+    }
+  }
+
+  await Promise.all(alterationStatements.map(statement => executeStatement(statement, tableAndClusterProps)));
 
   return tableName;
 }
+
+function getSortKeyColumnsString(sortKeyColumns: Column[]) {
+  return sortKeyColumns.map(column => column.name).join();
+}

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

@@ -0,0 +1,26 @@
+import { DatabaseQueryHandlerProps, TableHandlerProps } from '../handler-props';
+
+export type ClusterProps = Omit<DatabaseQueryHandlerProps, 'handler'>;
+export type TableAndClusterProps = TableHandlerProps & ClusterProps;
+
+/**
+ * The sort style of a table.
+ * This has been duplicated here to exporting private types.
+ */
+export enum TableSortStyle {
+  /**
+   * Amazon Redshift assigns an optimal sort key based on the table data.
+   */
+  AUTO = 'AUTO',
+
+  /**
+   * Specifies that the data is sorted using a compound key made up of all of the listed columns,
+   * in the order they are listed.
+   */
+  COMPOUND = 'COMPOUND',
+
+  /**
+   * Specifies that the data is sorted using an interleaved sort key.
+   */
+  INTERLEAVED = 'INTERLEAVED',
+}

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

@@ -3,7 +3,8 @@ import * as AWSLambda from 'aws-lambda';
 /* eslint-disable-next-line import/no-extraneous-dependencies */
 import * as SecretsManager from 'aws-sdk/clients/secretsmanager';
 import { UserHandlerProps } from '../handler-props';
-import { ClusterProps, executeStatement, makePhysicalId } from './util';
+import { ClusterProps } from './types';
+import { executeStatement, makePhysicalId } from './util';
 
 const secretsManager = new SecretsManager();
 

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

@@ -1,11 +1,10 @@
 /* eslint-disable-next-line import/no-extraneous-dependencies */
 import * as RedshiftData from 'aws-sdk/clients/redshiftdata';
-import { DatabaseQueryHandlerProps } from '../handler-props';
+import { Column } from '../../table';
+import { ClusterProps } from './types';
 
 const redshiftData = new RedshiftData();
 
-export type ClusterProps = Omit<DatabaseQueryHandlerProps, 'handler'>;
-
 export async function executeStatement(statement: string, clusterProps: ClusterProps): Promise<void> {
   const executeStatementProps = {
     ClusterIdentifier: clusterProps.clusterName,
@@ -38,3 +37,30 @@ async function waitForStatementComplete(statementId: string): Promise<void> {
 export function makePhysicalId(resourceName: string, clusterProps: ClusterProps, requestId: string): string {
   return `${clusterProps.clusterName}:${clusterProps.databaseName}:${resourceName}:${requestId}`;
 }
+
+export function getDistKeyColumn(columns: Column[]): Column | undefined {
+  // string comparison is required for custom resource since everything is passed as string
+  const distKeyColumns = columns.filter(column => column.distKey === true || (column.distKey as unknown as string) === 'true');
+
+  if (distKeyColumns.length === 0) {
+    return undefined;
+  } else if (distKeyColumns.length > 1) {
+    throw new Error('Multiple dist key columns found');
+  }
+
+  return distKeyColumns[0];
+}
+
+export function getSortKeyColumns(columns: Column[]): Column[] {
+  // string comparison is required for custom resource since everything is passed as string
+  return columns.filter(column => column.sortKey === true || (column.sortKey as unknown as string) === 'true');
+}
+
+export function areColumnsEqual(columnsA: Column[], columnsB: Column[]): boolean {
+  if (columnsA.length !== columnsB.length) {
+    return false;
+  }
+  return columnsA.every(columnA => {
+    return columnsB.find(column => column.name === columnA.name && column.dataType === columnA.dataType);
+  });
+}

packages/@aws-cdk/aws-redshift/lib/private/handler-props.ts

@@ -1,4 +1,4 @@
-import { Column } from '../table';
+import { Column, TableDistStyle, TableSortStyle } from '../table';
 
 export interface DatabaseQueryHandlerProps {
   readonly handler: string;
@@ -18,6 +18,8 @@ export interface TableHandlerProps {
     readonly generateSuffix: string;
   };
   readonly tableColumns: Column[];
+  readonly distStyle?: TableDistStyle;
+  readonly sortStyle: TableSortStyle;
 }
 
 export interface TablePrivilege {