add a waiter function

Author: anthony-nhs

packages/cdk/constructs/LambdaFunction.ts

@@ -24,7 +24,7 @@ export interface LambdaFunctionProps {
   readonly functionName: string
   readonly packageBasePath: string
   readonly handler: string
-  readonly environmentVariables: {[key: string]: string}
+  readonly environmentVariables?: {[key: string]: string}
   readonly additionalPolicies?: Array<IManagedPolicy>
   readonly logRetentionInDays: number
   readonly logLevel: string

packages/cdk/nagSuppressions.ts

@@ -92,6 +92,30 @@ export const nagSuppressions = (stack: Stack) => {
     ]
   )
 
+  // Suppress IAM wildcard permissions for waiter function execution role policy
+  safeAddNagSuppression(
+    stack,
+    "/EpsAssistMeStack/VectorIndex/SlackBotLambda/LambdaPutLogsManagedPolicy/Resource",
+    [
+      {
+        id: "AwsSolutions-IAM5",
+        reason: "Lambda needs access to all OpenSearch collections and indexes to create and manage indexes."
+      }
+    ]
+  )
+
+  // Suppress IAM wildcard permissions for waiter on event role policy
+  safeAddNagSuppression(
+    stack,
+    "/EpsAssistMeStack/VectorIndex/IndexWaiterProvider/framework-onEvent/ServiceRole/DefaultPolicy/Resource",
+    [
+      {
+        id: "AwsSolutions-IAM5",
+        reason: "Lambda needs access to all OpenSearch collections and indexes to create and manage indexes."
+      }
+    ]
+  )
+
   // Suppress wildcard permissions for SlackBot policy
   safeAddNagSuppression(
     stack,
@@ -142,6 +166,17 @@ export const nagSuppressions = (stack: Stack) => {
     ]
   )
 
+  safeAddNagSuppression(
+    stack,
+    "/EpsAssistMeStack/VectorIndex/IndexWaiterProvider/framework-onEvent/ServiceRole/Resource",
+    [
+      {
+        id: "AwsSolutions-IAM4",
+        reason: "Waiter function on event using managed policies is fine"
+      }
+    ]
+  )
+
   // Suppress AWS managed policy usage in BucketNotificationsHandler (wildcard for any hash)
   const bucketNotificationHandlers = stack.node.findAll().filter(node =>
     node.node.id.startsWith("BucketNotificationsHandler")

packages/cdk/resources/VectorIndex.ts

@@ -1,14 +1,24 @@
 import {Construct} from "constructs"
 import {CfnCollection, CfnIndex} from "aws-cdk-lib/aws-opensearchserverless"
+import {CustomResource} from "aws-cdk-lib"
+import {ManagedPolicy, PolicyStatement} from "aws-cdk-lib/aws-iam"
+import {Provider} from "aws-cdk-lib/custom-resources"
+import {LambdaFunction} from "../constructs/LambdaFunction"
 
 export interface VectorIndexProps {
   readonly indexName: string
   readonly collection: CfnCollection
   readonly endpoint: string
+  readonly account: string
+  readonly region: string
+  readonly stackName: string
+  readonly logRetentionInDays: number
+  readonly logLevel: string
 }
 
 export class VectorIndex extends Construct {
   public readonly cfnIndex: CfnIndex
+  public readonly indexReady: CustomResource
 
   constructor(scope: Construct, id: string, props: VectorIndexProps) {
     super(scope, id)
@@ -61,8 +71,56 @@ export class VectorIndex extends Construct {
       }
     })
 
+    const collectionArn = `arn:aws:aoss:${props.region}:${props.account}:collection/${props.collection.name}`
+    const indexArn = `arn:aws:aoss:${props.region}:${props.account}:index/${props.collection.name}/${props.indexName}`
+
+    const getCollectionPolicy = new PolicyStatement({
+      actions: [
+        "opensearchserverless:BatchGetCollection"
+      ],
+      resources: [collectionArn]
+    })
+    const getIndexPolicy = new PolicyStatement({
+      actions: [
+        "opensearchserverless:BatchGetIndex"
+      ],
+      resources: [indexArn]
+    })
+    const waiterFnManagedPolicy = new ManagedPolicy(this, "Policy", {
+      description: "Policy for Bedrock Knowledge Base to access S3 and OpenSearch",
+      statements: [
+        getCollectionPolicy,
+        getIndexPolicy
+      ]
+    })
+
+    const waiterFn = new LambdaFunction(this, "SlackBotLambda", {
+      stackName: props.stackName,
+      functionName: `${props.stackName}-VectorIndexWaiter`,
+      packageBasePath: "packages/cdk/resources/lambda",
+      handler: "index_waiter.handler",
+      logRetentionInDays: props.logRetentionInDays,
+      logLevel: props.logLevel,
+      additionalPolicies: [waiterFnManagedPolicy]
+    })
+
+    const provider = new Provider(this, "IndexWaiterProvider", {
+      onEventHandler: waiterFn.function
+    })
+
+    const indexReady = new CustomResource(this, "IndexReady", {
+      serviceToken: provider.serviceToken,
+      properties: {
+        CollectionName: props.collection.name,
+        IndexName: props.indexName
+      }
+    })
     // Ensure collection exists before creating index
     cfnIndex.node.addDependency(props.collection)
+    indexReady.node.addDependency(props.collection)
+    indexReady.node.addDependency(cfnIndex)
+
     this.cfnIndex = cfnIndex
+    this.indexReady = indexReady
   }
 }

