v2 stream provider

giograno · giograno · commit 5c1fe2d037d3 · 2025-05-20T18:21:09.000+02:00
diff --git a/localstack-core/localstack/services/dynamodbstreams/dynamodbstreams_api.py b/localstack-core/localstack/services/dynamodbstreams/dynamodbstreams_api.py
@@ -5,8 +5,10 @@
 from bson.json_util import dumps
 
 from localstack import config
+from localstack.aws.api import RequestContext
 from localstack.aws.api.dynamodbstreams import StreamStatus, StreamViewType, TableName
 from localstack.aws.connect import connect_to
+from localstack.services.dynamodb.v2.provider import DynamoDBProvider
 from localstack.services.dynamodbstreams.models import DynamoDbStreamsStore, dynamodbstreams_stores
 from localstack.utils.aws import arns, resources
 from localstack.utils.common import now_utc
@@ -211,3 +213,23 @@ def get_shard_id(stream: Dict, kinesis_shard_id: str) -> str:
         stream["shards_id_map"][kinesis_shard_id] = ddb_stream_shard_id
 
     return ddb_stream_shard_id
+
+
+def get_original_region(
+    context: RequestContext, stream_arn: str | None = None, table_name: str | None = None
+) -> str:
+    """
+    In DDB Global tables, we forward all the requests to the original region, instead of really replicating the data.
+    Since each table has a separate stream associated, we need to have a similar forwarding logic for DDB Streams.
+    To determine the original region, we need the table name, that can be either provided here or determined from the
+    ARN of the stream.
+    """
+    if not stream_arn and not table_name:
+        LOG.debug(
+            "No Stream ARN or table name provided. Returning region '%s' from the request",
+            context.region,
+        )
+        return context.region
+
+    table_name = table_name or table_name_from_stream_arn(stream_arn)
+    return DynamoDBProvider.get_global_table_region(context=context, table_name=table_name)
diff --git a/localstack-core/localstack/services/dynamodbstreams/provider.py b/localstack-core/localstack/services/dynamodbstreams/provider.py
@@ -29,6 +29,7 @@
     get_dynamodbstreams_store,
     get_kinesis_client,
     get_kinesis_stream_name,
+    get_original_region,
     get_shard_id,
     kinesis_shard_id,
     stream_name_from_stream_arn,
@@ -47,26 +48,6 @@
 }
 
 
-def get_original_region(
-    context: RequestContext, stream_arn: str | None = None, table_name: str | None = None
-) -> str:
-    """
-    In DDB Global tables, we forward all the requests to the original region, instead of really replicating the data.
-    Since each table has a separate stream associated, we need to have a similar forwarding logic for DDB Streams.
-    To determine the original region, we need the table name, that can be either provided here or determined from the
-    ARN of the stream.
-    """
-    if not stream_arn and not table_name:
-        LOG.debug(
-            "No Stream ARN or table name provided. Returning region '%s' from the request",
-            context.region,
-        )
-        return context.region
-
-    table_name = table_name or table_name_from_stream_arn(stream_arn)
-    return DynamoDBProvider.get_global_table_region(context=context, table_name=table_name)
-
-
 class DynamoDBStreamsProvider(DynamodbstreamsApi, ServiceLifecycleHook):
     def describe_stream(
         self,
diff --git a/localstack-core/localstack/services/dynamodbstreams/v2/provider.py b/localstack-core/localstack/services/dynamodbstreams/v2/provider.py
@@ -15,7 +15,8 @@
 )
 from localstack.services.dynamodb.server import DynamodbServer
 from localstack.services.dynamodb.utils import modify_ddblocal_arns
-from localstack.services.dynamodb.v2.provider import DynamoDBProvider
+from localstack.services.dynamodb.v2.provider import DynamoDBProvider, modify_context_region
+from localstack.services.dynamodbstreams.dynamodbstreams_api import get_original_region
 from localstack.services.plugins import ServiceLifecycleHook
 from localstack.utils.aws.arns import parse_arn
 
@@ -33,6 +34,14 @@ def on_after_init(self):
     def on_before_start(self):
         self.server.start_dynamodb()
 
+    def _forward_request(
+        self, context: RequestContext, region: str | None, service_request: ServiceRequest
+    ) -> ServiceResponse:
+        if region:
+            with modify_context_region(context, region):
+                return self.forward_request(context, service_request=service_request)
+        return self.forward_request(context, service_request=service_request)
+
     def forward_request(
         self, context: RequestContext, service_request: ServiceRequest = None
     ) -> ServiceResponse:
@@ -55,12 +64,19 @@ def describe_stream(
         context: RequestContext,
         payload: DescribeStreamInput,
     ) -> DescribeStreamOutput:
+        global_table_region = get_original_region(context=context, stream_arn=payload["StreamArn"])
         request = payload.copy()
         request["StreamArn"] = self.modify_stream_arn_for_ddb_local(request.get("StreamArn", ""))
-        return self.forward_request(context, request)
+        return self._forward_request(
+            context=context, service_request=request, region=global_table_region
+        )
 
     @handler("GetRecords", expand=False)
     def get_records(self, context: RequestContext, payload: GetRecordsInput) -> GetRecordsOutput:
+        # Limitation note: With this current implementation, we are not able to get the records from a stream of a
+        # replicated table. To do so, we would need to kept track of the emitted ShardIterators and the originating
+        # region in `GetShardIterator`.
+
         request = payload.copy()
         request["ShardIterator"] = self.modify_stream_arn_for_ddb_local(
             request.get("ShardIterator", "")
@@ -77,5 +93,8 @@ def get_shard_iterator(
 
     @handler("ListStreams", expand=False)
     def list_streams(self, context: RequestContext, payload: ListStreamsInput) -> ListStreamsOutput:
+        global_table_region = get_original_region(context=context, stream_arn=payload["TableName"])
         # TODO: look into `ExclusiveStartStreamArn` param
-        return self.forward_request(context, payload)
+        return self._forward_request(
+            context=context, service_request=payload, region=global_table_region
+        )
