Add Kinesis Data Streams trigger by aaron-y-chen · Pull Request #71135 · apache/airflow

aaron-y-chen · 2026-08-05T03:05:44Z

related: #52712

This is the first PR for Kinesis part of #52712, I will integrate the trigger with the common-messaging interface and add a system/e2e test in the next PR.

Summary

Add an asynchronous Kinesis Data Streams trigger for consuming records without occupying worker slots.

The trigger polls all shards, emits JSON-serializable record batches, handles iterator expiry, throttling, pagination, and in-run resharding, and checkpoints per-shard sequence numbers when an asset state store is available. New shards first discovered after a restart honor the configured initial position, so LATEST may skip records written during the downtime.

This is the provider-side foundation for the common-messaging integration tracked in #52712.

Was generative AI tooling used to co-author this PR?

Yes (please specify the tool below)

Generated-by: [GPT 5.6-sol] following the guidelines

aaron-y-chen · 2026-08-05T03:56:27Z

Those failures seem to be caused by the same issue:

flit build: error: unrecognized arguments: --no-setup-py

There is already a PR to fix it: #71118

aaron-y-chen · 2026-08-05T13:33:39Z

Waiting for another PR #71145 to be merged.

aaron-y-chen · 2026-08-06T01:20:30Z

Okay, all CI failures have been fixed.

aaron-y-chen · 2026-08-17T02:53:30Z

Hi @jason810496, I'd appreciate it if you could take a look when you have time. :)

vincbeck · 2026-08-17T21:10:41Z

@seanghaeli @ramitkataria

aaron-y-chen · 2026-08-19T09:13:25Z

Thanks for the reviews, I believe this version is more concrete than before :)

jason810496

Thanks for adding the new trigger to compelete #52712.

jason810496 · 2026-08-27T05:45:23Z

+    def _save_checkpoint(self, sequence_numbers: dict[str, str]) -> None:
+        store = getattr(self, "asset_state_store", None)


IIUC, this means we're calling a blocking call in the trigger run coroutine.
We should follow the same pattern (asyncio.to_thread) in #71387 to workaround.

Or wait until #72127 get merged and invoke the async native method from the accessor.

jason810496 · 2026-08-27T05:47:19Z

+        try:
+            store.set(self._build_checkpoint_key(), dict(sequence_numbers))
+        except ValueError:


This is the second PR that touched the Task Store within the trigger. #71387 is the first one.

cc @amoghrajesh to double check regarding the usage.
Just in case using Task Store within trigger isn't an expected pattern.

It's a good pattern to use (but slight correction: its asset state store not task state store), he Triggerer itself injects it at: https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/jobs/triggerer_job_runner.py#L1405-L1408, so a BaseEventTrigger attached to a watched asset is designed to read and write asset state. Watermarking a stream cursor is one of the use case.

jason810496 · 2026-08-27T05:50:59Z

+    @property
+    def hook(self) -> KinesisHook:


Not sure do we need to define as cached property or not? We might need to check the lifecycle of the Kinesis client. IIRC, this might be related to the https://lists.apache.org/thread/owp7869xmr5cyks9gx3lp66x53w1owyg discussion ( #71701 PR).

Just in case Kinesis client invoke eager external calls during construction.

jason810496 · 2026-08-27T05:55:00Z

+            config=self.botocore_config,
+        )
+
+    def _build_checkpoint_key(self) -> str:


Would it be better to make this as a cache property? (e.g. asset_store_checkpoint_key

amoghrajesh

Blocking due to sync access to asset state store

amoghrajesh · 2026-08-27T07:26:03Z

+    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]:
+        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,
+            )


A second ExpiredIteratorException will propagate through and kill the trigger.

amoghrajesh · 2026-08-27T07:26:28Z

+    :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.


AT_TIMESTAMP and AT_SEQUENCE_NUMBER are rejected?

boring-cyborg Bot added area:providers provider:amazon AWS/Amazon - related issues labels Aug 5, 2026

aaron-y-chen mentioned this pull request Aug 5, 2026

AIP-82: Add common messaging system providers #52712

Open

10 tasks

aaron-y-chen force-pushed the add-kinesis-trigger branch from 4e1733b to d7e4df5 Compare August 5, 2026 07:39

aaron-y-chen marked this pull request as ready for review August 5, 2026 09:18

aaron-y-chen requested a review from o-nikolas as a code owner August 5, 2026 09:18

Add Kinesis Data Streams trigger

5589794

aaron-y-chen force-pushed the add-kinesis-trigger branch from d7e4df5 to 5589794 Compare August 5, 2026 14:50

eladkal requested a review from vincbeck August 6, 2026 05:49

vincbeck reviewed Aug 10, 2026

View reviewed changes

Comment thread providers/amazon/tests/unit/amazon/aws/triggers/test_kinesis.py

eladkal requested a review from vincbeck August 17, 2026 20:09

ramitkataria reviewed Aug 18, 2026

View reviewed changes

Comment thread providers/amazon/src/airflow/providers/amazon/aws/triggers/kinesis.py Outdated

Comment thread providers/amazon/src/airflow/providers/amazon/aws/triggers/kinesis.py Outdated

Update KinesisTrigger parameters and improve error handling in tests

72b152f

eladkal requested a review from jason810496 August 21, 2026 14:54

jason810496 mentioned this pull request Aug 27, 2026

Add async asset store accessors for async tasks and watcher triggers #72127

Open

1 task

jason810496 reviewed Aug 27, 2026

View reviewed changes

jason810496 requested a review from amoghrajesh August 27, 2026 06:18

amoghrajesh requested changes Aug 27, 2026

View reviewed changes

		def _save_checkpoint(self, sequence_numbers: dict[str, str]) -> None:
		store = getattr(self, "asset_state_store", None)

		: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.

Uh oh!

Conversation

aaron-y-chen commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Was generative AI tooling used to co-author this PR?

Uh oh!

aaron-y-chen commented Aug 5, 2026

Uh oh!

aaron-y-chen commented Aug 5, 2026

Uh oh!

aaron-y-chen commented Aug 6, 2026

Uh oh!

Uh oh!

aaron-y-chen commented Aug 17, 2026

Uh oh!

vincbeck commented Aug 17, 2026

Uh oh!

Uh oh!

Uh oh!

aaron-y-chen commented Aug 19, 2026

Uh oh!

jason810496 left a comment

Choose a reason for hiding this comment

Uh oh!

jason810496 Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jason810496 Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amoghrajesh Aug 27, 2026

Choose a reason for hiding this comment

Uh oh!

jason810496 Aug 27, 2026

Choose a reason for hiding this comment

Uh oh!

jason810496 Aug 27, 2026

Choose a reason for hiding this comment

Uh oh!

amoghrajesh left a comment

Choose a reason for hiding this comment

Uh oh!

amoghrajesh Aug 27, 2026

Choose a reason for hiding this comment

Uh oh!

amoghrajesh Aug 27, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

aaron-y-chen commented Aug 5, 2026 •

edited

Loading

jason810496 Aug 27, 2026 •

edited

Loading

jason810496 Aug 27, 2026 •

edited

Loading