packages/cdk/resources/lambda/index_waiter.py

@@ -0,0 +1,38 @@
+import boto3
+import time
+import json
+
+aoss = boto3.client("opensearchserverless")
+
+
+def handler(event, context):
+    print("Received event:", json.dumps(event))
+    request_type = event["RequestType"]
+    collection_name = event["ResourceProperties"]["CollectionName"]
+    index_name = event["ResourceProperties"]["IndexName"]
+
+    if request_type == "Delete":
+        return {"PhysicalResourceId": f"{collection_name}/{index_name}", "Data": {"Status": "DELETED"}}
+
+    # Poll until both collection + index become ACTIVE
+    for _ in range(60):  # 10 minutes max
+        # 1. Check collection
+        coll_resp = aoss.batch_get_collection(names=[collection_name])
+        coll = next((c for c in coll_resp.get("collectionDetails", []) if c["name"] == collection_name), None)
+        if not coll or coll["status"] != "ACTIVE":
+            print(f"Collection {collection_name} status={coll['status'] if coll else 'MISSING'}")
+            time.sleep(10)
+            continue
+
+        # 2. Check index
+        idx_resp = aoss.batch_get_index(names=[index_name], collectionName=collection_name)
+        idx = next((i for i in idx_resp.get("indexDetails", []) if i["name"] == index_name), None)
+
+        if idx and idx["status"] == "ACTIVE":
+            print(f"Index {index_name} in {collection_name} is ACTIVE")
+            return {"PhysicalResourceId": f"{collection_name}/{index_name}", "Data": {"Status": "READY"}}
+
+        print(f"Index {index_name} status={idx['status'] if idx else 'MISSING'}")
+        time.sleep(10)
+
+    raise Exception(f"Collection {collection_name} / Index {index_name} not ready after timeout")

packages/cdk/stacks/EpsAssistMeStack.ts

@@ -137,11 +137,17 @@ export class EpsAssistMeStack extends Stack {
     const vectorIndex = new VectorIndex(this, "VectorIndex", {
       indexName: VECTOR_INDEX_NAME,
       collection: openSearchResources.collection.collection,
-      endpoint
+      endpoint,
+      region: this.region,
+      account: this.account,
+      stackName: props.stackName,
+      logRetentionInDays: logRetentionInDays,
+      logLevel: logLevel
     })
 
     // Ensure knowledge base waits for vector index
     vectorKB.knowledgeBase.node.addDependency(vectorIndex.cfnIndex)
+    vectorKB.knowledgeBase.node.addDependency(vectorIndex.indexReady)
 
     // Add S3 notification to trigger sync Lambda function
     new S3LambdaNotification(this, "S3LambdaNotification", {