From 6aaac9325990d8e6c7edb126725692dc77c36f69 Mon Sep 17 00:00:00 2001
From: Lee Hannigan <lhnng@amazon.com>
Date: Fri, 20 Sep 2024 18:00:22 +0100
Subject: [PATCH] Add TableV2 Stream Resource Policy

---
 packages/aws-cdk-lib/aws-dynamodb/README.md   | 24 +++++++++++++++++++
 .../aws-cdk-lib/aws-dynamodb/lib/table-v2.ts  | 21 +++++++++++++++-
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/aws-dynamodb/README.md b/packages/aws-cdk-lib/aws-dynamodb/README.md
index 24e09c418b6ae..0a1f1c4d173d9 100644
--- a/packages/aws-cdk-lib/aws-dynamodb/README.md
+++ b/packages/aws-cdk-lib/aws-dynamodb/README.md
@@ -687,6 +687,30 @@ Using `resourcePolicy` you can add a [resource policy](https://docs.aws.amazon.c
 TableV2 doesn’t support creating a replica and adding a resource-based policy to that replica in the same stack update in Regions other than the Region where you deploy the stack update.
 To incorporate a resource-based policy into a replica, you'll need to initially deploy the replica without the policy, followed by a subsequent update to include the desired policy.
 
+You can also add a resource policy to a DynamoDB stream, using the `streamResourcePolicy` parameter:
+
+```
+const streamDoc = new iam.PolicyDocument({
+  statements: [
+    new iam.PolicyStatement({
+      actions: ['dynamodb:GetRecords', 'dynamodb:DescribeStream'],
+      principals: [new iam.AccountRootPrincipal()],
+      resources: ['*'],
+    }),
+  ],
+});
+
+new dynamodb.TableV2(this, 'StreamPolicyTable', {
+  partitionKey: {
+    name: 'id',
+    type: dynamodb.AttributeType.STRING,
+  },
+  removalPolicy: RemovalPolicy.DESTROY,
+  streamResourcePolicy: streamDoc,
+  dynamoStream: dynamodb.StreamViewType.NEW_AND_OLD_IMAGES,
+});
+```
+
 ## Grants
 
 Using any of the `grant*` methods on an instance of the `TableV2` construct will only apply to the primary table, its indexes, and any associated `encryptionKey`. As an example, `grantReadData` used below will only apply the table in `us-west-2`:
diff --git a/packages/aws-cdk-lib/aws-dynamodb/lib/table-v2.ts b/packages/aws-cdk-lib/aws-dynamodb/lib/table-v2.ts
index 65e3c588968e5..dd56b9897c9da 100644
--- a/packages/aws-cdk-lib/aws-dynamodb/lib/table-v2.ts
+++ b/packages/aws-cdk-lib/aws-dynamodb/lib/table-v2.ts
@@ -156,6 +156,13 @@ export interface TableOptionsV2 {
    * @default - No resource policy statements are added to the created table.
    */
   readonly resourcePolicy?: PolicyDocument;
+
+  /**
+ * Resource policy to assign to DynamoDB Stream.
+ * @see https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-dynamodb-globaltable-replicaspecification.html#cfn-dynamodb-globaltable-replicaspecification-resourcepolicy
+ * @default - No resource policy statements are added to the stream.
+ */
+  readonly streamResourcePolicy?: PolicyDocument;
 }
 
 /**
@@ -399,13 +406,15 @@ export class TableV2 extends TableBaseV2 {
       public readonly tableStreamArn?: string;
       public readonly encryptionKey?: IKey;
       public readonly resourcePolicy?: PolicyDocument;
+      public readonly streamResourcePolicy?: PolicyDocument;
 
       protected readonly region: string;
       protected readonly hasIndex = (attrs.grantIndexPermissions ?? false) ||
         (attrs.globalIndexes ?? []).length > 0 ||
         (attrs.localIndexes ?? []).length > 0;
 
-      public constructor(tableArn: string, tableName: string, tableId?: string, tableStreamArn?: string, resourcePolicy?: PolicyDocument) {
+      public constructor(tableArn: string, tableName: string, tableId?: string, tableStreamArn?: string,
+        resourcePolicy?: PolicyDocument, streamResourcePolicy?: PolicyDocument) {
         super(scope, id, { environmentFromArn: tableArn });
 
         const resourceRegion = stack.splitArn(tableArn, ArnFormat.SLASH_RESOURCE_NAME).region;
@@ -420,6 +429,7 @@ export class TableV2 extends TableBaseV2 {
         this.tableStreamArn = tableStreamArn;
         this.encryptionKey = attrs.encryptionKey;
         this.resourcePolicy = resourcePolicy;
+        this.streamResourcePolicy = streamResourcePolicy;
       }
     }
 
@@ -480,6 +490,11 @@ export class TableV2 extends TableBaseV2 {
    */
   public resourcePolicy?: PolicyDocument;
 
+  /**
+   * @attribute
+   */
+  public streamResourcePolicy?: PolicyDocument;
+
   protected readonly region: string;
 
   private readonly billingMode: string;
@@ -665,6 +680,7 @@ export class TableV2 extends TableBaseV2 {
     const pointInTimeRecovery = props.pointInTimeRecovery ?? this.tableOptions.pointInTimeRecovery;
     const contributorInsights = props.contributorInsights ?? this.tableOptions.contributorInsights;
     const resourcePolicy = props.resourcePolicy ?? this.tableOptions.resourcePolicy;
+    const streamResourcePolicy = props.region === this.region ? this.tableOptions.streamResourcePolicy : props.streamResourcePolicy || undefined;
 
     return {
       region: props.region,
@@ -693,6 +709,9 @@ export class TableV2 extends TableBaseV2 {
       resourcePolicy: resourcePolicy
         ? { policyDocument: resourcePolicy }
         : undefined,
+      replicaStreamSpecification: streamResourcePolicy
+        ? { resourcePolicy: { policyDocument: streamResourcePolicy } }
+        : undefined,
     };
   }