Merge pull request #16 from beam-pyio/feature/prepare-release

Preparation for release

Author: jaehyeon-kim

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 <!--next-version-placeholder-->
 
-## v0.1.0 (11/06/2024)
+## v0.1.0 (23/07/2024)
 
-- First release of `firehose_pyio`!
+✨NEW
+
+- Add a composite transform (`WriteToFirehose`) that puts records into a Firehose delivery stream in batch, using the [`put_record_batch`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/firehose/client/put_record_batch.html) method of the boto3 package.
+- Create a dedicated pipeline option (`FirehoseOptions`) that reads AWS related values (e.g. `aws_access_key_id`) from pipeline arguments.
+- Implement metric objects that record the total, succeeded and failed elements counts.
+- Add unit and integration testing cases. The [moto](https://github.com/getmoto/moto) and [localstack-utils](https://docs.localstack.cloud/user-guide/tools/testing-utils/) are used for unit and integration testing respectively. Also, a custom test client is created for testing retry behavior, which is not supported by the moto package.
+- Integrate with GitHub Actions by adding workflows for testing, documentation and release management.

README.md

@@ -6,7 +6,7 @@
 ![python](https://img.shields.io/badge/python-3.8%2C%203.9%2C%203.10%2C%203.11%2C%203.12-blue)
 ![os](https://img.shields.io/badge/OS-Ubuntu%2C%20Mac%2C%20Windows-purple)
 
-Apache Beam Python I/O connector for Amazon Firehose
+[Amazon Data Firehose](https://aws.amazon.com/firehose/) is a fully managed service for delivering real-time streaming data to destinations such as Amazon Simple Storage Service (Amazon S3), Amazon Redshift, Amazon OpenSearch Service and Amazon OpenSearch Serverless. The Apache Beam Python I/O connector for Amazon Data Firehose (`firehose_pyio`) provides a data sink feature that facilitates integration with those services.
 
 ## Installation
 
@@ -16,7 +16,61 @@ $ pip install firehose_pyio
 
 ## Usage
 
-- TODO
+The connector has the main composite transform ([`WriteToFirehose`](https://beam-pyio.github.io/firehose_pyio/autoapi/firehose_pyio/io/index.html#firehose_pyio.io.WriteToFirehose)), and it expects a list or tuple _PCollection_ element. If the element is a tuple, the tuple's first element is taken. If the element is not of the accepted types, you can apply the [`GroupIntoBatches`](https://beam.apache.org/documentation/transforms/python/aggregation/groupintobatches/) or [`BatchElements`](https://beam.apache.org/releases/pydoc/current/apache_beam.transforms.util.html#apache_beam.transforms.util.BatchElements) transform beforehand. Then, the element is sent into a Firehose delivery stream using the [`put_record_batch`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/firehose/client/put_record_batch.html) method of the boto3 package. Note that the above batch transforms can also be useful to overcome the API limitation listed below.
+
+- Each `PutRecordBatch` request supports up to 500 records. Each record in the request can be as large as 1,000 KB (before base64 encoding), up to a limit of 4 MB for the entire request. These limits cannot be changed.
+
+The transform also has options that control individual records or handling failed records.
+
+- _jsonify_ - A flag that indicates whether to convert a record into Json. Note that a record should be of _bytes_, _bytearray_ or file-like object, and, if it is not of a supported type (e.g. integer), we can convert it into a Json string by specifying this flag to _True_.
+- _multiline_ - A flag that indicates whether to add a new line character (`\n`) to each record. It is useful to save records into a _CSV_ or _Jsonline_ file.
+- _max_trials_ - The maximum number of trials when there is one or more failed records - it defaults to 3. Note that failed records after all trials are returned, which allows users to determine how to handle them subsequently.
+
+### Example
+
+If a _PCollection_ element is key-value pair (i.e. keyed stream), it can be batched in group using the `GroupIntoBatches` transform before it is connected into the main transform.
+
+```python
+import apache_beam as beam
+from apache_beam import GroupIntoBatches
+from firehose_pyio.io import WriteToFirehose
+
+with beam.Pipeline(options=pipeline_options) as p:
+    (
+        p
+        | beam.Create([(1, "one"), (2, "three"), (1, "two"), (2, "four")])
+        | GroupIntoBatches(batch_size=2)
+        | WriteToFirehose(
+            delivery_stream_name=delivery_stream_name,
+            jsonify=True,
+            multiline=True,
+            max_trials=3
+        )
+    )
+```
+
+For a list element (i.e. unkeyed stream), we can apply the `BatchElements` transform instead.
+
+```python
+import apache_beam as beam
+from apache_beam.transforms.util import BatchElements
+from firehose_pyio.io import WriteToFirehose
+
+with beam.Pipeline(options=pipeline_options) as p:
+    (
+        p
+        | beam.Create(["one", "two", "three", "four"])
+        | BatchElements(min_batch_size=2, max_batch_size=2)
+        | WriteToFirehose(
+            delivery_stream_name=delivery_stream_name,
+            jsonify=True,
+            multiline=True,
+            max_trials=3
+        )
+    )
+```
+
+See [this post](https://beam-pyio.github.io/blog/2024/firehose-pyio-intro/) for more examples.
 
 ## Contributing
 

examples/create_resources.py

@@ -0,0 +1,85 @@
+#
+# 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.
+#
+
+import json
+import boto3
+from botocore.exceptions import ClientError
+
+COMMON_NAME = "firehose-pyio-test"
+
+
+def create_destination_bucket(bucket_name):
+    client = boto3.client("s3")
+    suffix = client._client_config.region_name
+    client.create_bucket(Bucket=f"{bucket_name}-{suffix}")
+
+
+def create_firehose_iam_role(role_name):
+    assume_role_policy_document = json.dumps(
+        {
+            "Version": "2012-10-17",
+            "Statement": [
+                {
+                    "Effect": "Allow",
+                    "Principal": {"Service": "firehose.amazonaws.com"},
+                    "Action": "sts:AssumeRole",
+                }
+            ],
+        }
+    )
+    client = boto3.client("iam")
+    try:
+        return client.get_role(RoleName=role_name)
+    except ClientError as error:
+        if error.response["Error"]["Code"] == "NoSuchEntity":
+            resp = client.create_role(
+                RoleName=role_name, AssumeRolePolicyDocument=assume_role_policy_document
+            )
+            client.attach_role_policy(
+                RoleName=role_name,
+                PolicyArn="arn:aws:iam::aws:policy/AmazonS3FullAccess",
+            )
+            return resp
+
+
+def create_delivery_stream(delivery_stream_name, role_arn, bucket_name):
+    client = boto3.client("firehose")
+    suffix = client._client_config.region_name
+    try:
+        client.create_delivery_stream(
+            DeliveryStreamName=delivery_stream_name,
+            DeliveryStreamType="DirectPut",
+            S3DestinationConfiguration={
+                "RoleARN": role_arn,
+                "BucketARN": f"arn:aws:s3:::{bucket_name}-{suffix}",
+                "BufferingHints": {"SizeInMBs": 1, "IntervalInSeconds": 0},
+            },
+        )
+    except ClientError as error:
+        if error.response["Error"]["Code"] == "ResourceInUseException":
+            pass
+        else:
+            raise error
+
+
+if __name__ == "__main__":
+    print("create a destination bucket...")
+    create_destination_bucket(COMMON_NAME)
+    print("create an iam role...")
+    iam_resp = create_firehose_iam_role(COMMON_NAME)
+    print("create a delivery stream...")
+    create_delivery_stream(COMMON_NAME, iam_resp["Role"]["Arn"], COMMON_NAME)

examples/pipeline.py

@@ -0,0 +1,129 @@
+#
+# 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.
+#
+
+import argparse
+import datetime
+import random
+import string
+import logging
+import boto3
+import time
+
+import apache_beam as beam
+from apache_beam.transforms.util import BatchElements
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.pipeline_options import SetupOptions
+
+from firehose_pyio.io import WriteToFirehose
+
+
+def get_all_contents(bucket_name):
+    client = boto3.client("s3")
+    bucket_objects = client.list_objects_v2(
+        Bucket=f"{bucket_name}-{client._client_config.region_name}"
+    )
+    return bucket_objects.get("Contents") or []
+
+
+def delete_all_objects(bucket_name):
+    client = boto3.client("s3")
+    contents = get_all_contents(bucket_name)
+    for content in contents:
+        client.delete_object(
+            Bucket=f"{bucket_name}-{client._client_config.region_name}",
+            Key=content["Key"],
+        )
+
+
+def print_bucket_contents(bucket_name):
+    client = boto3.client("s3")
+    contents = get_all_contents(bucket_name)
+    for content in contents:
+        resp = client.get_object(
+            Bucket=f"{bucket_name}-{client._client_config.region_name}",
+            Key=content["Key"],
+        )
+        print(f"Key - {content['Key']}")
+        print(resp["Body"].read().decode())
+
+
+def create_records(n=100):
+    return [
+        {
+            "id": i,
+            "name": "".join(random.choices(string.ascii_letters, k=5)).lower(),
+            "created_at": datetime.datetime.now(),
+        }
+        for i in range(n)
+    ]
+
+
+def convert_ts(record: dict):
+    record["created_at"] = record["created_at"].isoformat(timespec="milliseconds")
+    return record
+
+
+def mask_secrets(d: dict):
+    return {k: (v if k.find("aws") < 0 else "x" * len(v)) for k, v in d.items()}
+
+
+def run(argv=None, save_main_session=True):
+    parser = argparse.ArgumentParser(description="Beam pipeline arguments")
+    parser.add_argument(
+        "--stream_name",
+        default="firehose-pyio-test",
+        type=str,
+        help="Delivery stream name",
+    )
+    parser.add_argument(
+        "--num_records", default="100", type=int, help="Number of records"
+    )
+    known_args, pipeline_args = parser.parse_known_args(argv)
+
+    pipeline_options = PipelineOptions(pipeline_args)
+    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
+    print(f"known_args - {known_args}")
+    print(f"pipeline options - {mask_secrets(pipeline_options.display_data())}")
+
+    with beam.Pipeline(options=pipeline_options) as p:
+        (
+            p
+            | "CreateElements" >> beam.Create(create_records(known_args.num_records))
+            | "DatetimeToStr" >> beam.Map(convert_ts)
+            | "BatchElements" >> BatchElements(min_batch_size=50)
+            | "WriteToFirehose"
+            >> WriteToFirehose(
+                delivery_stream_name=known_args.stream_name,
+                jsonify=True,
+                multiline=True,
+                max_trials=3,
+            )
+        )
+
+        logging.getLogger().setLevel(logging.WARN)
+        logging.info("Building pipeline ...")
+
+
+if __name__ == "__main__":
+    BUCKET_NAME = "firehose-pyio-test"
+    print(">> delete existing objects...")
+    delete_all_objects(BUCKET_NAME)
+    print(">> start pipeline...")
+    run()
+    time.sleep(1)
+    print(">> print bucket contents...")
+    print_bucket_contents(BUCKET_NAME)