7056d8d9-7428-4543-abf2-c5dedb2c98c0

Committed 30 May 2025 02:03PM UTC coverage: 86.654% (-0.003%) from 86.657%

Build # 7056d8d9-7428-4543-abf2-c5dedb2c98c0

Build Type

push

circleci

Committed by

web-flow

Commit Message

Add stack option for CLI start command (#12675)

Co-authored-by: Silvio Vasiljevic <silvio.vasiljevic@gmail.com>
Co-authored-by: Erudit Morina <83708693+eruditmorina@users.noreply.github.com>

Coverage Stats

2 of 9 new or added lines in 1 file covered. (22.22%)

37 existing lines in 4 files now uncovered.

64641 of 74597 relevant lines covered (86.65%)

0.87 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.65

/localstack-core/localstack/services/lambda_/event_source_mapping/pollers/kinesis_poller.py

import base64
import json
import logging
from copy import deepcopy
from datetime import datetime

from botocore.client import BaseClient

from localstack.aws.api.kinesis import StreamStatus
from localstack.aws.api.pipes import (
    KinesisStreamStartPosition,
)
from localstack.services.lambda_.event_source_mapping.event_processor import (
    EventProcessor,
)
from localstack.services.lambda_.event_source_mapping.pollers.stream_poller import StreamPoller
from localstack.utils.strings import to_str

LOG = logging.getLogger(__name__)


class KinesisPoller(StreamPoller):
    # The role ARN of the processor (e.g., role ARN of the Pipe)
    invoke_identity_arn: str | None
    # Flag to enable nested kinesis namespace when formatting events to support the nested `kinesis` field structure
    # used for Lambda ESM: https://docs.aws.amazon.com/lambda/latest/dg/with-kinesis.html#services-kinesis-event-example
    # EventBridge Pipes uses no nesting: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-kinesis.html
    kinesis_namespace: bool

    def __init__(
        self,
        source_arn: str,
        source_parameters: dict | None = None,
        source_client: BaseClient | None = None,
        processor: EventProcessor | None = None,
        partner_resource_arn: str | None = None,
        invoke_identity_arn: str | None = None,
        kinesis_namespace: bool = False,
        esm_uuid: str | None = None,
        shards: dict[str, str] | None = None,
    ):
        super().__init__(
            source_arn,
            source_parameters,
            source_client,
            processor,
            esm_uuid=esm_uuid,
            partner_resource_arn=partner_resource_arn,
            shards=shards,
        )
        self.invoke_identity_arn = invoke_identity_arn
        self.kinesis_namespace = kinesis_namespace

    @property
    def stream_parameters(self) -> dict:
        return self.source_parameters["KinesisStreamParameters"]

    def initialize_shards(self) -> dict[str, str]:
        # TODO: cache this and update/re-try upon failures
        stream_info = self.source_client.describe_stream(StreamARN=self.source_arn)
        stream_status = stream_info["StreamDescription"]["StreamStatus"]
        if stream_status != StreamStatus.ACTIVE:
            LOG.warning(
                "Stream %s is not active. Current status: %s",
                self.source_arn,
                stream_status,
            )
            return {}

        # NOTICE: re-sharding might require updating this periodically (unknown how Pipes does it!?)
        # Mapping of shard id => shard iterator
        shards = {}
        for shard in stream_info["StreamDescription"]["Shards"]:
            shard_id = shard["ShardId"]
            starting_position = self.stream_parameters["StartingPosition"]
            kwargs = {}
            # TODO: test StartingPosition=AT_TIMESTAMP (only supported for Kinesis!)
            if starting_position == KinesisStreamStartPosition.AT_TIMESTAMP:
                kwargs["StartingSequenceNumber"] = self.stream_parameters[
                    "StartingPositionTimestamp"
                ]
            get_shard_iterator_response = self.source_client.get_shard_iterator(
                StreamARN=self.source_arn,
                ShardId=shard_id,
                ShardIteratorType=starting_position,
                **kwargs,
            )
            shards[shard_id] = get_shard_iterator_response["ShardIterator"]

        LOG.debug("Event source %s has %d shards.", self.source_arn, len(self.shards))
        return shards

    def stream_arn_param(self) -> dict:
        return {"StreamARN": self.source_arn}

    def event_source(self) -> str:
        return "aws:kinesis"

    def extra_metadata(self) -> dict:
        return {
            "eventVersion": "1.0",
            "eventName": "aws:kinesis:record",
            "invokeIdentityArn": self.invoke_identity_arn,
        }

    def transform_into_events(self, records: list[dict], shard_id) -> list[dict]:
        events = []
        for record in records:
            # TODO: consolidate with Kinesis event source listener:
            #  localstack.services.lambda_.event_source_listeners.kinesis_event_source_listener.KinesisEventSourceListener._create_lambda_event_payload
            #  check `encryptionType` leading to serialization errors by Dotnet Lambdas
            sequence_number = record["SequenceNumber"]
            event = {
                # TODO: add this metadata after filtering.
                #  This requires some design adjustment because the sequence number depends on the record.
                "eventID": f"{shard_id}:{sequence_number}",
            }
            kinesis_fields = {
                "kinesisSchemaVersion": "1.0",
                "partitionKey": record["PartitionKey"],
                "sequenceNumber": sequence_number,
                # TODO: implement heuristic based on content type: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html#pipes-filter-sqs
                # boto3 automatically decodes records in get_records(), so we must re-encode
                "data": to_str(base64.b64encode(record["Data"])),
                "approximateArrivalTimestamp": record["ApproximateArrivalTimestamp"].timestamp(),
            }
            if self.kinesis_namespace:
                event["kinesis"] = kinesis_fields
            else:
                event.update(kinesis_fields)
            events.append(event)
        return events

    def failure_payload_details_field_name(self) -> str:
        return "KinesisBatchInfo"

    def get_approximate_arrival_time(self, record: dict) -> float:
        if self.kinesis_namespace:
            return record["kinesis"]["approximateArrivalTimestamp"]
        else:
            return record["approximateArrivalTimestamp"]

    def format_datetime(self, time: datetime) -> str:
        return f"{time.isoformat(timespec='milliseconds')}Z"

    def get_sequence_number(self, record: dict) -> str:
        if self.kinesis_namespace:
            return record["kinesis"]["sequenceNumber"]
        else:
            return record["sequenceNumber"]

    def pre_filter(self, events: list[dict]) -> list[dict]:
        # TODO: test what happens with a mixture of data and non-data filters?
        if has_data_filter_criteria_parsed(self.filter_patterns):
            parsed_events = []
            for event in events:
                raw_data = self.get_data(event)
                try:
                    data = self.parse_data(raw_data)
                    # TODO: test "data" key remapping
                    # Filtering remaps "kinesis.data" in ESM to "data (idempotent for Pipes using "data" directly)
                    # ESM: https://docs.aws.amazon.com/lambda/latest/dg/with-kinesis-filtering.html
                    # Pipes: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-kinesis.html
                    # Pipes filtering: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html
                    parsed_event = deepcopy(event)
                    parsed_event["data"] = data

                    parsed_events.append(parsed_event)
                except json.JSONDecodeError:
                    LOG.warning(
                        "Unable to convert event data '%s' to json... Record will be dropped.",
                        raw_data,
                        exc_info=LOG.isEnabledFor(logging.DEBUG),
                    )
            return parsed_events
        else:
            return events

    def post_filter(self, events: list[dict]) -> list[dict]:
        if has_data_filter_criteria_parsed(self.filter_patterns):
            # convert them back (HACK for fixing parity with v1 and getting regression tests passing)
            for event in events:
                parsed_data = event.pop("data")
                encoded_data = self.encode_data(parsed_data)
                self.set_data(event, encoded_data)
        return events

    def get_data(self, event: dict) -> str:
        if self.kinesis_namespace:
            return event["kinesis"]["data"]
        else:
            return event["data"]

    def set_data(self, event: dict, data: bytes) -> None:
        if self.kinesis_namespace:
            event["kinesis"]["data"] = data
        else:
            event["data"] = data

    def parse_data(self, raw_data: str) -> dict | str:
        decoded_data = base64.b64decode(raw_data)
        return json.loads(decoded_data)

    def encode_data(self, parsed_data: dict) -> str:
        return base64.b64encode(json.dumps(parsed_data).encode()).decode()


def has_data_filter_criteria_parsed(parsed_filters: list[dict]) -> bool:
    for filter in parsed_filters:
        if "data" in filter:
            return True
    return False

1	import base64	1✔
2	import json	1✔
3	import logging	1✔
4	from copy import deepcopy	1✔
5	from datetime import datetime	1✔
6
7	from botocore.client import BaseClient	1✔
8
9	from localstack.aws.api.kinesis import StreamStatus	1✔
10	from localstack.aws.api.pipes import (	1✔
11	KinesisStreamStartPosition,
12	)
13	from localstack.services.lambda_.event_source_mapping.event_processor import (	1✔
14	EventProcessor,
15	)
16	from localstack.services.lambda_.event_source_mapping.pollers.stream_poller import StreamPoller	1✔
17	from localstack.utils.strings import to_str	1✔
18
19	LOG = logging.getLogger(__name__)	1✔
20
21
22	class KinesisPoller(StreamPoller):	1✔
23	# The role ARN of the processor (e.g., role ARN of the Pipe)
24	invoke_identity_arn: str \| None	1✔
25	# Flag to enable nested kinesis namespace when formatting events to support the nested `kinesis` field structure
26	# used for Lambda ESM: https://docs.aws.amazon.com/lambda/latest/dg/with-kinesis.html#services-kinesis-event-example
27	# EventBridge Pipes uses no nesting: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-kinesis.html
28	kinesis_namespace: bool	1✔
29
30	def __init__(	1✔
31	self,
32	source_arn: str,
33	source_parameters: dict \| None = None,
34	source_client: BaseClient \| None = None,
35	processor: EventProcessor \| None = None,
36	partner_resource_arn: str \| None = None,
37	invoke_identity_arn: str \| None = None,
38	kinesis_namespace: bool = False,
39	esm_uuid: str \| None = None,
40	shards: dict[str, str] \| None = None,
41	):
42	super().__init__(	1✔
43	source_arn,
44	source_parameters,
45	source_client,
46	processor,
47	esm_uuid=esm_uuid,
48	partner_resource_arn=partner_resource_arn,
49	shards=shards,
50	)
51	self.invoke_identity_arn = invoke_identity_arn	1✔
52	self.kinesis_namespace = kinesis_namespace	1✔
53
54	@property	1✔
55	def stream_parameters(self) -> dict:	1✔
56	return self.source_parameters["KinesisStreamParameters"]	1✔
57
58	def initialize_shards(self) -> dict[str, str]:	1✔
59	# TODO: cache this and update/re-try upon failures
60	stream_info = self.source_client.describe_stream(StreamARN=self.source_arn)	1✔
61	stream_status = stream_info["StreamDescription"]["StreamStatus"]	1✔
62	if stream_status != StreamStatus.ACTIVE:	1✔
UNCOV 63	LOG.warning(	×
64	"Stream %s is not active. Current status: %s",
65	self.source_arn,
66	stream_status,
67	)
UNCOV 68	return {}	×
69
70	# NOTICE: re-sharding might require updating this periodically (unknown how Pipes does it!?)
71	# Mapping of shard id => shard iterator
72	shards = {}	1✔
73	for shard in stream_info["StreamDescription"]["Shards"]:	1✔
74	shard_id = shard["ShardId"]	1✔
75	starting_position = self.stream_parameters["StartingPosition"]	1✔
76	kwargs = {}	1✔
77	# TODO: test StartingPosition=AT_TIMESTAMP (only supported for Kinesis!)
78	if starting_position == KinesisStreamStartPosition.AT_TIMESTAMP:	1✔
UNCOV 79	kwargs["StartingSequenceNumber"] = self.stream_parameters[	×
80	"StartingPositionTimestamp"
81	]
82	get_shard_iterator_response = self.source_client.get_shard_iterator(	1✔
83	StreamARN=self.source_arn,
84	ShardId=shard_id,
85	ShardIteratorType=starting_position,
86	**kwargs,
87	)
88	shards[shard_id] = get_shard_iterator_response["ShardIterator"]	1✔
89
90	LOG.debug("Event source %s has %d shards.", self.source_arn, len(self.shards))	1✔
91	return shards	1✔
92
93	def stream_arn_param(self) -> dict:	1✔
94	return {"StreamARN": self.source_arn}	1✔
95
96	def event_source(self) -> str:	1✔
97	return "aws:kinesis"	1✔
98
99	def extra_metadata(self) -> dict:	1✔
100	return {	1✔
101	"eventVersion": "1.0",
102	"eventName": "aws:kinesis:record",
103	"invokeIdentityArn": self.invoke_identity_arn,
104	}
105
106	def transform_into_events(self, records: list[dict], shard_id) -> list[dict]:	1✔
107	events = []	1✔
108	for record in records:	1✔
109	# TODO: consolidate with Kinesis event source listener:
110	# localstack.services.lambda_.event_source_listeners.kinesis_event_source_listener.KinesisEventSourceListener._create_lambda_event_payload
111	# check `encryptionType` leading to serialization errors by Dotnet Lambdas
112	sequence_number = record["SequenceNumber"]	1✔
113	event = {	1✔
114	# TODO: add this metadata after filtering.
115	# This requires some design adjustment because the sequence number depends on the record.
116	"eventID": f"{shard_id}:{sequence_number}",
117	}
118	kinesis_fields = {	1✔
119	"kinesisSchemaVersion": "1.0",
120	"partitionKey": record["PartitionKey"],
121	"sequenceNumber": sequence_number,
122	# TODO: implement heuristic based on content type: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html#pipes-filter-sqs
123	# boto3 automatically decodes records in get_records(), so we must re-encode
124	"data": to_str(base64.b64encode(record["Data"])),
125	"approximateArrivalTimestamp": record["ApproximateArrivalTimestamp"].timestamp(),
126	}
127	if self.kinesis_namespace:	1✔
128	event["kinesis"] = kinesis_fields	1✔
129	else:
UNCOV 130	event.update(kinesis_fields)	×
131	events.append(event)	1✔
132	return events	1✔
133
134	def failure_payload_details_field_name(self) -> str:	1✔
135	return "KinesisBatchInfo"	1✔
136
137	def get_approximate_arrival_time(self, record: dict) -> float:	1✔
138	if self.kinesis_namespace:	1✔
139	return record["kinesis"]["approximateArrivalTimestamp"]	1✔
140	else:
UNCOV 141	return record["approximateArrivalTimestamp"]	×
142
143	def format_datetime(self, time: datetime) -> str:	1✔
144	return f"{time.isoformat(timespec='milliseconds')}Z"	1✔
145
146	def get_sequence_number(self, record: dict) -> str:	1✔
147	if self.kinesis_namespace:	1✔
148	return record["kinesis"]["sequenceNumber"]	1✔
149	else:
UNCOV 150	return record["sequenceNumber"]	×
151
152	def pre_filter(self, events: list[dict]) -> list[dict]:	1✔
153	# TODO: test what happens with a mixture of data and non-data filters?
154	if has_data_filter_criteria_parsed(self.filter_patterns):	1✔
155	parsed_events = []	1✔
156	for event in events:	1✔
157	raw_data = self.get_data(event)	1✔
158	try:	1✔
159	data = self.parse_data(raw_data)	1✔
160	# TODO: test "data" key remapping
161	# Filtering remaps "kinesis.data" in ESM to "data (idempotent for Pipes using "data" directly)
162	# ESM: https://docs.aws.amazon.com/lambda/latest/dg/with-kinesis-filtering.html
163	# Pipes: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-kinesis.html
164	# Pipes filtering: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html
165	parsed_event = deepcopy(event)	1✔
166	parsed_event["data"] = data	1✔
167
168	parsed_events.append(parsed_event)	1✔
UNCOV 169	except json.JSONDecodeError:	×
UNCOV 170	LOG.warning(	×
171	"Unable to convert event data '%s' to json... Record will be dropped.",
172	raw_data,
173	exc_info=LOG.isEnabledFor(logging.DEBUG),
174	)
175	return parsed_events	1✔
176	else:
177	return events	1✔
178
179	def post_filter(self, events: list[dict]) -> list[dict]:	1✔
180	if has_data_filter_criteria_parsed(self.filter_patterns):	1✔
181	# convert them back (HACK for fixing parity with v1 and getting regression tests passing)
182	for event in events:	1✔
183	parsed_data = event.pop("data")	1✔
184	encoded_data = self.encode_data(parsed_data)	1✔
185	self.set_data(event, encoded_data)	1✔
186	return events	1✔
187
188	def get_data(self, event: dict) -> str:	1✔
189	if self.kinesis_namespace:	1✔
190	return event["kinesis"]["data"]	1✔
191	else:
UNCOV 192	return event["data"]	×
193
194	def set_data(self, event: dict, data: bytes) -> None:	1✔
195	if self.kinesis_namespace:	1✔
196	event["kinesis"]["data"] = data	1✔
197	else:
UNCOV 198	event["data"] = data	×
199
200	def parse_data(self, raw_data: str) -> dict \| str:	1✔
201	decoded_data = base64.b64decode(raw_data)	1✔
202	return json.loads(decoded_data)	1✔
203
204	def encode_data(self, parsed_data: dict) -> str:	1✔
205	return base64.b64encode(json.dumps(parsed_data).encode()).decode()	1✔
206
207
208	def has_data_filter_criteria_parsed(parsed_filters: list[dict]) -> bool:	1✔
209	for filter in parsed_filters:	1✔
210	if "data" in filter:	1✔
211	return True	1✔
212	return False	1✔

localstack / localstack / 7056d8d9-7428-4543-abf2-c5dedb2c98c0

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous