[AIRFLOW-1560] Add AWS DynamoDB hook and operator for inserting batch items

airflow/contrib/hooks/aws_dynamodb_hook.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed 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 airflow.exceptions import AirflowException
+from airflow.contrib.hooks.aws_hook import AwsHook
+
+
+class AwsDynamoDBHook(AwsHook):
+    """
+    Interact with AWS DynamoDB.
+    """
+
+    def __init__(self, table_keys=None, table_name=None, *args, **kwargs):
+        self.table_keys = table_keys
+        self.table_name = table_name
+        super(AwsDynamoDBHook, self).__init__(*args, **kwargs)
+
+    def get_conn(self):
+        self.conn = self.get_resource_type('dynamodb')
+        return self.conn
+
+    def write_batch_data(self, items):
+        """
+        Write batch items to dynamodb table with provisioned throughout capacity.
+        """
+
+        dynamodb_conn = self.get_conn()
+
+        try:
+            table = dynamodb_conn.Table(self.table_name)
+
+            with table.batch_writer(overwrite_by_pkeys=self.table_keys) as batch:
+                for item in items:
+                    batch.put_item(Item=item)
+            return True
+        except Exception as general_error:
+            raise AirflowException(
+                'Failed to insert items in dynamodb, error: {error}'.format(
+                    error=str(general_error)
+                )
+            )

airflow/contrib/hooks/aws_hook.py

