Implement amazon s3 to dynamodb transfer operator (apache#39654)

airflow/providers/amazon/aws/hooks/dynamodb.py

@@ -19,11 +19,17 @@
 
 from __future__ import annotations
 
-from typing import Iterable
+from functools import cached_property
+from typing import TYPE_CHECKING, Iterable
+
+from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
+if TYPE_CHECKING:
+ from botocore.client import BaseClient
+
 
 class DynamoDBHook(AwsBaseHook):
  """
@@ -50,6 +56,11 @@ def __init__(
  kwargs["resource_type"] = "dynamodb"
  super().__init__(*args, **kwargs)
 
+ @cached_property
+ def client(self) -> BaseClient:
+ """Return boto3 client."""
+ return self.get_conn().meta.client
+
  def write_batch_data(self, items: Iterable) -> bool:
  """
  Write batch items to DynamoDB table with provisioned throughout capacity.
@@ -70,3 +81,25 @@ def write_batch_data(self, items: Iterable) -> bool:
  return True
  except Exception as general_error:
  raise AirflowException(f"Failed to insert items in dynamodb, error: {general_error}")
+
+ def get_import_status(self, import_arn: str) -> tuple[str, str | None, str | None]:
+ """
+ Get import status from Dynamodb.
+
+ :param import_arn: The Amazon Resource Name (ARN) for the import.
+ :return: Import status, Error code and Error message
+ """
+ self.log.info("Poking for Dynamodb import %s", import_arn)
+
+ try:
+ describe_import = self.client.describe_import(ImportArn=import_arn)
+ status = describe_import["ImportTableDescription"]["ImportStatus"]
+ error_code = describe_import["ImportTableDescription"].get("FailureCode")
+ error_msg = describe_import["ImportTableDescription"].get("FailureMessage")
+ return status, error_code, error_msg
+ except ClientError as e:
+ error_code = e.response.get("Error", {}).get("Code")
+ if error_code == "ImportNotFoundException":
+ raise AirflowException("S3 import into Dynamodb job not found.")
+ else:
+ raise e

airflow/providers/amazon/aws/transfers/s3_to_dynamodb.py

@@ -0,0 +1,257 @@
+#
+# 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 typing import TYPE_CHECKING, Any, Literal, Sequence, TypedDict
+
+from botocore.exceptions import ClientError, WaiterError
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
+
+if TYPE_CHECKING:
+ from airflow.utils.context import Context
+
+
+class AttributeDefinition(TypedDict):
+ """Attribute Definition Type."""
+
+ AttributeName: str
+ AttributeType: Literal["S", "N", "B"]
+
+
+class KeySchema(TypedDict):
+ """Key Schema Type."""
+
+ AttributeName: str
+ KeyType: Literal["HASH", "RANGE"]
+
+
+class S3ToDynamoDBOperator(BaseOperator):
+ """Load Data from S3 into a DynamoDB.
+
+ Data stored in S3 can be uploaded to a new or existing DynamoDB. Supported file formats CSV, DynamoDB JSON and
+ Amazon ION.
+
+
+ :param s3_bucket: The S3 bucket that is imported
+ :param s3_key: Key prefix that imports single or multiple objects from S3
+ :param dynamodb_table_name: Name of the table that shall be created
+ :param dynamodb_key_schema: Primary key and sort key. Each element represents one primary key
+ attribute. AttributeName is the name of the attribute. KeyType is the role for the attribute. Valid values
+ HASH or RANGE
+ :param dynamodb_attributes: Name of the attributes of a table. AttributeName is the name for the attribute
+ AttributeType is the data type for the attribute. Valid values for AttributeType are
+ S - attribute is of type String
+ N - attribute is of type Number
+ B - attribute is of type Binary
+ :param dynamodb_tmp_table_prefix: Prefix for the temporary DynamoDB table
+ :param delete_on_error: If set, the new DynamoDB table will be deleted in case of import errors
+ :param use_existing_table: Whether to import to an existing non new DynamoDB table. If set to
+ true data is loaded first into a temporary DynamoDB table (using the AWS ImportTable Service),
+ then retrieved as chunks into memory and loaded into the target table. If set to false, a new
+ DynamoDB table will be created and S3 data is bulk loaded by the AWS ImportTable Service.
+ :param input_format: The format for the imported data. Valid values for InputFormat are CSV, DYNAMODB_JSON
+ or ION
+ :param billing_mode: Billing mode for the table. Valid values are PROVISIONED or PAY_PER_REQUEST
+ :param on_demand_throughput: Extra options for maximum number of read and write units
+ :param import_table_kwargs: Any additional optional import table parameters to pass, such as ClientToken,
+ InputCompressionType, or InputFormatOptions. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/client/import_table.html
+ :param import_table_creation_kwargs: Any additional optional import table creation parameters to pass, such as
+ ProvisionedThroughput, SSESpecification, or GlobalSecondaryIndexes. See:
+ https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/client/import_table.html
+ :param wait_for_completion: Whether to wait for cluster to stop
+ :param check_interval: Time in seconds to wait between status checks
+ :param max_attempts: Maximum number of attempts to check for job completion
+ :param aws_conn_id: The reference to the AWS connection details
+ """
+
+ template_fields: Sequence[str] = (
+ "s3_bucket",
+ "s3_key",
+ "dynamodb_table_name",
+ "dynamodb_key_schema",
+ "dynamodb_attributes",
+ "dynamodb_tmp_table_prefix",
+ "delete_on_error",
+ "use_existing_table",
+ "input_format",
+ "billing_mode",
+ "import_table_kwargs",
+ "import_table_creation_kwargs",
+ )
+ ui_color = "#e2e8f0"
+
+ def __init__(
+ self,
+ *,
+ s3_bucket: str,
+ s3_key: str,
+ dynamodb_table_name: str,
+ dynamodb_key_schema: list[KeySchema],
+ dynamodb_attributes: list[AttributeDefinition] | None = None,
+ dynamodb_tmp_table_prefix: str = "tmp",
+ delete_on_error: bool = False,
+ use_existing_table: bool = False,
+ input_format: Literal["CSV", "DYNAMODB_JSON", "ION"] = "DYNAMODB_JSON",
+ billing_mode: Literal["PROVISIONED", "PAY_PER_REQUEST"] = "PAY_PER_REQUEST",
+ import_table_kwargs: dict[str, Any] | None = None,
+ import_table_creation_kwargs: dict[str, Any] | None = None,
+ wait_for_completion: bool = True,
+ check_interval: int = 30,
+ max_attempts: int = 240,
+ aws_conn_id: str | None = "aws_default",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.s3_bucket = s3_bucket
+ self.s3_key = s3_key
+ self.dynamodb_table_name = dynamodb_table_name
+ self.dynamodb_attributes = dynamodb_attributes
+ self.dynamodb_tmp_table_prefix = dynamodb_tmp_table_prefix
+ self.delete_on_error = delete_on_error
+ self.use_existing_table = use_existing_table
+ self.dynamodb_key_schema = dynamodb_key_schema
+ self.input_format = input_format
+ self.billing_mode = billing_mode
+ self.import_table_kwargs = import_table_kwargs
+ self.import_table_creation_kwargs = import_table_creation_kwargs
+ self.wait_for_completion = wait_for_completion
+ self.check_interval = check_interval
+ self.max_attempts = max_attempts
+ self.aws_conn_id = aws_conn_id
+
+ @property
+ def tmp_table_name(self):
+ """Temporary table name."""
+ return f"{self.dynamodb_tmp_table_prefix}_{self.dynamodb_table_name}"
+
+ def _load_into_new_table(self, table_name: str, delete_on_error: bool) -> str:
+ """
+ Import S3 key or keys into a new DynamoDB table.
+
+ :param table_name: Name of the table that shall be created
+ :param delete_on_error: If set, the new DynamoDB table will be deleted in case of import errors
+ :return: The Amazon resource number (ARN)
+ """
+ dynamodb_hook = DynamoDBHook(aws_conn_id=self.aws_conn_id)
+ client = dynamodb_hook.client
+
+ import_table_config = self.import_table_kwargs or {}
+ import_table_creation_config = self.import_table_creation_kwargs or {}
+
+ try:
+ response = client.import_table(
+ S3BucketSource={
+ "S3Bucket": self.s3_bucket,
+ "S3KeyPrefix": self.s3_key,
+ },
+ InputFormat=self.input_format,
+ TableCreationParameters={
+ "TableName": table_name,
+ "AttributeDefinitions": self.dynamodb_attributes,
+ "KeySchema": self.dynamodb_key_schema,
+ "BillingMode": self.billing_mode,
+ **import_table_creation_config,
+ },
+ **import_table_config,
+ )
+ except ClientError as e:
+ self.log.error("Error: failed to load from S3 into DynamoDB table. Error: %s", str(e))
+ raise AirflowException(f"S3 load into DynamoDB table failed with error: {e}")
+
+ if response["ImportTableDescription"]["ImportStatus"] == "FAILED":
+ raise AirflowException(
+ "S3 into Dynamodb job creation failed. Code: "
+ f"{response['ImportTableDescription']['FailureCode']}. "
+ f"Failure: {response['ImportTableDescription']['FailureMessage']}"
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for S3 into Dynamodb job to complete")
+ waiter = dynamodb_hook.get_waiter("import_table")
+ try:
+ waiter.wait(
+ ImportArn=response["ImportTableDescription"]["ImportArn"],
+ WaiterConfig={"Delay": self.check_interval, "MaxAttempts": self.max_attempts},
+ )
+ except WaiterError:
+ status, error_code, error_msg = dynamodb_hook.get_import_status(
+ response["ImportTableDescription"]["ImportArn"]
+ )
+ if delete_on_error:
+ client.delete_table(TableName=table_name)
+ raise AirflowException(
+ f"S3 import into Dynamodb job failed: Status: {status}. Error: {error_code}. Error message: {error_msg}"
+ )
+ return response["ImportTableDescription"]["ImportArn"]
+
+ def _load_into_existing_table(self) -> str:
+ """
+ Import S3 key or keys in an existing DynamoDB table.
+
+ :return:The Amazon resource number (ARN)
+ """
+ if not self.wait_for_completion:
+ raise ValueError("wait_for_completion must be set to True when loading into an existing table")
+ table_keys = [key["AttributeName"] for key in self.dynamodb_key_schema]
+
+ dynamodb_hook = DynamoDBHook(
+ aws_conn_id=self.aws_conn_id, table_name=self.dynamodb_table_name, table_keys=table_keys
+ )
+ client = dynamodb_hook.client
+
+ self.log.info("Loading from S3 into a tmp DynamoDB table %s", self.tmp_table_name)
+ self._load_into_new_table(table_name=self.tmp_table_name, delete_on_error=self.delete_on_error)
+ total_items = 0
+ try:
+ paginator = client.get_paginator("scan")
+ paginate = paginator.paginate(
+ TableName=self.tmp_table_name,
+ Select="ALL_ATTRIBUTES",
+ ReturnConsumedCapacity="NONE",
+ ConsistentRead=True,
+ )
+ self.log.info(
+ "Loading data from %s to %s DynamoDB table", self.tmp_table_name, self.dynamodb_table_name
+ )
+ for page in paginate:
+ total_items += page.get("Count", 0)
+ dynamodb_hook.write_batch_data(items=page["Items"])
+ self.log.info("Number of items loaded: %s", total_items)
+ finally:
+ self.log.info("Delete tmp DynamoDB table %s", self.tmp_table_name)
+ client.delete_table(TableName=self.tmp_table_name)
+ return dynamodb_hook.get_conn().Table(self.dynamodb_table_name).table_arn
+
+ def execute(self, context: Context) -> str:
+ """
+ Execute S3 to DynamoDB Job from Airflow.
+
+ :param context: The current context of the task instance
+ :return: The Amazon resource number (ARN)
+ """
+ if self.use_existing_table:
+ self.log.info("Loading from S3 into new DynamoDB table %s", self.dynamodb_table_name)
+ return self._load_into_existing_table()
+ self.log.info("Loading from S3 into existing DynamoDB table %s", self.dynamodb_table_name)
+ return self._load_into_new_table(
+ table_name=self.dynamodb_table_name, delete_on_error=self.delete_on_error
+ )

airflow/providers/amazon/aws/waiters/dynamodb.json

@@ -25,6 +25,43 @@
  "state": "retry"
  }
  ]
+ },
+ "import_table": {
+ "operation": "DescribeImport",
+ "delay": 30,
+ "maxAttempts": 240,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "expected": "COMPLETED",
+ "argument": "ImportTableDescription.ImportStatus",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "expected": "CANCELLING",
+ "argument": "ImportTableDescription.ImportStatus",
+ "state": "failure"
+ },
+ {
+ "matcher": "path",
+ "expected": "CANCELLED",
+ "argument": "ImportTableDescription.ImportStatus",
+ "state": "failure"
+ },
+ {
+ "matcher": "path",
+ "expected": "FAILED",
+ "argument": "ImportTableDescription.ImportStatus",
+ "state": "failure"
+ },
+ {
+ "matcher": "path",
+ "expected": "IN_PROGRESS",
+ "argument": "ImportTableDescription.ImportStatus",
+ "state": "retry"
+ }
+ ]
  }
  }
 }

airflow/providers/amazon/provider.yaml

@@ -822,6 +822,10 @@ transfers:
  target-integration-name: Amazon Simple Storage Service (S3)
  how-to-guide: /docs/apache-airflow-providers-amazon/transfer/azure_blob_to_s3.rst
  python-module: airflow.providers.amazon.aws.transfers.azure_blob_to_s3
+ - source-integration-name: Amazon Simple Storage Service (S3)
+ target-integration-name: Amazon DynamoDB
+ how-to-guide: /docs/apache-airflow-providers-amazon/transfer/s3_to_dynamodb.rst
+ python-module: airflow.providers.amazon.aws.transfers.s3_to_dynamodb
 
 extra-links:
  - airflow.providers.amazon.aws.links.athena.AthenaQueryResultsLink