Merge pull request #6 from beam-pyio/feature/implement-ptransform

Feature/implement ptransform

Author: jaehyeon-kim

.github/workflows/doc.yml

@@ -25,7 +25,7 @@ jobs:
         uses: ammaraskar/[email protected]
 
       - name: Upload artifacts
-        uses: actions/upload-artifact@v1
+        uses: actions/upload-artifact@v4
         with:
           name: html-docs
           path: docs/_build/html/
@@ -34,4 +34,4 @@ jobs:
         uses: peaceiris/actions-gh-pages@v3
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
-          publish_dir: docs/_build/html
+          publish_dir: docs/_build/html

.github/workflows/test.yml

@@ -6,7 +6,6 @@ on:
 
 jobs:
   test:
-
     runs-on: ${{ matrix.os }}
 
     strategy:
@@ -30,9 +29,9 @@ jobs:
         run: poetry install
 
       - name: Test with pytest
-        run: poetry run pytest tests/ --cov=dynamodb_pyio --cov-report=xml
+        run: poetry run pytest tests/ -svv --cov=dynamodb_pyio --cov-report=xml -m "not integration"
 
       - name: Use Codecov to track coverage
         uses: codecov/codecov-action@v3
         with:
-          files: ./coverage.xml   # coverage report
+          files: ./coverage.xml # coverage report

.gitignore

@@ -1,3 +1,5 @@
+.vscode
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

poetry.lock

@@ -28,7 +28,7 @@ apache-beam = {extras = ["aws"], version = ">=2.19.0"}
 pytest = ">=7.1.2,<8.0"
 pytest-cov = "^5.0.0"
 apache-beam = {extras = ["test"], version = ">=2.19.0"}
-moto = {extras = ["firehose", "s3", "sts"], version = "^5.0.9"}
+moto = {extras = ["all"], version = "^5.0.9"}
 localstack-utils = ">=1.0.0"
 
 [tool.pytest.ini_options]

pyproject.toml

@@ -0,0 +1,107 @@
+#
+# 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 typing
+import boto3
+from apache_beam.options import pipeline_options
+
+from dynamodb_pyio.options import DynamoDBOptions
+
+__all__ = ["DynamoDBClient", "DynamoDBClientError"]
+
+
+def get_http_error_code(exc):
+    if hasattr(exc, "response"):
+        return exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
+    return None
+
+
+class DynamoDBClientError(Exception):
+    def __init__(self, message=None, code=None):
+        self.message = message
+        self.code = code
+
+
+class DynamoDBClient(object):
+    """
+    Wrapper for boto3 library.
+    """
+
+    def __init__(self, options: typing.Union[DynamoDBOptions, dict]):
+        """Constructor of the DynamoDBClient.
+
+        Args:
+            options (Union[DynamoDBOptions, dict]): Options to create a boto3 DynamoDB client.
+        """
+        assert boto3 is not None, "Missing boto3 requirement"
+        if isinstance(options, pipeline_options.PipelineOptions):
+            options = options.view_as(DynamoDBOptions)
+            access_key_id = options.aws_access_key_id
+            secret_access_key = options.aws_secret_access_key
+            session_token = options.aws_session_token
+            endpoint_url = options.endpoint_url
+            use_ssl = not options.disable_ssl
+            region_name = options.region_name
+            api_version = options.api_version
+            verify = options.verify
+        else:
+            access_key_id = options.get("aws_access_key_id")
+            secret_access_key = options.get("aws_secret_access_key")
+            session_token = options.get("aws_session_token")
+            endpoint_url = options.get("endpoint_url")
+            use_ssl = not options.get("disable_ssl", False)
+            region_name = options.get("region_name")
+            api_version = options.get("api_version")
+            verify = options.get("verify")
+
+        session = boto3.session.Session()
+        self.resource = session.resource(
+            service_name="dynamodb",
+            region_name=region_name,
+            api_version=api_version,
+            use_ssl=use_ssl,
+            verify=verify,
+            endpoint_url=endpoint_url,
+            aws_access_key_id=access_key_id,
+            aws_secret_access_key=secret_access_key,
+            aws_session_token=session_token,
+        )
+
+    def put_items_batch(self, records: list, table_name: str, dedup_pkeys: list = None):
+        """Put records to an Amazon DynamoDB table using a batch writer object.
+
+        Args:
+            records (list): Records to send into an Amazon SQS queue.
+            table_name (str): Amazon DynamoDB table name.
+            dedup_pkeys (list, Optional): List of keys to be used for de-duplicating items in buffer.
+
+        Raises:
+            DynamoDBClientError: DynamoDB client error.
+
+        Returns:
+            (Object): Boto3 response message.
+        """
+
+        if not isinstance(records, list):
+            raise DynamoDBClientError("Records should be a list.")
+        try:
+            table = self.resource.Table(table_name)
+            with table.batch_writer(overwrite_by_pkeys=dedup_pkeys or []) as batch:
+                for record in records:
+                    batch.put_item(Item=record)
+        except Exception as e:
+            raise DynamoDBClientError(str(e), get_http_error_code(e))