@@ -24,6 +24,7 @@ class AwsHook(BaseHook):
     Interact with AWS.
     This class is a thin wrapper around the boto3 python library.
     """
+
     def __init__(self, aws_conn_id='aws_default'):
         self.aws_conn_id = aws_conn_id
 
@@ -48,3 +49,25 @@ def get_client_type(self, client_type, region_name=None):
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key
         )
+
+    def get_resource_type(self, resource_type, region_name=None):
+        try:
+            connection_object = self.get_connection(self.aws_conn_id)
+            aws_access_key_id = connection_object.login
+            aws_secret_access_key = connection_object.password
+
+            if region_name is None:
+                region_name = connection_object.extra_dejson.get('region_name')
+
+        except AirflowException:
+            # No connection found: fallback on boto3 credential strategy
+            # http://boto3.readthedocs.io/en/latest/guide/configuration.html
+            aws_access_key_id = None
+            aws_secret_access_key = None
+
+        return boto3.resource(
+            resource_type,
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key
+        )

airflow/contrib/operators/hive_to_dynamodb.py

@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed 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
+
+from airflow.hooks.hive_hooks import HiveServer2Hook
+from airflow.contrib.hooks.aws_dynamodb_hook import AwsDynamoDBHook
+
+from airflow.models import BaseOperator
+from airflow.utils.decorators import apply_defaults
+
+
+class HiveToDynamoDBTransfer(BaseOperator):
+    """
+    Moves data from Hive to DynamoDB, note that for now the data is loaded
+    into memory before being pushed to DynamoDB, so this operator should
+    be used for smallish amount of data.
+
+    :param sql: SQL query to execute against the hive database
+    :type sql: str
+    :param table_name: target DynamoDB table
+    :type table_name: str
+    :param table_keys: partition key and sort key
+    :type table_keys: list
+    :param aws_conn_id: aws connection
+    :type aws_conn_id: str
+    :param hiveserver2_conn_id: source hive connection
+    :type hiveserver2_conn_id: str
+    """
+
+    template_fields = ('sql')
+    template_ext = ('.sql',)
+    ui_color = '#a0e08c'
+
+    @apply_defaults
+    def __init__(
+            self,
+            sql,
+            table_name,
+            table_keys,
+            pre_process=None,
+            hiveserver2_conn_id='hiveserver2_default',
+            aws_conn_id='aws_default',
+            *args, **kwargs):
+        super(HiveToDynamoDBTransfer, self).__init__(*args, **kwargs)
+        self.sql = sql
+        self.table_name = table_name
+        self.table_keys = table_keys
+        self.pre_process = pre_process
+        self.hiveserver2_conn_id = hiveserver2_conn_id
+        self.aws_conn_id = aws_conn_id
+
+    def execute(self, context):
+        hive = HiveServer2Hook(hiveserver2_conn_id=self.hiveserver2_conn_id)
+
+        logging.info("Extracting data from Hive")
+        logging.info(self.sql)
+
+        results = hive.get_records(self.sql)
+        dynamodb = AwsDynamoDBHook(aws_conn_id=self.aws_conn_id,
+                                   table_name=self.table_name, table_keys=self.table_keys)
+
+        logging.info("Inserting rows into dynamodb")
+
+        if self.pre_process is None:
+            dynamodb.write_batch_data(results)
+        else:
+            dynamodb.write_batch_data(
+                self.pre_process.process_data(results=results))
+
+        logging.info("Done.")

tests/contrib/hooks/test_aws_dynamodb_hook.py

@@ -0,0 +1,83 @@
+# -*- coding: utf-8 -*-
+#
+# Licensed 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 unittest
+import uuid
+
+from airflow.contrib.hooks.aws_dynamodb_hook import AwsDynamoDBHook
+
+try:
+    from moto import mock_dynamodb2
+except ImportError:
+    mock_dynamodb2 = None
+
+
+class TestDynamoDBHook(unittest.TestCase):
+
+    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
+    @mock_dynamodb2
+    def test_get_conn_returns_a_boto3_connection(self):
+        hook = AwsDynamoDBHook(aws_conn_id='aws_default')
+        self.assertIsNotNone(hook.get_conn())
+
+    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
+    @mock_dynamodb2
+    def test_insert_batch_items_dynamodb_table(self):
+
+        hook = AwsDynamoDBHook(aws_conn_id='aws_default',
+                               table_name="test_airflow", table_keys=['id'])
+
+        # this table needs to be created in production
+        table = hook.get_conn().create_table(
+            TableName="test_airflow",
+            KeySchema=[
+                {
+                    'AttributeName': 'id',
+                    'KeyType': 'HASH'
+                },
+            ],
+            AttributeDefinitions=[
+                {
+                    "AttributeName": "name",
+                    "AttributeType": "S"
+                }
+            ],
+            ProvisionedThroughput={
+                'ReadCapacityUnits': 10,
+                'WriteCapacityUnits': 10
+            }
+        )
+
+        table = hook.get_conn().Table('test_airflow')
+
+        items = []
+
+        for i in range(0, 10):
+            items.append(
+                {
+                    "id": str(uuid.uuid4()),
+                    "name": "airflow"
+                }
+            )
+
+        hook.write_batch_data(items)
+
+        table.meta.client.get_waiter(
+            'table_exists').wait(TableName='test_airflow')
+        self.assertEqual(table.item_count, 10)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/contrib/hooks/test_aws_hook.py

@@ -21,9 +21,10 @@
 
 
 try:
-    from moto import mock_emr
+    from moto import mock_emr, mock_dynamodb2
 except ImportError:
     mock_emr = None
+    mock_dynamodb2 = None
 
 
 class TestAwsHook(unittest.TestCase):
@@ -43,5 +44,39 @@ def test_get_client_type_returns_a_boto3_client_of_the_requested_type(self):
 
         self.assertEqual(client_from_hook.list_clusters()['Clusters'], [])
 
+    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
+    @mock_dynamodb2
+    def test_get_resource_type_returns_a_boto3_resource_of_the_requested_type(self):
+
+        hook = AwsHook(aws_conn_id='aws_default')
+        resource_from_hook = hook.get_resource_type('dynamodb')
+
+        # this table needs to be created in production
+        table = resource_from_hook.create_table(
+            TableName="test_airflow",
+            KeySchema=[
+                {
+                    'AttributeName': 'id',
+                    'KeyType': 'HASH'
+                },
+            ],
+            AttributeDefinitions=[
+                {
+                    "AttributeName": "name",
+                    "AttributeType": "S"
+                }
+            ],
+            ProvisionedThroughput={
+                'ReadCapacityUnits': 10,
+                'WriteCapacityUnits': 10
+            }
+        )
+
+        table.meta.client.get_waiter(
+            'table_exists').wait(TableName='test_airflow')
+
+        self.assertEqual(table.item_count, 0)
+
+
 if __name__ == '__main__':
     unittest.main()