d166ec6d-e321-44a4-963c-9d9dd94aa89d

Committed 11 Mar 2025 05:46PM UTC coverage: 86.936% (+0.04%) from 86.901%

Build # d166ec6d-e321-44a4-963c-9d9dd94aa89d

Build Type

push

circleci

Committed by

web-flow

Commit Message

chore: improve test snapshot (#12367)

Coverage Stats

62152 of 71492 relevant lines covered (86.94%)

0.87 hits per line

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

94.9

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

import json
import logging
from abc import ABC, abstractmethod
from typing import Any

from botocore.client import BaseClient

from localstack.aws.api.pipes import PipeStateReason
from localstack.services.lambda_.event_source_mapping.event_processor import EventProcessor
from localstack.services.lambda_.event_source_mapping.noops_event_processor import (
    NoOpsEventProcessor,
)
from localstack.services.lambda_.event_source_mapping.pipe_utils import get_internal_client
from localstack.utils.aws.arns import parse_arn
from localstack.utils.event_matcher import matches_event


class EmptyPollResultsException(Exception):
    service: str
    source_arn: str

    def __init__(self, service: str = "", source_arn: str = ""):
        self.service = service
        self.source_arn = source_arn


class PipeStateReasonValues(PipeStateReason):
    USER_INITIATED = "USER_INITIATED"
    NO_RECORDS_PROCESSED = "No records processed"
    # TODO: add others (e.g., failure)


LOG = logging.getLogger(__name__)


class Poller(ABC):
    source_arn: str | None
    aws_region: str | None
    source_parameters: dict
    filter_patterns: list[dict[str, Any]]
    source_client: BaseClient

    # Target processor (e.g., Pipe, EventSourceMapping)
    processor: EventProcessor

    def __init__(
        self,
        source_arn: str | None = None,
        source_parameters: dict | None = None,
        source_client: BaseClient | None = None,
        processor: EventProcessor | None = None,
    ):
        # TODO: handle pollers without an ARN (e.g., Apache Kafka)
        if source_arn:
            self.source_arn = source_arn
            self.aws_region = parse_arn(source_arn)["region"]
            self.source_client = source_client or get_internal_client(source_arn)

        self.source_parameters = source_parameters or {}
        filters = self.source_parameters.get("FilterCriteria", {}).get("Filters", [])
        self.filter_patterns = [json.loads(event_filter["Pattern"]) for event_filter in filters]

        # Target processor
        self.processor = processor or NoOpsEventProcessor()

    @abstractmethod
    def event_source(self) -> str:
        """Return the event source metadata (e.g., aws:sqs)"""
        pass

    # TODO: create an abstract fetch_records method that all children should implement. This will unify how poller's internally retreive data from an event
    # source and make for much easier error handling.
    @abstractmethod
    def poll_events(self) -> None:
        """Poll events polled from the event source and matching at least one filter criteria and invoke the target processor."""
        pass

    def close(self) -> None:
        """Closes a target poller alongside all associated internal polling/consuming clients.
        Only implemented for supported pollers. Therefore, the default implementation is empty."""
        pass

    def send_events_to_dlq(self, events, context) -> None:
        """Send failed events to a DLQ configured on the source.
        Only implemented for supported pollers. Therefore, the default implementation is empty."""
        pass

    def filter_events(self, events: list[dict]) -> list[dict]:
        """Filter events using the EventBridge event patterns:
        https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-event-patterns.html"""
        if len(self.filter_patterns) == 0:
            return events

        filtered_events = []
        for event in events:
            # TODO: add try/catch with default discard and error log for extra resilience
            if any(matches_event(pattern, event) for pattern in self.filter_patterns):
                filtered_events.append(event)
        return filtered_events

    def add_source_metadata(self, events: list[dict], extra_metadata=None) -> list[dict]:
        """Add event source metadata to each event for eventSource, eventSourceARN, and awsRegion.
        This metadata is added after filtering: https://repost.aws/knowledge-center/eventbridge-filter-events-with-pipes
        See "The following fields can't be used in event patterns":
        https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html
        """
        for event in events:
            event["eventSourceARN"] = self.source_arn
            event["eventSource"] = self.event_source()
            event["awsRegion"] = self.aws_region
            event.update(self.extra_metadata())
        return events

    def extra_metadata(self) -> dict:
        """Default implementation that subclasses can override to customize"""
        return {}


def has_batch_item_failures(
    result: dict | str | None, valid_item_ids: set[str] | None = None
) -> bool:
    """Returns False if no batch item failures are present and True otherwise (i.e., including parse exceptions)."""
    # TODO: validate correct behavior upon exceptions
    try:
        failed_items_ids = parse_batch_item_failures(result, valid_item_ids)
        return len(failed_items_ids) > 0
    except (KeyError, ValueError):
        return True


def get_batch_item_failures(
    result: dict | str | None, valid_item_ids: set[str] | None = None
) -> list[str] | None:
    """
    Returns a list of failed batch item IDs. If an empty list is returned, then the batch should be considered as a complete success.

    If `None` is returned, the batch should be considered a complete failure.
    """
    try:
        failed_items_ids = parse_batch_item_failures(result, valid_item_ids)
        return failed_items_ids
    except (KeyError, ValueError):
        return None


def parse_batch_item_failures(
    result: dict | str | None, valid_item_ids: set[str] | None = None
) -> list[str]:
    """
    Parses a partial batch failure response, that looks like this: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html

        {
          "batchItemFailures": [
                {
                    "itemIdentifier": "id2"
                },
                {
                    "itemIdentifier": "id4"
                }
            ]
        }

    If the response returns an empty list, then the batch should be considered as a complete success. If an exception
    is raised, the batch should be considered a complete failure.

    Pipes partial batch failure: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html
    Lambda ESM with SQS: https://docs.aws.amazon.com/lambda/latest/dg/services-sqs-errorhandling.html
    Special cases: https://repost.aws/knowledge-center/lambda-sqs-report-batch-item-failures
    Kinesis: https://docs.aws.amazon.com/lambda/latest/dg/services-kinesis-batchfailurereporting.html

    :param result: the process status (e.g., invocation result from Lambda)
    :param valid_item_ids: the set of valid item ids in the batch
    :raises KeyError: if the itemIdentifier value is missing or not in the batch
    :raises Exception: any other exception related to parsing (e.g., JSON parser error)
    :return: a list of item IDs that failed
    """
    if not result:
        return []

    if isinstance(result, dict):
        partial_batch_failure = result
    else:
        partial_batch_failure = json.loads(result)

    if not partial_batch_failure:
        return []

    batch_item_failures = partial_batch_failure.get("batchItemFailures")

    if not batch_item_failures:
        return []

    failed_items = []
    for item in batch_item_failures:
        if "itemIdentifier" not in item:
            raise KeyError(f"missing itemIdentifier in batchItemFailure record {item}")

        item_identifier = item["itemIdentifier"]
        if not item_identifier:
            raise ValueError("itemIdentifier cannot be empty or null")

        # Optionally validate whether the item_identifier is part of the batch
        if valid_item_ids and item_identifier not in valid_item_ids:
            raise KeyError(f"itemIdentifier '{item_identifier}' not in the batch")

        failed_items.append(item_identifier)

    return failed_items

1	import json	1✔
2	import logging	1✔
3	from abc import ABC, abstractmethod	1✔
4	from typing import Any	1✔
5
6	from botocore.client import BaseClient	1✔
7
8	from localstack.aws.api.pipes import PipeStateReason	1✔
9	from localstack.services.lambda_.event_source_mapping.event_processor import EventProcessor	1✔
10	from localstack.services.lambda_.event_source_mapping.noops_event_processor import (	1✔
11	NoOpsEventProcessor,
12	)
13	from localstack.services.lambda_.event_source_mapping.pipe_utils import get_internal_client	1✔
14	from localstack.utils.aws.arns import parse_arn	1✔
15	from localstack.utils.event_matcher import matches_event	1✔
16
17
18	class EmptyPollResultsException(Exception):	1✔
19	service: str	1✔
20	source_arn: str	1✔
21
22	def __init__(self, service: str = "", source_arn: str = ""):	1✔
23	self.service = service	1✔
24	self.source_arn = source_arn	1✔
25
26
27	class PipeStateReasonValues(PipeStateReason):	1✔
28	USER_INITIATED = "USER_INITIATED"	1✔
29	NO_RECORDS_PROCESSED = "No records processed"	1✔
30	# TODO: add others (e.g., failure)
31
32
33	LOG = logging.getLogger(__name__)	1✔
34
35
36	class Poller(ABC):	1✔
37	source_arn: str \| None	1✔
38	aws_region: str \| None	1✔
39	source_parameters: dict	1✔
40	filter_patterns: list[dict[str, Any]]	1✔
41	source_client: BaseClient	1✔
42
43	# Target processor (e.g., Pipe, EventSourceMapping)
44	processor: EventProcessor	1✔
45
46	def __init__(	1✔
47	self,
48	source_arn: str \| None = None,
49	source_parameters: dict \| None = None,
50	source_client: BaseClient \| None = None,
51	processor: EventProcessor \| None = None,
52	):
53	# TODO: handle pollers without an ARN (e.g., Apache Kafka)
54	if source_arn:	1✔
55	self.source_arn = source_arn	1✔
56	self.aws_region = parse_arn(source_arn)["region"]	1✔
57	self.source_client = source_client or get_internal_client(source_arn)	1✔
58
59	self.source_parameters = source_parameters or {}	1✔
60	filters = self.source_parameters.get("FilterCriteria", {}).get("Filters", [])	1✔
61	self.filter_patterns = [json.loads(event_filter["Pattern"]) for event_filter in filters]	1✔
62
63	# Target processor
64	self.processor = processor or NoOpsEventProcessor()	1✔
65
66	@abstractmethod	1✔
67	def event_source(self) -> str:	1✔
68	"""Return the event source metadata (e.g., aws:sqs)"""
69	pass	×
70
71	# TODO: create an abstract fetch_records method that all children should implement. This will unify how poller's internally retreive data from an event
72	# source and make for much easier error handling.
73	@abstractmethod	1✔
74	def poll_events(self) -> None:	1✔
75	"""Poll events polled from the event source and matching at least one filter criteria and invoke the target processor."""
76	pass	×
77
78	def close(self) -> None:	1✔
79	"""Closes a target poller alongside all associated internal polling/consuming clients.
80	Only implemented for supported pollers. Therefore, the default implementation is empty."""
81	pass	1✔
82
83	def send_events_to_dlq(self, events, context) -> None:	1✔
84	"""Send failed events to a DLQ configured on the source.
85	Only implemented for supported pollers. Therefore, the default implementation is empty."""
86	pass	×
87
88	def filter_events(self, events: list[dict]) -> list[dict]:	1✔
89	"""Filter events using the EventBridge event patterns:
90	https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-event-patterns.html"""
91	if len(self.filter_patterns) == 0:	1✔
92	return events	1✔
93
94	filtered_events = []	1✔
95	for event in events:	1✔
96	# TODO: add try/catch with default discard and error log for extra resilience
97	if any(matches_event(pattern, event) for pattern in self.filter_patterns):	1✔
98	filtered_events.append(event)	1✔
99	return filtered_events	1✔
100
101	def add_source_metadata(self, events: list[dict], extra_metadata=None) -> list[dict]:	1✔
102	"""Add event source metadata to each event for eventSource, eventSourceARN, and awsRegion.
103	This metadata is added after filtering: https://repost.aws/knowledge-center/eventbridge-filter-events-with-pipes
104	See "The following fields can't be used in event patterns":
105	https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html
106	"""
107	for event in events:	1✔
108	event["eventSourceARN"] = self.source_arn	1✔
109	event["eventSource"] = self.event_source()	1✔
110	event["awsRegion"] = self.aws_region	1✔
111	event.update(self.extra_metadata())	1✔
112	return events	1✔
113
114	def extra_metadata(self) -> dict:	1✔
115	"""Default implementation that subclasses can override to customize"""
116	return {}	1✔
117
118
119	def has_batch_item_failures(	1✔
120	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
121	) -> bool:
122	"""Returns False if no batch item failures are present and True otherwise (i.e., including parse exceptions)."""
123	# TODO: validate correct behavior upon exceptions
124	try:	1✔
125	failed_items_ids = parse_batch_item_failures(result, valid_item_ids)	1✔
126	return len(failed_items_ids) > 0	1✔
127	except (KeyError, ValueError):	1✔
128	return True	1✔
129
130
131	def get_batch_item_failures(	1✔
132	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
133	) -> list[str] \| None:
134	"""
135	Returns a list of failed batch item IDs. If an empty list is returned, then the batch should be considered as a complete success.
136
137	If `None` is returned, the batch should be considered a complete failure.
138	"""
139	try:	1✔
140	failed_items_ids = parse_batch_item_failures(result, valid_item_ids)	1✔
141	return failed_items_ids	1✔
142	except (KeyError, ValueError):	1✔
143	return None	1✔
144
145
146	def parse_batch_item_failures(	1✔
147	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
148	) -> list[str]:
149	"""
150	Parses a partial batch failure response, that looks like this: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html
151
152	{
153	"batchItemFailures": [
154	{
155	"itemIdentifier": "id2"
156	},
157	{
158	"itemIdentifier": "id4"
159	}
160	]
161	}
162
163	If the response returns an empty list, then the batch should be considered as a complete success. If an exception
164	is raised, the batch should be considered a complete failure.
165
166	Pipes partial batch failure: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html
167	Lambda ESM with SQS: https://docs.aws.amazon.com/lambda/latest/dg/services-sqs-errorhandling.html
168	Special cases: https://repost.aws/knowledge-center/lambda-sqs-report-batch-item-failures
169	Kinesis: https://docs.aws.amazon.com/lambda/latest/dg/services-kinesis-batchfailurereporting.html
170
171	:param result: the process status (e.g., invocation result from Lambda)
172	:param valid_item_ids: the set of valid item ids in the batch
173	:raises KeyError: if the itemIdentifier value is missing or not in the batch
174	:raises Exception: any other exception related to parsing (e.g., JSON parser error)
175	:return: a list of item IDs that failed
176	"""
177	if not result:	1✔
178	return []	1✔
179
180	if isinstance(result, dict):	1✔
181	partial_batch_failure = result	1✔
182	else:
183	partial_batch_failure = json.loads(result)	×
184
185	if not partial_batch_failure:	1✔
186	return []	×
187
188	batch_item_failures = partial_batch_failure.get("batchItemFailures")	1✔
189
190	if not batch_item_failures:	1✔
191	return []	1✔
192
193	failed_items = []	1✔
194	for item in batch_item_failures:	1✔
195	if "itemIdentifier" not in item:	1✔
196	raise KeyError(f"missing itemIdentifier in batchItemFailure record {item}")	1✔
197
198	item_identifier = item["itemIdentifier"]	1✔
199	if not item_identifier:	1✔
200	raise ValueError("itemIdentifier cannot be empty or null")	1✔
201
202	# Optionally validate whether the item_identifier is part of the batch
203	if valid_item_ids and item_identifier not in valid_item_ids:	1✔
204	raise KeyError(f"itemIdentifier '{item_identifier}' not in the batch")	1✔
205
206	failed_items.append(item_identifier)	1✔
207
208	return failed_items	1✔

localstack / localstack / d166ec6d-e321-44a4-963c-9d9dd94aa89d

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