src/dynamodb_pyio/boto3_client.py

@@ -1,2 +1,101 @@
-def my_fn():
-    return 1
+#
+# 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 logging
+import typing
+import apache_beam as beam
+from apache_beam import metrics
+from apache_beam.pvalue import PCollection
+
+from dynamodb_pyio.boto3_client import DynamoDBClient
+from dynamodb_pyio.options import DynamoDBOptions
+
+
+__all__ = ["WriteToDynamoDB"]
+
+
+class _DynamoDBWriteFn(beam.DoFn):
+    """Create the connector can send messages in batch to an Amazon DynamoDB table.
+
+    Args:
+        table_name (str): Amazon DynamoDB table name.
+        dedup_pkeys (list): List of keys to be used for de-duplicating records in buffer.
+        options (Union[DynamoDBOptions, dict]): Options to create a boto3 dynamodb client.
+    """
+
+    total_elements_count = metrics.Metrics.counter(
+        "_DynamoDBWriteFn", "total_elements_count"
+    )
+
+    def __init__(
+        self,
+        table_name: str,
+        dedup_pkeys: list,
+        options: typing.Union[DynamoDBOptions, dict],
+    ):
+        """Constructor of _DynamoDBWriteFn
+
+        Args:
+            table_name (str): Amazon DynamoDB table name.
+            dedup_pkeys (list): List of keys to be used for de-duplicating records in buffer.
+            options (Union[DynamoDBOptions, dict]): Options to create a boto3 dynamodb client.
+        """
+        super().__init__()
+        self.table_name = table_name
+        self.dedup_pkeys = dedup_pkeys
+        self.options = options
+
+    def start_bundle(self):
+        self.client = DynamoDBClient(self.options)
+
+    def process(self, element):
+        if isinstance(element, tuple):
+            element = element[1]
+        self.client.put_items_batch(element, self.table_name, self.dedup_pkeys)
+        self.total_elements_count.inc(len(element))
+        logging.info(f"total {len(element)} elements processed...")
+
+
+class WriteToDynamoDB(beam.PTransform):
+    """A transform that puts records to an Amazon DynamoDB table.
+
+    Takes an input PCollection and put them in batch using the boto3 package.
+    For more information, visit the `Boto3 Documentation <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/table/batch_writer.html>`__.
+
+    Note that, if the PCollection element is a tuple (i.e. keyed stream), only the value is used to put records in batch.
+
+    Args:
+        table_name (str): Amazon DynamoDB table name.
+        dedup_pkeys (list, Optional): List of keys to be used for de-duplicating records in buffer.
+    """
+
+    def __init__(self, table_name: str, dedup_pkeys: list = None):
+        """Constructor of the transform that puts records into an Amazon DynamoDB table.
+
+        Args:
+            table_name (str): Amazon DynamoDB table name.
+            dedup_pkeys (list, Optional): List of keys to be used for de-duplicating items in buffer.
+        """
+        super().__init__()
+        self.table_name = table_name
+        self.dedup_pkeys = dedup_pkeys
+
+    def expand(self, pcoll: PCollection):
+        options = pcoll.pipeline.options.view_as(DynamoDBOptions)
+        return pcoll | beam.ParDo(
+            _DynamoDBWriteFn(self.table_name, self.dedup_pkeys, options)
+        )

src/dynamodb_pyio/io.py

@@ -0,0 +1,68 @@
+#
+# 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 apache_beam.options.pipeline_options import PipelineOptions
+
+
+class DynamoDBOptions(PipelineOptions):
+    @classmethod
+    def _add_argparse_args(cls, parser):
+        # These options are passed to the Firehose PyIO Client
+        parser.add_argument(
+            "--aws_access_key_id",
+            default=None,
+            help="The secret key to use when creating the boto3 client.",
+        )
+        parser.add_argument(
+            "--aws_secret_access_key",
+            default=None,
+            help="The secret key to use when creating the boto3 client.",
+        )
+        parser.add_argument(
+            "--aws_session_token",
+            default=None,
+            help="The session token to use when creating the boto3 client.",
+        )
+        parser.add_argument(
+            "--endpoint_url",
+            default=None,
+            help="The complete URL to use for the constructed boto3 client.",
+        )
+        parser.add_argument(
+            "--region_name",
+            default=None,
+            help="The name of the region associated with the boto3 client.",
+        )
+        parser.add_argument(
+            "--api_version",
+            default=None,
+            help="The API version to use with the boto3 client.",
+        )
+        parser.add_argument(
+            "--verify",
+            default=None,
+            help="Whether or not to verify SSL certificates with the boto3 client.",
+        )
+        parser.add_argument(
+            "--disable_ssl",
+            default=False,
+            action="store_true",
+            help=(
+                "Whether or not to use SSL with the boto3 client. "
+                "By default, SSL is used."
+            ),
+        )