From 6c290993a38fdf6e5d8283a0d152ddba7c723702 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Tue, 30 Sep 2025 16:22:57 +0800 Subject: [PATCH 1/9] move firehose class Signed-off-by: Xch1 --- .../providers/amazon/aws/hooks/firehose.py | 56 +++++++++++++++++++ .../providers/amazon/aws/hooks/kinesis.py | 33 ++++++----- .../unit/amazon/aws/hooks/test_kinesis.py | 2 +- 3 files changed, 73 insertions(+), 18 deletions(-) create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py new file mode 100644 index 0000000000000..2601527ea5a7c --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py @@ -0,0 +1,56 @@ +# +# 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. +"""This module contains AWS Firehose hook.""" + +from __future__ import annotations + +from collections.abc import Iterable + +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook + + +class FirehoseHook(AwsBaseHook): + """ + Interact with Amazon Kinesis Firehose. + + Provide thick wrapper around :external+boto3:py:class:`boto3.client("firehose") `. + + :param delivery_stream: Name of the delivery stream + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, delivery_stream: str, *args, **kwargs) -> None: + self.delivery_stream = delivery_stream + kwargs["client_type"] = "firehose" + super().__init__(*args, **kwargs) + + def put_records(self, records: Iterable): + """ + Write batch records to Kinesis Firehose. + + .. seealso:: + - :external+boto3:py:meth:`Firehose.Client.put_record_batch` + + :param records: list of records + """ + return self.get_conn().put_record_batch(DeliveryStreamName=self.delivery_stream, Records=records) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py index 2601527ea5a7c..c73dee831d1f5 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py @@ -19,12 +19,13 @@ from __future__ import annotations -from collections.abc import Iterable +import warnings -from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.providers.amazon.aws.hooks.firehose import FirehoseHook as _FirehoseHook -class FirehoseHook(AwsBaseHook): +class FirehoseHook(_FirehoseHook): """ Interact with Amazon Kinesis Firehose. @@ -37,20 +38,18 @@ class FirehoseHook(AwsBaseHook): .. seealso:: - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + ..deprecated:: + This hook was moved. Import from + :class:`airflow.providers.amazon.aws.hooks.firehose.FirehoseHook` + instead of kinesis.py """ - def __init__(self, delivery_stream: str, *args, **kwargs) -> None: - self.delivery_stream = delivery_stream - kwargs["client_type"] = "firehose" + def __init__(self, *args, **kwargs) -> None: + warnings.warn( + "Importing FirehoseHook from kinesis.py is deprecated " + "and will be removed in a future release. " + "Please import it from firehose.py instead.", + AirflowProviderDeprecationWarning, + stacklevel=2, + ) super().__init__(*args, **kwargs) - - def put_records(self, records: Iterable): - """ - Write batch records to Kinesis Firehose. - - .. seealso:: - - :external+boto3:py:meth:`Firehose.Client.put_record_batch` - - :param records: list of records - """ - return self.get_conn().put_record_batch(DeliveryStreamName=self.delivery_stream, Records=records) diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py index 1080aa5379e31..58ec480379076 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py @@ -22,7 +22,7 @@ import boto3 from moto import mock_aws -from airflow.providers.amazon.aws.hooks.kinesis import FirehoseHook +from airflow.providers.amazon.aws.hooks.firehose import FirehoseHook @mock_aws From 02058f4286f866b2adde9ae4e4b5089bc32ba860 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Tue, 30 Sep 2025 17:26:07 +0800 Subject: [PATCH 2/9] add Kinesishook Signed-off-by: Xch1 --- .../providers/amazon/aws/hooks/kinesis.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py index c73dee831d1f5..bba49e5affa59 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py @@ -22,6 +22,7 @@ import warnings from airflow.exceptions import AirflowProviderDeprecationWarning +from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook from airflow.providers.amazon.aws.hooks.firehose import FirehoseHook as _FirehoseHook @@ -53,3 +54,95 @@ def __init__(self, *args, **kwargs) -> None: stacklevel=2, ) super().__init__(*args, **kwargs) + + +class KinesisHook(AwsBaseHook): + """ + Interact with Amazon Kinesis. + + Provide thin wrapper around :external+boto3:py:class:`boto3.client("kinesis") `. + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "kinesis" + super().__init__(*args, **kwargs) + + def create_stream(self, stream_name: str, shard_count: int) -> dict: + """ + Create a Kinesis stream. + + .. seealso:: + - :external+boto3:py:meth:`Kinesis.Client.create_stream` + + :param stream_name: Name of the stream to create. + :param shard_count: Number of shards for the stream. + :return: Response from the create_stream call. + """ + return self.get_conn().create_stream(StreamName=stream_name, ShardCount=shard_count) + + def put_record( + self, + stream_name: str, + data: bytes, + partition_key: str, + explicit_hash_key: str | None = None, + sequence_number_for_ordering: str | None = None, + ) -> dict: + """ + Put a single record into a Kinesis stream. + + .. seealso:: + - :external+boto3:py:meth:`Kinesis.Client.put_record` + + :param stream_name: Name of the stream to put the record into. + :param data: The data blob to put into the record, which is base64-encoded when the blob is serialized. + :param partition_key: Determines which shard in the stream the data record is assigned to. + :param explicit_hash_key: The hash value used to explicitly determine the shard the data record is assigned to. + :param sequence_number_for_ordering: Guarantees strictly increasing sequence numbers, for puts from the same client and to the same partition key. + :return: Response from the put_record call. + """ + params = { + "StreamName": stream_name, + "Data": data, + "PartitionKey": partition_key, + } + if explicit_hash_key: + params["ExplicitHashKey"] = explicit_hash_key + if sequence_number_for_ordering: + params["SequenceNumberForOrdering"] = sequence_number_for_ordering + + return self.get_conn().put_record(**params) + + def put_records(self, stream_name: str, records: list[dict]) -> dict: + """ + Put multiple records into a Kinesis stream. + + .. seealso:: + - :external+boto3:py:meth:`Kinesis.Client.put_records` + + :param stream_name: Name of the stream to put the records into. + :param records: List of records to put. Each record should be a dict with 'Data' and 'PartitionKey' keys. + :return: Response from the put_records call. + """ + return self.get_conn().put_records(Records=records, StreamName=stream_name) + + def delete_stream(self, stream_name: str, enforce_consumer_deletion: bool = False) -> dict: + """ + Delete a Kinesis stream. + + .. seealso:: + - :external+boto3:py:meth:`Kinesis.Client.delete_stream` + + :param stream_name: Name of the stream to delete. + :param enforce_consumer_deletion: If this parameter is set to true, the stream is deleted even if it has registered consumers. + :return: Response from the delete_stream call. + """ + return self.get_conn().delete_stream( + StreamName=stream_name, EnforceConsumerDeletion=enforce_consumer_deletion + ) From bfa97c8d1c1d7ec31260ba6ef119096d0749796b Mon Sep 17 00:00:00 2001 From: Xch1 Date: Tue, 30 Sep 2025 17:28:17 +0800 Subject: [PATCH 3/9] rename test_kinesis to test_firehose Signed-off-by: Xch1 --- .../unit/amazon/aws/hooks/{test_kinesis.py => test_firehose.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename providers/amazon/tests/unit/amazon/aws/hooks/{test_kinesis.py => test_firehose.py} (100%) diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_firehose.py similarity index 100% rename from providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py rename to providers/amazon/tests/unit/amazon/aws/hooks/test_firehose.py From f6e80fc5421820a20b235761351cd28ce66d6584 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Wed, 1 Oct 2025 09:28:57 +0800 Subject: [PATCH 4/9] add test_kinesis Signed-off-by: Xch1 --- .../unit/amazon/aws/hooks/test_kinesis.py | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py new file mode 100644 index 0000000000000..2866bbdd377a6 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py @@ -0,0 +1,29 @@ +# +# 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 + +from moto import mock_aws + +from airflow.providers.amazon.aws.hooks.kinesis import KinesisHook + + +class TestKinesisHook: + @mock_aws + def test_get_conn(self): + hook = KinesisHook(aws_conn_id="aws_default") + assert hook.get_conn() is not None From c02e41ad76b0fbf7334450120d5876d6d1ba32e4 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Wed, 1 Oct 2025 10:37:28 +0800 Subject: [PATCH 5/9] modify deprecate wapper Signed-off-by: Xch1 --- .../amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py | 1 + 1 file changed, 1 insertion(+) diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py index fc5654961def8..949ff487aea03 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_hooks_signature.py @@ -45,6 +45,7 @@ "EmrHook": {"emr_conn_id"}, "EmrContainerHook": {"virtual_cluster_id"}, "FirehoseHook": {"delivery_stream"}, + "_FirehoseHook": {"delivery_stream"}, "GlueJobHook": { "job_name", "concurrent_run_limit", From e2bc830b939f35d13adb0d6eb088564a095fdd19 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Wed, 1 Oct 2025 10:52:31 +0800 Subject: [PATCH 6/9] fix provider yaml and file check Signed-off-by: Xch1 --- providers/amazon/provider.yaml | 5 ++++- .../src/airflow/providers/amazon/get_provider_info.py | 6 +++++- scripts/in_container/run_provider_yaml_files_check.py | 1 + 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 6115efd5faaa4..1f2d0b4205ed1 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -629,9 +629,12 @@ hooks: - airflow.providers.amazon.aws.hooks.glue - airflow.providers.amazon.aws.hooks.glue_crawler - airflow.providers.amazon.aws.hooks.glue_catalog - - integration-name: Amazon Kinesis Data Firehose + - integration-name: Amazon Kinesis Data Stream python-modules: - airflow.providers.amazon.aws.hooks.kinesis + - integration-name: Amazon Kinesis Data Firehose + python-modules: + - airflow.providers.amazon.aws.hooks.firehose - integration-name: AWS Lambda python-modules: - airflow.providers.amazon.aws.hooks.lambda_function 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 fc546e5eb6a24..c4813094eee60 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -695,9 +695,13 @@ def get_provider_info(): ], }, { - "integration-name": "Amazon Kinesis Data Firehose", + "integration-name": "Amazon Kinesis Data Stream", "python-modules": ["airflow.providers.amazon.aws.hooks.kinesis"], }, + { + "integration-name": "Amazon Kinesis Data Firehose", + "python-modules": ["airflow.providers.amazon.aws.hooks.firehose"], + }, { "integration-name": "AWS Lambda", "python-modules": ["airflow.providers.amazon.aws.hooks.lambda_function"], diff --git a/scripts/in_container/run_provider_yaml_files_check.py b/scripts/in_container/run_provider_yaml_files_check.py index f2d8d4619603c..2b1ac8e6c67e5 100755 --- a/scripts/in_container/run_provider_yaml_files_check.py +++ b/scripts/in_container/run_provider_yaml_files_check.py @@ -72,6 +72,7 @@ "airflow.providers.google.cloud.operators.automl.AutoMLTablesListTableSpecsOperator", "airflow.providers.google.cloud.operators.automl.AutoMLTablesUpdateDatasetOperator", "airflow.providers.google.cloud.operators.automl.AutoMLDeployModelOperator", + "airflow.providers.amazon.aws.hooks.kinesis.FirehoseHook", ] if __name__ != "__main__": From adc618e8c29e74d50457a68ec80a420e9678ee96 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Wed, 1 Oct 2025 11:32:09 +0800 Subject: [PATCH 7/9] modify reStructuredText Signed-off-by: Xch1 --- .../amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py index bba49e5affa59..6f9be3d773925 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py @@ -39,7 +39,7 @@ class FirehoseHook(_FirehoseHook): .. seealso:: - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` - ..deprecated:: + .. deprecated:: This hook was moved. Import from :class:`airflow.providers.amazon.aws.hooks.firehose.FirehoseHook` instead of kinesis.py From 7f427b31ae4618fbba338af981287fb12db1868e Mon Sep 17 00:00:00 2001 From: Xch1 Date: Thu, 2 Oct 2025 09:10:16 +0800 Subject: [PATCH 8/9] add return type Signed-off-by: Xch1 --- .../amazon/src/airflow/providers/amazon/aws/hooks/firehose.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py index 2601527ea5a7c..fbe419fbd22c9 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py @@ -44,7 +44,7 @@ def __init__(self, delivery_stream: str, *args, **kwargs) -> None: kwargs["client_type"] = "firehose" super().__init__(*args, **kwargs) - def put_records(self, records: Iterable): + def put_records(self, records: Iterable) -> dict: """ Write batch records to Kinesis Firehose. From 818d92c5f96a782c42beca073ad5ca6a88ac8ef6 Mon Sep 17 00:00:00 2001 From: Xch1 Date: Wed, 15 Oct 2025 09:33:33 +0800 Subject: [PATCH 9/9] remove warp boto3 API methods Signed-off-by: Xch1 --- .../providers/amazon/aws/hooks/kinesis.py | 74 ------------------- 1 file changed, 74 deletions(-) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py index 6f9be3d773925..d022d49ac8531 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/kinesis.py @@ -72,77 +72,3 @@ class KinesisHook(AwsBaseHook): def __init__(self, *args, **kwargs) -> None: kwargs["client_type"] = "kinesis" super().__init__(*args, **kwargs) - - def create_stream(self, stream_name: str, shard_count: int) -> dict: - """ - Create a Kinesis stream. - - .. seealso:: - - :external+boto3:py:meth:`Kinesis.Client.create_stream` - - :param stream_name: Name of the stream to create. - :param shard_count: Number of shards for the stream. - :return: Response from the create_stream call. - """ - return self.get_conn().create_stream(StreamName=stream_name, ShardCount=shard_count) - - def put_record( - self, - stream_name: str, - data: bytes, - partition_key: str, - explicit_hash_key: str | None = None, - sequence_number_for_ordering: str | None = None, - ) -> dict: - """ - Put a single record into a Kinesis stream. - - .. seealso:: - - :external+boto3:py:meth:`Kinesis.Client.put_record` - - :param stream_name: Name of the stream to put the record into. - :param data: The data blob to put into the record, which is base64-encoded when the blob is serialized. - :param partition_key: Determines which shard in the stream the data record is assigned to. - :param explicit_hash_key: The hash value used to explicitly determine the shard the data record is assigned to. - :param sequence_number_for_ordering: Guarantees strictly increasing sequence numbers, for puts from the same client and to the same partition key. - :return: Response from the put_record call. - """ - params = { - "StreamName": stream_name, - "Data": data, - "PartitionKey": partition_key, - } - if explicit_hash_key: - params["ExplicitHashKey"] = explicit_hash_key - if sequence_number_for_ordering: - params["SequenceNumberForOrdering"] = sequence_number_for_ordering - - return self.get_conn().put_record(**params) - - def put_records(self, stream_name: str, records: list[dict]) -> dict: - """ - Put multiple records into a Kinesis stream. - - .. seealso:: - - :external+boto3:py:meth:`Kinesis.Client.put_records` - - :param stream_name: Name of the stream to put the records into. - :param records: List of records to put. Each record should be a dict with 'Data' and 'PartitionKey' keys. - :return: Response from the put_records call. - """ - return self.get_conn().put_records(Records=records, StreamName=stream_name) - - def delete_stream(self, stream_name: str, enforce_consumer_deletion: bool = False) -> dict: - """ - Delete a Kinesis stream. - - .. seealso:: - - :external+boto3:py:meth:`Kinesis.Client.delete_stream` - - :param stream_name: Name of the stream to delete. - :param enforce_consumer_deletion: If this parameter is set to true, the stream is deleted even if it has registered consumers. - :return: Response from the delete_stream call. - """ - return self.get_conn().delete_stream( - StreamName=stream_name, EnforceConsumerDeletion=enforce_consumer_deletion - )