diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index de693ee7d3abb..28477b7de7a18 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -630,9 +630,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/aws/hooks/firehose.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/firehose.py new file mode 100644 index 0000000000000..fbe419fbd22c9 --- /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) -> dict: + """ + 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..d022d49ac8531 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,14 @@ from __future__ import annotations -from collections.abc import Iterable +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 -class FirehoseHook(AwsBaseHook): +class FirehoseHook(_FirehoseHook): """ Interact with Amazon Kinesis Firehose. @@ -37,20 +39,36 @@ 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` +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. - :param records: list of records - """ - return self.get_conn().put_record_batch(DeliveryStreamName=self.delivery_stream, Records=records) + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "kinesis" + super().__init__(*args, **kwargs) 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/providers/amazon/tests/unit/amazon/aws/hooks/test_firehose.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_firehose.py new file mode 100644 index 0000000000000..58ec480379076 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_firehose.py @@ -0,0 +1,61 @@ +# +# 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 uuid + +import boto3 +from moto import mock_aws + +from airflow.providers.amazon.aws.hooks.firehose import FirehoseHook + + +@mock_aws +class TestFirehoseHook: + def test_get_conn_returns_a_boto3_connection(self): + hook = FirehoseHook( + aws_conn_id="aws_default", delivery_stream="test_airflow", region_name="us-east-1" + ) + assert hook.get_conn() is not None + + def test_insert_batch_records_kinesis_firehose(self): + boto3.client("s3").create_bucket(Bucket="kinesis-test") + hook = FirehoseHook( + aws_conn_id="aws_default", delivery_stream="test_airflow", region_name="us-east-1" + ) + + response = hook.get_conn().create_delivery_stream( + DeliveryStreamName="test_airflow", + S3DestinationConfiguration={ + "RoleARN": "arn:aws:iam::123456789012:role/firehose_delivery_role", + "BucketARN": "arn:aws:s3:::kinesis-test", + "Prefix": "airflow/", + "BufferingHints": {"SizeInMBs": 123, "IntervalInSeconds": 124}, + "CompressionFormat": "UNCOMPRESSED", + }, + ) + + stream_arn = response["DeliveryStreamARN"] + assert stream_arn == "arn:aws:firehose:us-east-1:123456789012:deliverystream/test_airflow" + + records = [{"Data": str(uuid.uuid4())} for _ in range(100)] + + response = hook.put_records(records) + + assert response["FailedPutCount"] == 0 + assert response["ResponseMetadata"]["HTTPStatusCode"] == 200 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", 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..2866bbdd377a6 100644 --- a/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_kinesis.py @@ -17,45 +17,13 @@ # under the License. from __future__ import annotations -import uuid - -import boto3 from moto import mock_aws -from airflow.providers.amazon.aws.hooks.kinesis import FirehoseHook +from airflow.providers.amazon.aws.hooks.kinesis import KinesisHook -@mock_aws -class TestFirehoseHook: - def test_get_conn_returns_a_boto3_connection(self): - hook = FirehoseHook( - aws_conn_id="aws_default", delivery_stream="test_airflow", region_name="us-east-1" - ) +class TestKinesisHook: + @mock_aws + def test_get_conn(self): + hook = KinesisHook(aws_conn_id="aws_default") assert hook.get_conn() is not None - - def test_insert_batch_records_kinesis_firehose(self): - boto3.client("s3").create_bucket(Bucket="kinesis-test") - hook = FirehoseHook( - aws_conn_id="aws_default", delivery_stream="test_airflow", region_name="us-east-1" - ) - - response = hook.get_conn().create_delivery_stream( - DeliveryStreamName="test_airflow", - S3DestinationConfiguration={ - "RoleARN": "arn:aws:iam::123456789012:role/firehose_delivery_role", - "BucketARN": "arn:aws:s3:::kinesis-test", - "Prefix": "airflow/", - "BufferingHints": {"SizeInMBs": 123, "IntervalInSeconds": 124}, - "CompressionFormat": "UNCOMPRESSED", - }, - ) - - stream_arn = response["DeliveryStreamARN"] - assert stream_arn == "arn:aws:firehose:us-east-1:123456789012:deliverystream/test_airflow" - - records = [{"Data": str(uuid.uuid4())} for _ in range(100)] - - response = hook.put_records(records) - - assert response["FailedPutCount"] == 0 - assert response["ResponseMetadata"]["HTTPStatusCode"] == 200 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__":