fix(glue-alpha): race condition when creating mutiple indices

package.json

@@ -131,20 +131,20 @@
       "@aws-cdk/core/table/**",
       "@aws-cdk/core/yaml",
       "@aws-cdk/core/yaml/**",
-      "@aws-cdk/cx-api/semver",
-      "@aws-cdk/cx-api/semver/**",
       "@aws-cdk/cx-api/@aws-cdk/cloud-assembly-api",
       "@aws-cdk/cx-api/@aws-cdk/cloud-assembly-api/**",
+      "@aws-cdk/cx-api/semver",
+      "@aws-cdk/cx-api/semver/**",
       "@aws-cdk/mixins-preview/minimatch",
       "@aws-cdk/mixins-preview/minimatch/**",
       "@aws-cdk/pipelines/aws-sdk",
       "@aws-cdk/pipelines/aws-sdk/**",
       "@aws-cdk/yaml-cfn/yaml",
       "@aws-cdk/yaml-cfn/yaml/**",
-      "aws-cdk-lib/@balena/dockerignore",
-      "aws-cdk-lib/@balena/dockerignore/**",
       "aws-cdk-lib/@aws-cdk/cloud-assembly-api",
       "aws-cdk-lib/@aws-cdk/cloud-assembly-api/**",
+      "aws-cdk-lib/@balena/dockerignore",
+      "aws-cdk-lib/@balena/dockerignore/**",
       "aws-cdk-lib/case",
       "aws-cdk-lib/case/**",
       "aws-cdk-lib/chalk",

packages/@aws-cdk/aws-glue-alpha/lib/partition-index-handler/index.ts

@@ -0,0 +1,78 @@
+// eslint-disable-next-line import/no-extraneous-dependencies
+import { GlueClient, CreatePartitionIndexCommand, DeletePartitionIndexCommand, GetPartitionIndexesCommand } from '@aws-sdk/client-glue';
+
+class PartitionIndexError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = 'PartitionIndexError';
+  }
+}
+
+const glue = new GlueClient({});
+
+export async function onEvent(event: any) {
+  const { DatabaseName, TableName, IndexName, Keys } = event.ResourceProperties;
+
+  if (event.RequestType === 'Create') {
+    await glue.send(new CreatePartitionIndexCommand({
+      DatabaseName,
+      TableName,
+      PartitionIndex: { IndexName, Keys },
+    }));
+    return { PhysicalResourceId: IndexName };
+  } else if (event.RequestType === 'Update') {
+    const oldKeys = event.OldResourceProperties?.Keys;
+    if (JSON.stringify(oldKeys) !== JSON.stringify(Keys)) {
+    // For auto-generated index names, this should be unreachable. CDK derives index names from keys, so key changes always produce a new resource
+      throw new PartitionIndexError('Partition index keys cannot be updated. Delete and recreate the index instead.');
+    }
+    return { PhysicalResourceId: IndexName };
+  } else if (event.RequestType === 'Delete') {
+    try {
+      await glue.send(new DeletePartitionIndexCommand({
+        DatabaseName,
+        TableName,
+        IndexName,
+      }));
+    } catch (e: any) {
+      if (e.name === 'EntityNotFoundException') {
+        // eslint-disable-next-line no-console
+        console.log(`Partition index ${IndexName} not found on ${DatabaseName}.${TableName} - may have been deleted out-of-band`);
+      } else {
+        throw e;
+      }
+    }
+  }
+
+  return {};
+}
+
+export async function isComplete(event: any) {
+  const { DatabaseName, TableName, IndexName } = event.ResourceProperties;
+
+  if (event.RequestType === 'Delete') {
+    const resp = await glue.send(new GetPartitionIndexesCommand({ DatabaseName, TableName }));
+    const index = (resp.PartitionIndexDescriptorList || []).find(
+      (i) => i.IndexName!.toLowerCase() === IndexName.toLowerCase(),
+    );
+    if (!index || index.IndexStatus !== 'DELETING') {
+      return { IsComplete: true };
+    }
+    return { IsComplete: false };
+  }
+
+  const resp = await glue.send(new GetPartitionIndexesCommand({ DatabaseName, TableName }));
+  // Glue lowercases index names, so compare case-insensitively
+  const index = (resp.PartitionIndexDescriptorList || []).find(
+    (i) => i.IndexName!.toLowerCase() === IndexName.toLowerCase(),
+  );
+
+  if (!index) return { IsComplete: false };
+  if (index.IndexStatus === 'ACTIVE') return { IsComplete: true };
+  if (index.IndexStatus === 'CREATING') return { IsComplete: false };
+
+  const errorDetails = index.BackfillErrors?.length
+    ? ` Backfill errors: ${JSON.stringify(index.BackfillErrors)}`
+    : '';
+  throw new PartitionIndexError(`Partition index ${IndexName} in unexpected state: ${index.IndexStatus}.${errorDetails}`);
+}

