This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 30800029afd Add Kinesis Data Streams trigger (#71135)
30800029afd is described below
commit 30800029afd6465a7383aa040799710568142962
Author: Aaron Chen <[email protected]>
AuthorDate: Wed Sep 16 13:27:45 2026 +0800
Add Kinesis Data Streams trigger (#71135)
* Add Kinesis Data Streams trigger
* Update KinesisTrigger parameters and improve error handling in tests
* Refactor KinesisTrigger to use cached properties for checkpoint keys and
update checkpoint methods to be asynchronous
* support `.aget` and `.aset` method
---
providers/amazon/provider.yaml | 3 +
.../providers/amazon/aws/triggers/kinesis.py | 355 ++++++++++
.../airflow/providers/amazon/get_provider_info.py | 4 +
.../tests/unit/amazon/aws/triggers/test_kinesis.py | 733 +++++++++++++++++++++
4 files changed, 1095 insertions(+)
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 97c218c2f84..4ed477b213e 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -845,6 +845,9 @@ triggers:
- integration-name: AWS Lambda
python-modules:
- airflow.providers.amazon.aws.triggers.lambda_function
+ - integration-name: Amazon Kinesis Data Stream
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.kinesis
- integration-name: Amazon Managed Workflows for Apache Airflow (MWAA)
python-modules:
- airflow.providers.amazon.aws.triggers.mwaa
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/kinesis.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/kinesis.py
new file mode 100644
index 00000000000..e21e577f076
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/kinesis.py
@@ -0,0 +1,355 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import json
+from collections.abc import AsyncIterator
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.amazon.aws.hooks.kinesis import KinesisHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.triggers.base import BaseEventTrigger, TriggerEvent
+else:
+ from airflow.triggers.base import ( # type: ignore
+ BaseTrigger as BaseEventTrigger,
+ TriggerEvent,
+ )
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import BaseAwsConnection
+
+_CHECKPOINT_KEY_PREFIX = "kinesis_shard_sequence_numbers"
+_EXPIRED_ITERATOR_RETRIES = 2
+_ITERATOR_TYPES_WITHOUT_EXTRA_ARGS = frozenset({"LATEST", "TRIM_HORIZON"})
+
+
+class KinesisTrigger(BaseEventTrigger):
+ """
+ Wait asynchronously for records on an Amazon Kinesis Data Stream.
+
+ The trigger is long-running and emits one event for each non-empty shard
response. Record data is
+ base64-encoded in the event payload and must be decoded by the consumer.
Delivery is best-effort:
+ a triggerer failure can cause records to be repeated or missed around the
failure window.
+
+ When Airflow provides an asset state store for a single watched asset, the
trigger checkpoints the
+ last sequence number read from each shard. The same asset and stream
identity share one logical cursor;
+ do not configure multiple watchers that require independent progress for
the same stream on one asset.
+
+ :param stream_name: Name of the Kinesis Data Stream to watch.
+ :param aws_conn_id: AWS connection id.
+ :param shard_iterator_type: Position used when a shard has no checkpoint.
``LATEST`` only sees records
+ that arrive after the watcher starts; ``TRIM_HORIZON`` starts from the
oldest retained record.
+ Only these types are supported because ``AT_TIMESTAMP`` requires a
timestamp and sequence-number
+ types require a starting sequence number for each shard. Checkpoints
handle resuming each shard.
+ :param batch_size: Maximum records per ``GetRecords`` call and trigger
event. Must be between 1 and
+ 10,000. Record data is base64-encoded before it is stored in the
metadata database, so use a
+ conservative value for large records.
+ :param waiter_delay: Seconds between complete polling sweeps. Must be less
than the five-minute
+ shard iterator lifetime. Kinesis permits at most five ``GetRecords``
calls per second per shard.
+ When reading a backlog with ``TRIM_HORIZON``, draining ``N`` records
from one shard takes roughly
+ ``ceil(N / batch_size) * waiter_delay`` seconds when calls return full
batches; with the defaults,
+ 10,000 records take about 1,000 seconds.
+ :param region_name: AWS region for the Kinesis client.
+ :param verify: Whether to verify SSL certificates, or the path to a CA
bundle.
+ :param botocore_config: Botocore configuration passed to the Kinesis
client.
+ """
+
+ def __init__(
+ self,
+ stream_name: str,
+ aws_conn_id: str | None = "aws_default",
+ shard_iterator_type: str = "LATEST",
+ batch_size: int = 100,
+ waiter_delay: int = 10,
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict | None = None,
+ ) -> None:
+ super().__init__()
+ if shard_iterator_type not in _ITERATOR_TYPES_WITHOUT_EXTRA_ARGS:
+ raise ValueError(
+ "shard_iterator_type must be one of "
+ f"{sorted(_ITERATOR_TYPES_WITHOUT_EXTRA_ARGS)}; got
{shard_iterator_type!r}"
+ )
+ if not 1 <= batch_size <= 10_000:
+ raise ValueError("batch_size must be between 1 and 10000")
+ if not 0 < waiter_delay < 300:
+ raise ValueError("waiter_delay must be between 1 and 299 seconds")
+
+ self.stream_name = stream_name
+ self.aws_conn_id = aws_conn_id
+ self.shard_iterator_type = shard_iterator_type
+ self.batch_size = batch_size
+ self.waiter_delay = waiter_delay
+ self.region_name = region_name
+ self.verify = verify
+ self.botocore_config = botocore_config
+ self._checkpoint_warning_logged = False
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ return (
+ self.__class__.__module__ + "." + self.__class__.__qualname__,
+ {
+ "stream_name": self.stream_name,
+ "aws_conn_id": self.aws_conn_id,
+ "shard_iterator_type": self.shard_iterator_type,
+ "batch_size": self.batch_size,
+ "waiter_delay": self.waiter_delay,
+ "region_name": self.region_name,
+ "verify": self.verify,
+ "botocore_config": self.botocore_config,
+ },
+ )
+
+ @property
+ def hook(self) -> KinesisHook:
+ return KinesisHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+ @cached_property
+ def _asset_store_checkpoint_key(self) -> str:
+ identity = json.dumps(
+ {
+ "stream_name": self.stream_name,
+ "aws_conn_id": self.aws_conn_id,
+ "region_name": self.region_name,
+ },
+ sort_keys=True,
+ separators=(",", ":"),
+ ).encode()
+ return
f"{_CHECKPOINT_KEY_PREFIX}:{hashlib.sha256(identity).hexdigest()}"
+
+ def _log_checkpoint_warning_once(self, message: str) -> None:
+ if self._checkpoint_warning_logged:
+ return
+ self.log.warning(message)
+ self._checkpoint_warning_logged = True
+
+ async def _load_checkpoint(self) -> dict[str, str]:
+ store = getattr(self, "asset_state_store", None)
+ if store is None:
+ self._log_checkpoint_warning_once(
+ "Kinesis checkpointing is unavailable; using an in-memory
cursor"
+ )
+ return {}
+
+ try:
+ # aget/aset landed in Airflow 3.3.2; 3.3.0 and 3.3.1 only expose
the blocking API.
+ if hasattr(store, "aget"):
+ checkpoint = await
store.aget(self._asset_store_checkpoint_key, default={})
+ else:
+ checkpoint = await asyncio.to_thread(store.get,
self._asset_store_checkpoint_key, default={})
+ checkpoint = checkpoint or {}
+ except ValueError:
+ self._log_checkpoint_warning_once(
+ "Kinesis checkpointing requires a single watched asset; using
an in-memory cursor"
+ )
+ return {}
+
+ if not isinstance(checkpoint, dict) or not all(
+ isinstance(shard_id, str) and isinstance(sequence_number, str)
+ for shard_id, sequence_number in checkpoint.items()
+ ):
+ self._log_checkpoint_warning_once(
+ "Kinesis checkpoint data is invalid; using the configured
initial position"
+ )
+ return {}
+ return dict(checkpoint)
+
+ async def _save_checkpoint(self, sequence_numbers: dict[str, str]) -> None:
+ store = getattr(self, "asset_state_store", None)
+ if store is None:
+ self._log_checkpoint_warning_once(
+ "Kinesis checkpointing is unavailable; using an in-memory
cursor"
+ )
+ return
+
+ try:
+ if hasattr(store, "aset"):
+ await store.aset(self._asset_store_checkpoint_key,
dict(sequence_numbers))
+ else:
+ await asyncio.to_thread(
+ store.set,
+ self._asset_store_checkpoint_key,
+ dict(sequence_numbers),
+ )
+ except ValueError:
+ self._log_checkpoint_warning_once(
+ "Kinesis checkpointing requires a single watched asset; using
an in-memory cursor"
+ )
+
+ async def _find_shard_ids(self, client: BaseAwsConnection) -> list[str]:
+ paginator = client.get_paginator("list_shards")
+ shard_ids: list[str] = []
+ async for page in paginator.paginate(StreamName=self.stream_name):
+ shard_ids.extend(shard["ShardId"] for shard in page["Shards"])
+ return shard_ids
+
+ async def _get_shard_iterator(
+ self,
+ client: BaseAwsConnection,
+ shard_id: str,
+ after_sequence_number: str | None,
+ fallback_iterator_type: str,
+ ) -> str:
+ request: dict[str, Any] = {"StreamName": self.stream_name, "ShardId":
shard_id}
+ if after_sequence_number:
+ request.update(
+ ShardIteratorType="AFTER_SEQUENCE_NUMBER",
+ StartingSequenceNumber=after_sequence_number,
+ )
+ else:
+ request["ShardIteratorType"] = fallback_iterator_type
+
+ try:
+ response = await client.get_shard_iterator(**request)
+ except client.exceptions.InvalidArgumentException:
+ if not after_sequence_number:
+ raise
+ self.log.warning(
+ "Stored Kinesis checkpoint for shard %s is no longer valid;
using the configured initial position",
+ shard_id,
+ )
+ response = await client.get_shard_iterator(
+ StreamName=self.stream_name,
+ ShardId=shard_id,
+ ShardIteratorType=fallback_iterator_type,
+ )
+ return response["ShardIterator"]
+
+ async def _get_records(
+ self,
+ client: BaseAwsConnection,
+ shard_id: str,
+ shard_iterator: str,
+ after_sequence_number: str | None,
+ fallback_iterator_type: str,
+ ) -> dict[str, Any]:
+ for _ in range(_EXPIRED_ITERATOR_RETRIES):
+ try:
+ return await client.get_records(
+ ShardIterator=shard_iterator,
+ Limit=self.batch_size,
+ )
+ except client.exceptions.ExpiredIteratorException:
+ shard_iterator = await self._get_shard_iterator(
+ client,
+ shard_id,
+ after_sequence_number,
+ fallback_iterator_type,
+ )
+
+ return await client.get_records(
+ ShardIterator=shard_iterator,
+ Limit=self.batch_size,
+ )
+
+ @staticmethod
+ def _build_event_records(shard_id: str, records: list[dict[str, Any]]) ->
list[dict[str, Any]]:
+ event_records = []
+ for record in records:
+ timestamp = record.get("ApproximateArrivalTimestamp")
+ event_record = {
+ "ShardId": shard_id,
+ "SequenceNumber": record["SequenceNumber"],
+ "PartitionKey": record["PartitionKey"],
+ "ApproximateArrivalTimestamp": timestamp.isoformat() if
timestamp else None,
+ "Data": base64.b64encode(record["Data"]).decode("ascii"),
+ }
+ event_records.append(event_record)
+ return event_records
+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ loaded_sequence_numbers = await self._load_checkpoint()
+
+ async with await self.hook.get_async_conn() as client:
+ shard_ids = await self._find_shard_ids(client)
+ known_shard_ids = set(shard_ids)
+ sequence_numbers = {
+ shard_id: sequence_number
+ for shard_id, sequence_number in
loaded_sequence_numbers.items()
+ if shard_id in known_shard_ids
+ }
+ checkpoint_dirty = sequence_numbers != loaded_sequence_numbers
+ iterators: dict[str, str] = {}
+ fallback_iterator_types: dict[str, str] = {}
+
+ for shard_id in shard_ids:
+ iterators[shard_id] = await self._get_shard_iterator(
+ client,
+ shard_id,
+ sequence_numbers.get(shard_id),
+ self.shard_iterator_type,
+ )
+
+ while True:
+ for shard_id, shard_iterator in list(iterators.items()):
+ try:
+ response = await self._get_records(
+ client,
+ shard_id,
+ shard_iterator,
+ sequence_numbers.get(shard_id),
+ fallback_iterator_types.get(shard_id,
self.shard_iterator_type),
+ )
+ except
client.exceptions.ProvisionedThroughputExceededException:
+ self.log.warning("Kinesis read throughput exceeded for
shard %s", shard_id)
+ continue
+
+ next_shard_iterator = response.get("NextShardIterator")
+ records = response.get("Records", [])
+ if records:
+ sequence_numbers[shard_id] =
records[-1]["SequenceNumber"]
+ checkpoint_dirty = True
+ yield TriggerEvent(
+ {
+ "status": "success",
+ "message_batch":
self._build_event_records(shard_id, records),
+ }
+ )
+
+ if next_shard_iterator is None:
+ for child in response.get("ChildShards", []):
+ child_id = child["ShardId"]
+ if child_id not in iterators:
+ fallback_iterator_types[child_id] =
"TRIM_HORIZON"
+ iterators[child_id] = await
self._get_shard_iterator(
+ client,
+ child_id,
+ None,
+ fallback_iterator_types[child_id],
+ )
+ iterators.pop(shard_id, None)
+ fallback_iterator_types.pop(shard_id, None)
+ else:
+ iterators[shard_id] = next_shard_iterator
+
+ if checkpoint_dirty:
+ await self._save_checkpoint(sequence_numbers)
+ checkpoint_dirty = False
+
+ await asyncio.sleep(self.waiter_delay)
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 0a093ba5b5c..82ed31ca7a8 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -944,6 +944,10 @@ def get_provider_info():
"integration-name": "AWS Lambda",
"python-modules":
["airflow.providers.amazon.aws.triggers.lambda_function"],
},
+ {
+ "integration-name": "Amazon Kinesis Data Stream",
+ "python-modules":
["airflow.providers.amazon.aws.triggers.kinesis"],
+ },
{
"integration-name": "Amazon Managed Workflows for Apache
Airflow (MWAA)",
"python-modules":
["airflow.providers.amazon.aws.triggers.mwaa"],
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_kinesis.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_kinesis.py
new file mode 100644
index 00000000000..f42bdee7aac
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_kinesis.py
@@ -0,0 +1,733 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import datetime
+from types import SimpleNamespace
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+
+from airflow.providers.amazon.aws.triggers.kinesis import KinesisTrigger
+from airflow.triggers.base import TriggerEvent
+
+MODULE = "airflow.providers.amazon.aws.triggers.kinesis"
+STREAM_NAME = "test-stream"
+AWS_CONN_ID = "test-aws-conn"
+REGION_NAME = "us-east-1"
+SHARD_ID = "shardId-000000000000"
+CHILD_SHARD_ID = "shardId-000000000001"
+
+
+class ExpiredIteratorException(Exception):
+ pass
+
+
+class InvalidArgumentException(Exception):
+ pass
+
+
+class ProvisionedThroughputExceededException(Exception):
+ pass
+
+
+class StopPolling(Exception):
+ pass
+
+
+class AsyncPages:
+ def __init__(self, pages):
+ self._pages = iter(pages)
+
+ def __aiter__(self):
+ return self
+
+ async def __anext__(self):
+ try:
+ return next(self._pages)
+ except StopIteration:
+ raise StopAsyncIteration
+
+
+class AsyncClientContext:
+ def __init__(self, client):
+ self.client = client
+
+ async def __aenter__(self):
+ return self.client
+
+ async def __aexit__(self, exc_type, exc_value, traceback):
+ return None
+
+
[email protected]
+def trigger():
+ return KinesisTrigger(
+ stream_name=STREAM_NAME,
+ aws_conn_id=AWS_CONN_ID,
+ shard_iterator_type="LATEST",
+ batch_size=100,
+ waiter_delay=10,
+ region_name=REGION_NAME,
+ verify=True,
+ botocore_config={"retries": {"max_attempts": 3}},
+ )
+
+
+def create_client(pages):
+ client = mock.MagicMock(spec=["exceptions", "get_paginator",
"get_records", "get_shard_iterator"])
+ client.exceptions = SimpleNamespace(
+ ExpiredIteratorException=ExpiredIteratorException,
+ InvalidArgumentException=InvalidArgumentException,
+
ProvisionedThroughputExceededException=ProvisionedThroughputExceededException,
+ )
+ client.get_records = AsyncMock()
+ client.get_shard_iterator = AsyncMock()
+ paginator = mock.MagicMock(spec=["paginate"])
+ paginator.paginate.return_value = AsyncPages(pages)
+ client.get_paginator.return_value = paginator
+ return client, paginator
+
+
+def configure_hook(hook_property, client):
+ hook = mock.MagicMock(spec=["get_async_conn"])
+ hook.get_async_conn = AsyncMock(return_value=AsyncClientContext(client))
+ hook_property.return_value = hook
+ return hook
+
+
+def configure_checkpoint_store(trigger, checkpoint, *, supports_async=False):
+ methods = ["get", "set", "aget", "aset"] if supports_async else ["get",
"set"]
+ store = mock.MagicMock(spec=methods)
+ store.get.return_value = checkpoint
+ if supports_async:
+ store.aget = AsyncMock(spec=["__call__"], return_value=checkpoint)
+ store.aset = AsyncMock(spec=["__call__"])
+ trigger.asset_state_store = store
+ return store
+
+
+def create_record(sequence_number="1", data=b"message", *,
include_timestamp=True):
+ record = {
+ "SequenceNumber": sequence_number,
+ "PartitionKey": "partition-key",
+ "Data": data,
+ }
+ if include_timestamp:
+ record["ApproximateArrivalTimestamp"] = datetime.datetime(
+ 2026, 8, 4, 12, 30, tzinfo=datetime.timezone.utc
+ )
+ return record
+
+
+def test_serialize(trigger):
+ assert trigger.serialize() == (
+ f"{MODULE}.KinesisTrigger",
+ {
+ "stream_name": STREAM_NAME,
+ "aws_conn_id": AWS_CONN_ID,
+ "shard_iterator_type": "LATEST",
+ "batch_size": 100,
+ "waiter_delay": 10,
+ "region_name": REGION_NAME,
+ "verify": True,
+ "botocore_config": {"retries": {"max_attempts": 3}},
+ },
+ )
+
+
[email protected](
+ ("kwargs", "message"),
+ [
+ ({"shard_iterator_type": "AT_TIMESTAMP"}, "shard_iterator_type must be
one of"),
+ ({"batch_size": 0}, "batch_size must be between"),
+ ({"batch_size": 10_001}, "batch_size must be between"),
+ ({"waiter_delay": 0}, "waiter_delay must be between"),
+ ({"waiter_delay": 300}, "waiter_delay must be between"),
+ ],
+)
+def test_invalid_arguments_raise(kwargs, message):
+ with pytest.raises(ValueError, match=message):
+ KinesisTrigger(stream_name=STREAM_NAME, **kwargs)
+
+
[email protected](f"{MODULE}.KinesisHook", autospec=True)
+def test_hook(mock_hook, trigger):
+ assert trigger.hook == mock_hook.return_value
+ mock_hook.assert_called_once_with(
+ aws_conn_id=AWS_CONN_ID,
+ region_name=REGION_NAME,
+ verify=True,
+ config={"retries": {"max_attempts": 3}},
+ )
+
+
[email protected]("field", ["stream_name", "aws_conn_id",
"region_name"])
+def test_checkpoint_key_uses_stream_identity(trigger, field):
+ kwargs = trigger.serialize()[1]
+ kwargs[field] = f"different-{field}"
+ other = KinesisTrigger(**kwargs)
+
+ assert trigger._asset_store_checkpoint_key !=
other._asset_store_checkpoint_key
+
+
[email protected](
+ ("field", "value"),
+ [
+ ("shard_iterator_type", "TRIM_HORIZON"),
+ ("batch_size", 200),
+ ("waiter_delay", 5),
+ ("verify", False),
+ ("botocore_config", {"retries": {"max_attempts": 5}}),
+ ],
+)
+def test_checkpoint_key_ignores_polling_and_client_options(trigger, field,
value):
+ kwargs = trigger.serialize()[1]
+ kwargs[field] = value
+ other = KinesisTrigger(**kwargs)
+
+ assert trigger._asset_store_checkpoint_key ==
other._asset_store_checkpoint_key
+
+
[email protected](f"{MODULE}.hashlib.sha256", autospec=True)
+def test_checkpoint_key_is_cached(mock_sha256, trigger):
+ mock_sha256.return_value.hexdigest.return_value = "digest"
+
+ first_key = trigger._asset_store_checkpoint_key
+ second_key = trigger._asset_store_checkpoint_key
+
+ assert first_key == second_key == "kinesis_shard_sequence_numbers:digest"
+ mock_sha256.assert_called_once()
+
+
[email protected]
[email protected]("supports_async", [True, False], ids=["async",
"sync_fallback"])
[email protected](f"{MODULE}.asyncio.to_thread", autospec=True)
+async def test_load_and_save_checkpoint(mock_to_thread, trigger,
supports_async):
+ checkpoint = {SHARD_ID: "123"}
+ store = configure_checkpoint_store(trigger, checkpoint,
supports_async=supports_async)
+ mock_to_thread.side_effect = lambda func, *args, **kwargs: func(*args,
**kwargs)
+
+ assert await trigger._load_checkpoint() == checkpoint
+ await trigger._save_checkpoint(checkpoint)
+
+ if supports_async:
+
store.aget.assert_awaited_once_with(trigger._asset_store_checkpoint_key,
default={})
+
store.aset.assert_awaited_once_with(trigger._asset_store_checkpoint_key,
checkpoint)
+ store.get.assert_not_called()
+ store.set.assert_not_called()
+ mock_to_thread.assert_not_called()
+ else:
+ store.get.assert_called_once_with(trigger._asset_store_checkpoint_key,
default={})
+ store.set.assert_called_once_with(trigger._asset_store_checkpoint_key,
checkpoint)
+ assert mock_to_thread.await_args_list == [
+ mock.call(store.get, trigger._asset_store_checkpoint_key,
default={}),
+ mock.call(store.set, trigger._asset_store_checkpoint_key,
checkpoint),
+ ]
+
+
[email protected]
[email protected]("supports_async", [True, False], ids=["async",
"sync_fallback"])
[email protected](
+ "checkpoint",
+ [
+ ["not-a-dict"],
+ {1: "123"},
+ {SHARD_ID: 123},
+ ],
+)
[email protected](KinesisTrigger, "log", new_callable=mock.PropertyMock)
+async def test_invalid_checkpoint_falls_back_to_initial_position(
+ mock_log, trigger, checkpoint, supports_async
+):
+ store = configure_checkpoint_store(trigger, checkpoint,
supports_async=supports_async)
+
+ assert await trigger._load_checkpoint() == {}
+ mock_log.return_value.warning.assert_called_once()
+ if supports_async:
+
store.aget.assert_awaited_once_with(trigger._asset_store_checkpoint_key,
default={})
+ store.get.assert_not_called()
+
+
[email protected]
[email protected]("missing_attribute", [True, False])
[email protected](KinesisTrigger, "log", new_callable=mock.PropertyMock)
+async def test_checkpoint_falls_back_to_memory_without_store(mock_log,
trigger, missing_attribute):
+ if missing_attribute:
+ trigger.__dict__.pop("asset_state_store", None)
+ else:
+ trigger.asset_state_store = None
+
+ assert await trigger._load_checkpoint() == {}
+ await trigger._save_checkpoint({SHARD_ID: "123"})
+
+ mock_log.return_value.warning.assert_called_once()
+
+
[email protected]
[email protected]("supports_async", [True, False], ids=["async",
"sync_fallback"])
[email protected](KinesisTrigger, "log", new_callable=mock.PropertyMock)
+async def test_checkpoint_falls_back_to_memory_for_multiple_assets(mock_log,
trigger, supports_async):
+ store = configure_checkpoint_store(trigger, {},
supports_async=supports_async)
+ get_method = store.aget if supports_async else store.get
+ set_method = store.aset if supports_async else store.set
+ get_method.side_effect = ValueError
+
+ assert await trigger._load_checkpoint() == {}
+
+ get_method.side_effect = None
+ set_method.side_effect = ValueError
+ await trigger._save_checkpoint({SHARD_ID: "123"})
+
+ mock_log.return_value.warning.assert_called_once()
+ if supports_async:
+
store.aget.assert_awaited_once_with(trigger._asset_store_checkpoint_key,
default={})
+
store.aset.assert_awaited_once_with(trigger._asset_store_checkpoint_key,
{SHARD_ID: "123"})
+ store.get.assert_not_called()
+ store.set.assert_not_called()
+
+
[email protected]
+async def test_find_shards_follows_pagination(trigger):
+ client, paginator = create_client(
+ [
+ {"Shards": [{"ShardId": SHARD_ID}]},
+ {"Shards": [{"ShardId": CHILD_SHARD_ID}]},
+ ]
+ )
+
+ assert await trigger._find_shard_ids(client) == [SHARD_ID, CHILD_SHARD_ID]
+ paginator.paginate.assert_called_once_with(StreamName=STREAM_NAME)
+
+
[email protected]
[email protected](
+ ("sequence_number", "fallback_type", "expected_request"),
+ [
+ (
+ "123",
+ "LATEST",
+ {
+ "StreamName": STREAM_NAME,
+ "ShardId": SHARD_ID,
+ "ShardIteratorType": "AFTER_SEQUENCE_NUMBER",
+ "StartingSequenceNumber": "123",
+ },
+ ),
+ (
+ None,
+ "TRIM_HORIZON",
+ {
+ "StreamName": STREAM_NAME,
+ "ShardId": SHARD_ID,
+ "ShardIteratorType": "TRIM_HORIZON",
+ },
+ ),
+ ],
+)
+async def test_get_shard_iterator(trigger, sequence_number, fallback_type,
expected_request):
+ client, _ = create_client([])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator"}
+
+ assert await trigger._get_shard_iterator(client, SHARD_ID,
sequence_number, fallback_type) == "iterator"
+ client.get_shard_iterator.assert_awaited_once_with(**expected_request)
+
+
[email protected]
[email protected](KinesisTrigger, "log", new_callable=mock.PropertyMock)
+async def test_stale_checkpoint_falls_back_to_initial_position(mock_log,
trigger):
+ client, _ = create_client([])
+ client.get_shard_iterator.side_effect = [
+ InvalidArgumentException,
+ {"ShardIterator": "fallback-iterator"},
+ ]
+
+ assert await trigger._get_shard_iterator(client, SHARD_ID, "123",
"LATEST") == "fallback-iterator"
+ assert client.get_shard_iterator.await_args_list == [
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="AFTER_SEQUENCE_NUMBER",
+ StartingSequenceNumber="123",
+ ),
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="LATEST",
+ ),
+ ]
+ mock_log.return_value.warning.assert_called_once()
+
+
[email protected]
+async def test_invalid_initial_position_is_not_suppressed(trigger):
+ client, _ = create_client([])
+ client.get_shard_iterator.side_effect = InvalidArgumentException
+
+ with pytest.raises(InvalidArgumentException):
+ await trigger._get_shard_iterator(client, SHARD_ID, None, "LATEST")
+
+
+def test_build_event_records():
+ assert KinesisTrigger._build_event_records(
+ SHARD_ID,
+ [create_record(), create_record("2", b"second",
include_timestamp=False)],
+ ) == [
+ {
+ "ShardId": SHARD_ID,
+ "SequenceNumber": "1",
+ "PartitionKey": "partition-key",
+ "ApproximateArrivalTimestamp": "2026-08-04T12:30:00+00:00",
+ "Data": "bWVzc2FnZQ==",
+ },
+ {
+ "ShardId": SHARD_ID,
+ "SequenceNumber": "2",
+ "PartitionKey": "partition-key",
+ "ApproximateArrivalTimestamp": None,
+ "Data": "c2Vjb25k",
+ },
+ ]
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_run_yields_events_and_keeps_running(mock_sleep, hook_property,
trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator-1"}
+ client.get_records.side_effect = [
+ {"Records": [create_record("1")], "NextShardIterator": "iterator-2"},
+ {"Records": [create_record("2")], "NextShardIterator": "iterator-3"},
+ ]
+ configure_hook(hook_property, client)
+ store = configure_checkpoint_store(trigger, {})
+
+ generator = trigger.run()
+ first_event = await anext(generator)
+ second_event = await anext(generator)
+ await generator.aclose()
+
+ assert first_event.payload["message_batch"][0]["SequenceNumber"] == "1"
+ assert second_event.payload["message_batch"][0]["SequenceNumber"] == "2"
+ store.set.assert_called_once_with(trigger._asset_store_checkpoint_key,
{SHARD_ID: "1"})
+ mock_sleep.assert_awaited_once_with(10)
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_run_checkpoints_once_after_sweep(mock_sleep, hook_property,
trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}, {"ShardId":
CHILD_SHARD_ID}]}])
+ client.get_shard_iterator.side_effect = [
+ {"ShardIterator": "iterator-1"},
+ {"ShardIterator": "iterator-2"},
+ ]
+ client.get_records.side_effect = [
+ {"Records": [create_record("1")], "NextShardIterator":
"iterator-1-next"},
+ {"Records": [create_record("2")], "NextShardIterator":
"iterator-2-next"},
+ ]
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ store = configure_checkpoint_store(trigger, {})
+
+ generator = trigger.run()
+ assert await anext(generator) == TriggerEvent(
+ {
+ "status": "success",
+ "message_batch": KinesisTrigger._build_event_records(SHARD_ID,
[create_record("1")]),
+ }
+ )
+ assert await anext(generator) == TriggerEvent(
+ {
+ "status": "success",
+ "message_batch":
KinesisTrigger._build_event_records(CHILD_SHARD_ID, [create_record("2")]),
+ }
+ )
+ with pytest.raises(StopPolling):
+ await anext(generator)
+
+ store.set.assert_called_once_with(
+ trigger._asset_store_checkpoint_key,
+ {SHARD_ID: "1", CHILD_SHARD_ID: "2"},
+ )
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_checkpoint_prunes_expired_shards_on_startup(mock_sleep,
hook_property, trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator"}
+ client.get_records.return_value = {"Records": [], "NextShardIterator":
"next-iterator"}
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ store = configure_checkpoint_store(trigger, {SHARD_ID: "1",
"expired-shard": "2"})
+
+ with pytest.raises(StopPolling):
+ await anext(trigger.run())
+
+ store.set.assert_called_once_with(trigger._asset_store_checkpoint_key,
{SHARD_ID: "1"})
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_checkpoint_retains_closed_parent(mock_sleep, hook_property,
trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.side_effect = [
+ {"ShardIterator": "parent-iterator"},
+ {"ShardIterator": "child-iterator"},
+ ]
+ client.get_records.side_effect = [
+ {
+ "Records": [],
+ "ChildShards": [{"ShardId": CHILD_SHARD_ID}],
+ },
+ {
+ "Records": [create_record("child-sequence")],
+ "NextShardIterator": "child-next-iterator",
+ },
+ ]
+ mock_sleep.side_effect = [None, StopPolling]
+ configure_hook(hook_property, client)
+ store = configure_checkpoint_store(trigger, {SHARD_ID: "parent-sequence"})
+
+ generator = trigger.run()
+ event = await anext(generator)
+ with pytest.raises(StopPolling):
+ await anext(generator)
+
+ assert event.payload["message_batch"][0]["SequenceNumber"] ==
"child-sequence"
+ store.set.assert_called_once_with(
+ trigger._asset_store_checkpoint_key,
+ {SHARD_ID: "parent-sequence", CHILD_SHARD_ID: "child-sequence"},
+ )
+
+
[email protected]
[email protected]("iterator_type", ["LATEST", "TRIM_HORIZON"])
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_startup_without_checkpoint_honours_configured_position(
+ mock_sleep, hook_property, iterator_type
+):
+ trigger = KinesisTrigger(
+ stream_name=STREAM_NAME,
+ shard_iterator_type=iterator_type,
+ )
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator"}
+ client.get_records.return_value = {"Records": [], "NextShardIterator":
"next-iterator"}
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {})
+
+ with pytest.raises(StopPolling):
+ await anext(trigger.run())
+
+ client.get_shard_iterator.assert_awaited_once_with(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType=iterator_type,
+ )
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def
test_repeated_expired_iterators_are_recovered_from_checkpoint(mock_sleep,
hook_property, trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.side_effect = [
+ {"ShardIterator": "initial-iterator"},
+ {"ShardIterator": "first-recovered-iterator"},
+ {"ShardIterator": "second-recovered-iterator"},
+ ]
+ client.get_records.side_effect = [
+ ExpiredIteratorException,
+ ExpiredIteratorException,
+ {"Records": [create_record("124")], "NextShardIterator":
"next-iterator"},
+ ]
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {SHARD_ID: "123"})
+
+ event = await anext(trigger.run())
+
+ assert event.payload["message_batch"][0]["SequenceNumber"] == "124"
+ assert client.get_shard_iterator.await_args_list == [
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="AFTER_SEQUENCE_NUMBER",
+ StartingSequenceNumber="123",
+ ),
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="AFTER_SEQUENCE_NUMBER",
+ StartingSequenceNumber="123",
+ ),
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="AFTER_SEQUENCE_NUMBER",
+ StartingSequenceNumber="123",
+ ),
+ ]
+ mock_sleep.assert_not_awaited()
+
+
[email protected]
+async def test_expired_iterator_retry_limit(trigger):
+ client, _ = create_client([])
+ client.get_records.side_effect = ExpiredIteratorException
+ client.get_shard_iterator.return_value = {"ShardIterator":
"recovered-iterator"}
+
+ with pytest.raises(ExpiredIteratorException):
+ await trigger._get_records(
+ client,
+ SHARD_ID,
+ "initial-iterator",
+ "123",
+ "LATEST",
+ )
+
+ assert client.get_records.await_count == 3
+ assert client.get_shard_iterator.await_count == 2
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_child_iterator_recovery_preserves_trim_horizon(mock_sleep,
hook_property, trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.side_effect = [
+ {"ShardIterator": "parent-iterator"},
+ {"ShardIterator": "child-iterator"},
+ {"ShardIterator": "recovered-child-iterator"},
+ ]
+ client.get_records.side_effect = [
+ {
+ "Records": [],
+ "ChildShards": [{"ShardId": CHILD_SHARD_ID}],
+ },
+ ExpiredIteratorException,
+ {
+ "Records": [create_record("1")],
+ "NextShardIterator": "child-next-iterator",
+ },
+ ]
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {})
+
+ event = await anext(trigger.run())
+
+ assert event.payload["message_batch"][0]["SequenceNumber"] == "1"
+ assert client.get_shard_iterator.await_args_list[-2:] == [
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=CHILD_SHARD_ID,
+ ShardIteratorType="TRIM_HORIZON",
+ ),
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=CHILD_SHARD_ID,
+ ShardIteratorType="TRIM_HORIZON",
+ ),
+ ]
+ mock_sleep.assert_awaited_once_with(10)
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
[email protected](KinesisTrigger, "log", new_callable=mock.PropertyMock)
+async def test_throttling_is_logged_and_paced(mock_log, mock_sleep,
hook_property, trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator"}
+ client.get_records.side_effect = ProvisionedThroughputExceededException
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {})
+
+ with pytest.raises(StopPolling):
+ await anext(trigger.run())
+
+ mock_log.return_value.warning.assert_called_once()
+ mock_sleep.assert_awaited_once_with(10)
+
+
[email protected]
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_closed_shard_adds_each_child_once(mock_sleep, hook_property,
trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.side_effect = [
+ {"ShardIterator": "parent-iterator"},
+ {"ShardIterator": "child-iterator"},
+ ]
+ client.get_records.return_value = {
+ "Records": [],
+ "ChildShards": [
+ {"ShardId": CHILD_SHARD_ID},
+ {"ShardId": CHILD_SHARD_ID},
+ ],
+ }
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {})
+
+ with pytest.raises(StopPolling):
+ await anext(trigger.run())
+
+ assert client.get_shard_iterator.await_args_list == [
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=SHARD_ID,
+ ShardIteratorType="LATEST",
+ ),
+ mock.call(
+ StreamName=STREAM_NAME,
+ ShardId=CHILD_SHARD_ID,
+ ShardIteratorType="TRIM_HORIZON",
+ ),
+ ]
+
+
[email protected]
[email protected]("records", [[], [create_record()]])
[email protected](KinesisTrigger, "hook", new_callable=mock.PropertyMock)
[email protected](f"{MODULE}.asyncio.sleep", new_callable=AsyncMock)
+async def test_run_sleeps_after_empty_and_busy_sweeps(mock_sleep,
hook_property, records, trigger):
+ client, _ = create_client([{"Shards": [{"ShardId": SHARD_ID}]}])
+ client.get_shard_iterator.return_value = {"ShardIterator": "iterator"}
+ client.get_records.return_value = {
+ "Records": records,
+ "NextShardIterator": "next-iterator",
+ }
+ mock_sleep.side_effect = StopPolling
+ configure_hook(hook_property, client)
+ configure_checkpoint_store(trigger, {})
+
+ generator = trigger.run()
+ if records:
+ await anext(generator)
+ with pytest.raises(StopPolling):
+ await anext(generator)
+
+ mock_sleep.assert_awaited_once_with(10)
