527eb572-51be-41bc-a2a3-1b2142a8a8f1

Committed 26 Feb 2025 03:36PM UTC coverage: 86.885% (+0.02%) from 86.862%

Build # 527eb572-51be-41bc-a2a3-1b2142a8a8f1

Build Type

push

circleci

Committed by

web-flow

Commit Message

[ESM] Add backoff between poll events calls (#12304)

Coverage Stats

25 of 25 new or added lines in 4 files covered. (100.0%)

48 existing lines in 12 files now uncovered.

61741 of 71061 relevant lines covered (86.88%)

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

    @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)"""
UNCOV 69	pass	×
70
71	@abstractmethod	1✔
72	def poll_events(self) -> None:	1✔
73	"""Poll events polled from the event source and matching at least one filter criteria and invoke the target processor."""
UNCOV 74	pass	×
75
76	def close(self) -> None:	1✔
77	"""Closes a target poller alongside all associated internal polling/consuming clients.
78	Only implemented for supported pollers. Therefore, the default implementation is empty."""
79	pass	1✔
80
81	def send_events_to_dlq(self, events, context) -> None:	1✔
82	"""Send failed events to a DLQ configured on the source.
83	Only implemented for supported pollers. Therefore, the default implementation is empty."""
UNCOV 84	pass	×
85
86	def filter_events(self, events: list[dict]) -> list[dict]:	1✔
87	"""Filter events using the EventBridge event patterns:
88	https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-event-patterns.html"""
89	if len(self.filter_patterns) == 0:	1✔
90	return events	1✔
91
92	filtered_events = []	1✔
93	for event in events:	1✔
94	# TODO: add try/catch with default discard and error log for extra resilience
95	if any(matches_event(pattern, event) for pattern in self.filter_patterns):	1✔
96	filtered_events.append(event)	1✔
97	return filtered_events	1✔
98
99	def add_source_metadata(self, events: list[dict], extra_metadata=None) -> list[dict]:	1✔
100	"""Add event source metadata to each event for eventSource, eventSourceARN, and awsRegion.
101	This metadata is added after filtering: https://repost.aws/knowledge-center/eventbridge-filter-events-with-pipes
102	See "The following fields can't be used in event patterns":
103	https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-event-filtering.html
104	"""
105	for event in events:	1✔
106	event["eventSourceARN"] = self.source_arn	1✔
107	event["eventSource"] = self.event_source()	1✔
108	event["awsRegion"] = self.aws_region	1✔
109	event.update(self.extra_metadata())	1✔
110	return events	1✔
111
112	def extra_metadata(self) -> dict:	1✔
113	"""Default implementation that subclasses can override to customize"""
114	return {}	1✔
115
116
117	def has_batch_item_failures(	1✔
118	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
119	) -> bool:
120	"""Returns False if no batch item failures are present and True otherwise (i.e., including parse exceptions)."""
121	# TODO: validate correct behavior upon exceptions
122	try:	1✔
123	failed_items_ids = parse_batch_item_failures(result, valid_item_ids)	1✔
124	return len(failed_items_ids) > 0	1✔
125	except (KeyError, ValueError):	1✔
126	return True	1✔
127
128
129	def get_batch_item_failures(	1✔
130	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
131	) -> list[str] \| None:
132	"""
133	Returns a list of failed batch item IDs. If an empty list is returned, then the batch should be considered as a complete success.
134
135	If `None` is returned, the batch should be considered a complete failure.
136	"""
137	try:	1✔
138	failed_items_ids = parse_batch_item_failures(result, valid_item_ids)	1✔
139	return failed_items_ids	1✔
140	except (KeyError, ValueError):	1✔
141	return None	1✔
142
143
144	def parse_batch_item_failures(	1✔
145	result: dict \| str \| None, valid_item_ids: set[str] \| None = None
146	) -> list[str]:
147	"""
148	Parses a partial batch failure response, that looks like this: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html
149
150	{
151	"batchItemFailures": [
152	{
153	"itemIdentifier": "id2"
154	},
155	{
156	"itemIdentifier": "id4"
157	}
158	]
159	}
160
161	If the response returns an empty list, then the batch should be considered as a complete success. If an exception
162	is raised, the batch should be considered a complete failure.
163
164	Pipes partial batch failure: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-pipes-batching-concurrency.html
165	Lambda ESM with SQS: https://docs.aws.amazon.com/lambda/latest/dg/services-sqs-errorhandling.html
166	Special cases: https://repost.aws/knowledge-center/lambda-sqs-report-batch-item-failures
167	Kinesis: https://docs.aws.amazon.com/lambda/latest/dg/services-kinesis-batchfailurereporting.html
168
169	:param result: the process status (e.g., invocation result from Lambda)
170	:param valid_item_ids: the set of valid item ids in the batch
171	:raises KeyError: if the itemIdentifier value is missing or not in the batch
172	:raises Exception: any other exception related to parsing (e.g., JSON parser error)
173	:return: a list of item IDs that failed
174	"""
175	if not result:	1✔
176	return []	1✔
177
178	if isinstance(result, dict):	1✔
179	partial_batch_failure = result	1✔
180	else:
UNCOV 181	partial_batch_failure = json.loads(result)	×
182
183	if not partial_batch_failure:	1✔
UNCOV 184	return []	×
185
186	batch_item_failures = partial_batch_failure.get("batchItemFailures")	1✔
187
188	if not batch_item_failures:	1✔
189	return []	1✔
190
191	failed_items = []	1✔
192	for item in batch_item_failures:	1✔
193	if "itemIdentifier" not in item:	1✔
194	raise KeyError(f"missing itemIdentifier in batchItemFailure record {item}")	1✔
195
196	item_identifier = item["itemIdentifier"]	1✔
197	if not item_identifier:	1✔
198	raise ValueError("itemIdentifier cannot be empty or null")	1✔
199
200	# Optionally validate whether the item_identifier is part of the batch
201	if valid_item_ids and item_identifier not in valid_item_ids:	1✔
202	raise KeyError(f"itemIdentifier '{item_identifier}' not in the batch")	1✔
203
204	failed_items.append(item_identifier)	1✔
205
206	return failed_items	1✔

localstack / localstack / 527eb572-51be-41bc-a2a3-1b2142a8a8f1

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