packages/@aws-cdk/aws-glue-alpha/lib/table-base.ts

@@ -1,10 +1,11 @@
+import * as path from 'path';
 import type { CfnTable } from 'aws-cdk-lib/aws-glue';
 import * as iam from 'aws-cdk-lib/aws-iam';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
 import type { IResource } from 'aws-cdk-lib/core';
-import { ArnFormat, Fn, Lazy, Names, Resource, Stack, UnscopedValidationError, ValidationError } from 'aws-cdk-lib/core';
+import { ArnFormat, CustomResource, Duration, Fn, Lazy, Names, Resource, Stack, Token, UnscopedValidationError, ValidationError } from 'aws-cdk-lib/core';
 import { lit } from 'aws-cdk-lib/core/lib/helpers-internal';
 import * as cr from 'aws-cdk-lib/custom-resources';
-import type { AwsCustomResource } from 'aws-cdk-lib/custom-resources';
 import type { Construct } from 'constructs';
 import type { DataFormat } from './data-format';
 import type { IDatabase } from './database';
@@ -254,7 +255,7 @@ export abstract class TableBase extends Resource implements ITable {
    * race conditions, we store the resource and add dependencies
    * each time a new partition index is created.
    */
-  private partitionIndexCustomResources: AwsCustomResource[] = [];
+  private partitionIndexCustomResources: CustomResource[] = [];
 
   constructor(scope: Construct, id: string, props: TableBaseProps) {
     super(scope, id, {
@@ -298,37 +299,79 @@ export abstract class TableBase extends Resource implements ITable {
     this.validatePartitionIndex(index);
 
     const indexName = index.indexName ?? this.generateIndexName(index.keyNames);
-    const partitionIndexCustomResource = new cr.AwsCustomResource(this, `partition-index-${indexName}`, {
-      onCreate: {
-        service: 'Glue',
-        action: 'createPartitionIndex',
-        parameters: {
-          DatabaseName: this.database.databaseName,
-          TableName: this.tableName,
-          PartitionIndex: {
-            IndexName: indexName,
-            Keys: index.keyNames,
-          },
-        },
-        physicalResourceId: cr.PhysicalResourceId.of(
-          indexName,
-        ),
+    const { provider, handler, isCompleteHandler } = this.getOrCreatePartitionIndexProvider();
+
+    // Add scoped permissions for this table's Glue resources
+    // https://docs.aws.amazon.com/service-authorization/latest/reference/list_awsglue.html
+    const tablePolicy = new iam.PolicyStatement({
+      actions: ['glue:CreatePartitionIndex', 'glue:DeletePartitionIndex', 'glue:GetPartitionIndexes', 'glue:GetTable', 'glue:UpdateTable'],
+      resources: [this.tableArn, this.database.databaseArn, this.database.catalogArn],
+    });
+    handler.addToRolePolicy(tablePolicy);
+    isCompleteHandler.addToRolePolicy(tablePolicy);
+
+    const partitionIndexCustomResource = new CustomResource(this, `partition-index-${indexName}`, {
+      resourceType: 'Custom::GluePartitionIndex',
+      serviceToken: provider.serviceToken,
+      properties: {
+        DatabaseName: this.database.databaseName,
+        TableName: this.tableName,
+        IndexName: indexName,
+        Keys: index.keyNames,
       },
-      policy: cr.AwsCustomResourcePolicy.fromSdkCalls({
-        resources: cr.AwsCustomResourcePolicy.ANY_RESOURCE,
-      }),
-      // APIs are available in 2.1055.0
-      installLatestAwsSdk: false,
     });
-    this.grantToUnderlyingResources(partitionIndexCustomResource, ['glue:UpdateTable']);
 
-    // Depend on previous partition index if possible, to avoid race condition
+    // Ensure IAM policies are created before the custom resource invokes the handlers
+    partitionIndexCustomResource.node.addDependency(handler.role!);
+    partitionIndexCustomResource.node.addDependency(isCompleteHandler.role!);
+
+    // Depend on previous partition index to avoid race condition
     if (numPartitions > 0) {
-      this.partitionIndexCustomResources[numPartitions-1].node.addDependency(partitionIndexCustomResource);
+      partitionIndexCustomResource.node.addDependency(this.partitionIndexCustomResources[numPartitions-1]);
     }
     this.partitionIndexCustomResources.push(partitionIndexCustomResource);
   }
 
+  private getOrCreatePartitionIndexProvider(): { provider: cr.Provider; handler: lambda.Function; isCompleteHandler: lambda.Function } {
+    const providerId = 'GluePartitionIndexProvider';
+    const stack = Stack.of(this);
+    const existingProvider = stack.node.tryFindChild(providerId) as cr.Provider;
+    if (existingProvider) {
+      return {
+        provider: existingProvider,
+        handler: stack.node.findChild('GluePartitionIndexHandler') as lambda.Function,
+        isCompleteHandler: stack.node.findChild('GluePartitionIndexIsComplete') as lambda.Function,
+      };
+    }
+
+    const handler = new lambda.Function(stack, 'GluePartitionIndexHandler', {
+      runtime: lambda.determineLatestNodeRuntime(stack),
+      code: lambda.Code.fromAsset(path.join(__dirname, 'partition-index-handler'), {
+        exclude: ['*.ts', '*.d.ts'],
+      }),
+      handler: 'index.onEvent',
+      timeout: Duration.minutes(1),
+    });
+
+    const isCompleteHandler = new lambda.Function(stack, 'GluePartitionIndexIsComplete', {
+      runtime: lambda.determineLatestNodeRuntime(stack),
+      code: lambda.Code.fromAsset(path.join(__dirname, 'partition-index-handler'), {
+        exclude: ['*.ts', '*.d.ts'],
+      }),
+      handler: 'index.isComplete',
+      timeout: Duration.minutes(1),
+    });
+
+    const provider = new cr.Provider(stack, providerId, {
+      onEventHandler: handler,
+      isCompleteHandler: isCompleteHandler,
+      queryInterval: Duration.seconds(10),
+      totalTimeout: Duration.hours(1),
+    });
+
+    return { provider, handler, isCompleteHandler };
+  }
+
   private generateIndexName(keys: string[]): string {
     const prefix = keys.join('-') + '-';
     const uniqueId = Names.uniqueId(this);
@@ -338,7 +381,7 @@ export abstract class TableBase extends Resource implements ITable {
   }
 
   private validatePartitionIndex(index: PartitionIndex) {
-    if (index.indexName !== undefined && (index.indexName.length < 1 || index.indexName.length > 255)) {
+    if (index.indexName !== undefined && !Token.isUnresolved(index.indexName) && (index.indexName.length < 1 || index.indexName.length > 255)) {
       throw new ValidationError(lit`IndexNameLengthInvalid`, `Index name must be between 1 and 255 characters, but got ${index.indexName.length}`, this);
     }
     if (!this.partitionKeys || this.partitionKeys.length === 0) {

packages/@aws-cdk/aws-glue-alpha/package.json

@@ -86,6 +86,7 @@
     "@aws-cdk/integ-runner": "^2.198.0",
     "@aws-cdk/integ-tests-alpha": "0.0.0",
     "@aws-cdk/pkglint": "0.0.0",
+    "@aws-sdk/client-glue": "3.632.0",
     "@types/jest": "^29.5.14",
     "aws-cdk-lib": "0.0.0",
     "constructs": "^10.5.0",
@@ -112,5 +113,6 @@
       "naming/package-matches-directory",
       "assert/assert-dependency"
     ]
-  }
+  },
+  "dependencies": {}
 }