diff --git a/airflow/contrib/hooks/aws_dynamodb_hook.py b/airflow/contrib/hooks/aws_dynamodb_hook.py new file mode 100644 index 0000000000000..d7ae0c7ca747e --- /dev/null +++ b/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) + ) + ) diff --git a/airflow/contrib/hooks/aws_hook.py b/airflow/contrib/hooks/aws_hook.py index 3eced28895a35..61d0eb425e788 100644 --- a/airflow/contrib/hooks/aws_hook.py +++ b/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 + ) diff --git a/airflow/contrib/operators/hive_to_dynamodb.py b/airflow/contrib/operators/hive_to_dynamodb.py new file mode 100644 index 0000000000000..717e3b92b291d --- /dev/null +++ b/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.") diff --git a/tests/contrib/hooks/test_aws_dynamodb_hook.py b/tests/contrib/hooks/test_aws_dynamodb_hook.py new file mode 100644 index 0000000000000..e597e06de8eb1 --- /dev/null +++ b/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() diff --git a/tests/contrib/hooks/test_aws_hook.py b/tests/contrib/hooks/test_aws_hook.py index 6f13e58661458..d460da3c5f34f 100644 --- a/tests/contrib/hooks/test_aws_hook.py +++ b/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() diff --git a/tests/contrib/operators/test_hive_to_dynamodb_operator.py b/tests/contrib/operators/test_hive_to_dynamodb_operator.py new file mode 100644 index 0000000000000..025aeaaabdd4c --- /dev/null +++ b/tests/contrib/operators/test_hive_to_dynamodb_operator.py @@ -0,0 +1,151 @@ +# -*- 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 mock +from airflow import configuration, DAG +configuration.load_test_config() +import datetime +from airflow.contrib.hooks.aws_dynamodb_hook import AwsDynamoDBHook +import airflow.contrib.operators.hive_to_dynamodb + +DEFAULT_DATE = datetime.datetime(2015, 1, 1) +DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat() +DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10] + +try: + from moto import mock_dynamodb2 +except ImportError: + mock_dynamodb2 = None + + +class ProcessData: + def process_data(self, results): + return results + + +class HiveToDynamoDBTransferTest(unittest.TestCase): + + def setUp(self): + configuration.load_test_config() + args = {'owner': 'airflow', 'start_date': DEFAULT_DATE} + dag = DAG('test_dag_id', default_args=args) + self.dag = dag + self.nondefault_schema = "nondefault" + + @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()) + + @mock.patch('airflow.hooks.hive_hooks.HiveServer2Hook.get_results', return_value={'data': [{"id": "1", "name": "siddharth"}]}) + @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present') + @mock_dynamodb2 + def test_get_records_with_schema(self, get_results_mock): + + # Configure + sql = "SELECT 1" + + hook = AwsDynamoDBHook(aws_conn_id='aws_default') + # 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 + } + ) + + operator = airflow.contrib.operators.hive_to_dynamodb.HiveToDynamoDBTransfer( + sql=sql, + table_name="test_airflow", + task_id='hive_to_dynamodb_check', + table_keys=['id'], + dag=self.dag) + + operator.execute(None) + + table = hook.get_conn().Table('test_airflow') + table.meta.client.get_waiter( + 'table_exists').wait(TableName='test_airflow') + self.assertEqual(table.item_count, 1) + + @staticmethod + def pre_process(self, results): + return results + + @mock.patch('airflow.hooks.hive_hooks.HiveServer2Hook.get_results', return_value={'data': [{"id": "1", "name": "siddharth"}, {"id": "1", "name": "gupta"}]}) + @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present') + @mock_dynamodb2 + def test_pre_process_records_with_schema(self, get_results_mock): + + # Configure + sql = "SELECT 1" + + hook = AwsDynamoDBHook(aws_conn_id='aws_default') + # 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 + } + ) + + pre_process = ProcessData() + + operator = airflow.contrib.operators.hive_to_dynamodb.HiveToDynamoDBTransfer( + sql=sql, + table_name="test_airflow", + task_id='hive_to_dynamodb_check', + table_keys=['id'], + pre_process=pre_process, + dag=self.dag) + + operator.execute(None) + + table = hook.get_conn().Table('test_airflow') + table.meta.client.get_waiter( + 'table_exists').wait(TableName='test_airflow') + self.assertEqual(table.item_count, 1) + + +if __name__ == '__main__': + unittest.main()