You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2018/12/18 04:22:42 UTC
[GitHub] stale[bot] closed pull request #3250: [WIP] Add dms raw
stale[bot] closed pull request #3250: [WIP] Add dms raw
URL: https://github.com/apache/incubator-airflow/pull/3250
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/airflow/contrib/hooks/dms_hook.py b/airflow/contrib/hooks/dms_hook.py
new file mode 100644
index 0000000000..971c38e3bd
--- /dev/null
+++ b/airflow/contrib/hooks/dms_hook.py
@@ -0,0 +1,90 @@
+# -*- 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 DMSHook(AwsHook):
+ """
+ Interact with AWS Database Migration Service.
+
+ :param dms_conn_id: The connection to AWS DMS, only necessary for
+ the create_replication_task method
+ :param dms_conn_id: string
+ """
+
+ def __init__(self, dms_conn_id=None, *args, **kwargs):
+ self.dms_conn_id = dms_conn_id
+ super(DMSHook, self).__init__(*args, **kwargs)
+
+ def get_conn(self):
+ self.conn = self.get_client_type('dms')
+ return self.conn
+
+ def create_replication_task(self, replication_task_overrides):
+ """
+ Creates a replication task using the config from the DMS connection.
+ Keys of the json extra hash may have the arguments of the boto3
+ create_replication_task method. Overrides for this config may be
+ passed as the replication_task_overrides.
+ """
+
+ if not self.dms_conn_id:
+ raise AirflowException('dms_conn_id must be present '
+ 'to use create_replication_task')
+
+ dms_conn = self.get_connection(self.dms_conn_id)
+
+ config = dms_conn.extra_dejson.copy()
+ config.update(replication_task_overrides)
+
+ return self.get_conn().create_replication_task(
+ ReplicationTaskIdentifier=config.get('replication_task_identifier'),
+ SourceEndpointArn=config.get('source_endpoint_arn'),
+ TargetEndpointArn=config.get('target_endpoint_arn'),
+ ReplicationInstanceArn=config.get('replication_instance_arn'),
+ MigrationType=config.get('migration_type'),
+ TableMappings=config.get('table_mappings'),
+ ReplicationTaskSettings=config.get('replication_task_settings')
+ )
+
+ def stop_replication_task(self, replication_task_arn):
+ return self.get_conn().stop_replication_task(
+ ReplicationTaskArn=replication_task_arn
+ )
+
+ def delete_replication_task(self, replication_task_arn):
+ return self.get_conn().delete_replication_task(
+ ReplicationTaskArn=replication_task_arn
+ )
+
+ def start_replication_task(self, replication_task_arn, start_replication_task_type, cdc_start_time=None):
+ return self.get_conn().start_replication_task(
+ ReplicationTaskArn=replication_task_arn,
+ StartReplicationTaskType=start_replication_task_type
+ # ,CdcStartTime=self.cdc_start_time,
+ )
+
+ def describe_replication_task(self, replication_task_arn):
+ return self.get_conn().describe_replication_tasks(
+ Filters=[
+ {
+ 'Name': 'replication-task-arn',
+ 'Values': [
+ replication_task_arn,
+ ]
+ },
+ ],
+ )
diff --git a/airflow/contrib/operators/dms_create_replication_task_operator.py b/airflow/contrib/operators/dms_create_replication_task_operator.py
new file mode 100644
index 0000000000..c23b1d126a
--- /dev/null
+++ b/airflow/contrib/operators/dms_create_replication_task_operator.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.
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.utils import apply_defaults
+
+from airflow.contrib.hooks.dms_hook import DMSHook
+
+class DMSCreateReplicationTaskOperator(BaseOperator):
+ """
+ Execute a task on AWS Database Migration Service
+
+ :param task_definition: the task definition name on EC2 Container Service
+ :type task_definition: str
+ :param cluster: the cluster name on EC2 Container Service
+ :type cluster: str
+ :param: overrides: the same parameter that boto3 will receive:
+ http://boto3.readthedocs.org/en/latest/reference/services/ecs.html#ECS.Client.run_task
+ :type: overrides: dict
+ :param aws_conn_id: connection id of AWS credentials / region name. If None,
+ credential boto3 strategy will be used (http://boto3.readthedocs.io/en/latest/guide/configuration.html).
+ :type aws_conn_id: str
+ :param region_name: region name to use in AWS Hook. Override the region_name in connection (if provided)
+ """
+
+ ui_color = '#f0ede4'
+ # template_fields = ['replication_task_arn']
+ template_fields = []
+ template_ext = ()
+
+ @apply_defaults
+ def __init__(self, dms_conn_id='dms_default', aws_conn_id=None,
+ replication_task_overrides=None, region_name=None, **kwargs):
+ super(DMSCreateReplicationTaskOperator, self).__init__(**kwargs)
+
+ self.aws_conn_id = aws_conn_id
+ # self.region_name = region_name
+ self.dms_conn_id = dms_conn_id
+ if replication_task_overrides is None:
+ replication_task_overrides = {}
+ self.replication_task_overrides = replication_task_overrides
+
+ self.hook = self.get_hook()
+
+ def execute(self, context):
+ self.log.info(
+ 'Creating DMS task using aws-conn-id: %s, dms-conn-id: %s',
+ self.aws_conn_id, self.dms_conn_id
+ )
+ response = self.hook.create_replication_task(
+ self.replication_task_overrides
+ )
+
+ if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
+ raise AirflowException('Creating DMS task failed: %s' % response)
+ else:
+ task_arn = response['ReplicationTask']['ReplicationTaskArn']
+ self.log.info('Running Create DMS Task with task arn: %s', task_arn)
+ return task_arn
+
+ def get_hook(self):
+ return DMSHook(
+ dms_conn_id=self.dms_conn_id,
+ aws_conn_id=self.aws_conn_id
+ )
+
+ # def on_kill(self):
+ # response = self.get_hook().stop_replication_task(
+ # ReplicationTaskArn=self.replication_task_arn
+ # )
+ # self.log.info('Create DMS on_kill: %s', response)
diff --git a/airflow/contrib/operators/dms_delete_replication_task_operator.py b/airflow/contrib/operators/dms_delete_replication_task_operator.py
new file mode 100644
index 0000000000..fc1bd25d30
--- /dev/null
+++ b/airflow/contrib/operators/dms_delete_replication_task_operator.py
@@ -0,0 +1,80 @@
+# -*- 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.models import BaseOperator
+from airflow.utils import apply_defaults
+
+from airflow.contrib.hooks.dms_hook import DMSHook
+
+class DMSDeleteReplicationTaskOperator(BaseOperator):
+ """
+ Execute a task on AWS Database Migration Service
+
+ :param task_definition: the task definition name on EC2 Container Service
+ :type task_definition: str
+ :param cluster: the cluster name on EC2 Container Service
+ :type cluster: str
+ :param: overrides: the same parameter that boto3 will receive:
+ http://boto3.readthedocs.org/en/latest/reference/services/ecs.html#ECS.Client.run_task
+ :type: overrides: dict
+ :param aws_conn_id: connection id of AWS credentials / region name. If None,
+ credential boto3 strategy will be used (http://boto3.readthedocs.io/en/latest/guide/configuration.html).
+ :type aws_conn_id: str
+ :param region_name: region name to use in AWS Hook. Override the region_name in connection (if provided)
+ """
+
+ ui_color = '#f0ede4'
+ template_fields = ['replication_task_arn']
+ template_ext = ()
+
+ @apply_defaults
+ def __init__(self, replication_task_arn, dms_conn_id='dms_default',
+ aws_conn_id=None, **kwargs):
+ super(DMSDeleteReplicationTaskOperator, self).__init__(**kwargs)
+
+ self.aws_conn_id = aws_conn_id
+ # self.region_name = region_name
+ self.dms_conn_id = dms_conn_id
+ self.replication_task_arn = replication_task_arn
+
+ self.hook = self.get_hook()
+
+ def execute(self, context):
+ self.log.info(
+ 'Deleting DMS task using aws-conn-id: %s, dms-conn-id: %s',
+ self.aws_conn_id, self.dms_conn_id
+ )
+ response = self.hook.delete_replication_task(
+ ReplicationTaskArn=self.replication_task_arn
+ )
+
+ if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
+ raise AirflowException('Deleting DMS task failed: %s' % response)
+ else:
+ task_arn = response['ReplicationTask']['ReplicationTaskArn']
+ self.log.info('Running Delete DMS Task with arn: %s', task_arn)
+ return task_arn
+
+ def get_hook(self):
+ return DMSHook(
+ dms_conn_id=self.dms_conn_id,
+ aws_conn_id=self.aws_conn_id
+ )
+
+ def on_kill(self):
+ response = self.hook.stop_replication_task(
+ ReplicationTaskArn=self.replication_task_arn
+ )
+ self.log.info('Stop DMS task on_kill: %s', response)
diff --git a/airflow/contrib/operators/dms_start_replication_task_operator.py b/airflow/contrib/operators/dms_start_replication_task_operator.py
new file mode 100644
index 0000000000..442e270e67
--- /dev/null
+++ b/airflow/contrib/operators/dms_start_replication_task_operator.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.
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.utils import apply_defaults
+
+from airflow.contrib.hooks.dms_hook import DMSHook
+
+class DMSStartReplicationTaskOperator(BaseOperator):
+ """
+ Execute a task on AWS Database Migration Service
+
+ :param task_definition: the task definition name on EC2 Container Service
+ :type task_definition: str
+ :param cluster: the cluster name on EC2 Container Service
+ :type cluster: str
+ :param: overrides: the same parameter that boto3 will receive:
+ http://boto3.readthedocs.org/en/latest/reference/services/ecs.html#ECS.Client.run_task
+ :type: overrides: dict
+ :param aws_conn_id: connection id of AWS credentials / region name. If None,
+ credential boto3 strategy will be used (http://boto3.readthedocs.io/en/latest/guide/configuration.html).
+ :type aws_conn_id: str
+ :param region_name: region name to use in AWS Hook. Override the region_name in connection (if provided)
+ """
+
+ ui_color = '#f0ede4'
+ template_fields = ['replication_task_arn']
+ template_ext = ()
+
+ @apply_defaults
+ def __init__(self, replication_task_arn,
+ start_replication_task_type='start-replication',
+ dms_conn_id='dms_default', aws_conn_id=None, **kwargs):
+ super(DMSStartReplicationTaskOperator, self).__init__(**kwargs)
+ self.aws_conn_id = aws_conn_id
+ # self.region_name = region_name
+ self.start_replication_task_type = start_replication_task_type
+ self.replication_task_arn = replication_task_arn
+ self.dms_conn_id = dms_conn_id
+
+ self.hook = self.get_hook()
+
+ def execute(self, context):
+ self.log.info(
+ 'Starting DMS task using aws-conn-id: %s, dms-conn-id: %s',
+ self.aws_conn_id, self.dms_conn_id
+ )
+ response = self.hook.start_replication_task(
+ ReplicationTaskArn=self.replication_task_arn,
+ StartReplicationTaskType=self.start_replication_task_type
+ # ,CdcStartTime=self.cdc_start_time,
+ )
+
+ if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
+ raise AirflowException('Adding steps failed: %s' % response)
+ else:
+ task_arn = response['ReplicationTask']['ReplicationTaskArn']
+ self.log.info('Running Delete DMS Task with arn: %s', task_arn)
+ return task_arn
+
+ def get_hook(self):
+ return DMSHook(
+ dms_conn_id=self.dms_conn_id,
+ aws_conn_id=self.aws_conn_id
+ )
+
+ def on_kill(self):
+ response = self.hook.stop_replication_task(
+ ReplicationTaskArn=self.replication_task_arn
+ )
+ self.log.info('Stop DMS task on_kill: %s', response)
diff --git a/airflow/contrib/sensors/dms_base_sensor.py b/airflow/contrib/sensors/dms_base_sensor.py
new file mode 100644
index 0000000000..9f2f5862dc
--- /dev/null
+++ b/airflow/contrib/sensors/dms_base_sensor.py
@@ -0,0 +1,93 @@
+# -*- 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.operators.sensors import BaseSensorOperator
+from airflow.contrib.hooks.dms_hook import DMSHook
+from airflow.utils import apply_defaults
+from airflow.exceptions import AirflowException
+
+
+class DMSBaseSensor(BaseSensorOperator):
+ """
+ Contains general sensor behavior for EMR.
+ Subclasses should implement get_emr_response() and state_from_response() methods.
+ Subclasses should also implement NON_TERMINAL_STATES and FAILED_STATE constants.
+ """
+ ui_color = '#66c3ff'
+ template_fields = ['replication_task_arn']
+ template_ext = ()
+
+ @apply_defaults
+ def __init__(
+ self,
+ replication_task_arn,
+ aws_conn_id='aws_default',
+ *args, **kwargs):
+ super(DMSBaseSensor, self).__init__(*args, **kwargs)
+ self.aws_conn_id = aws_conn_id
+ self.replication_task_arn = replication_task_arn
+
+ def poke(self, context):
+ response = self.get_dms_response()
+
+ if not response['ResponseMetadata']['HTTPStatusCode'] == 200:
+ self.log.info('Bad HTTP response: %s', response)
+ return False
+
+ status = self.status_from_response(response)
+ self.log.info('DMS Replication task (%s) currently %s', self.replication_task_arn, status)
+
+ if status in self.NON_TERMINAL_STATUSES:
+ return False
+
+ if status in self.FAILED_STATUSES:
+ raise AirflowException('DMS replication task failed')
+
+ if status in self.TARGET_STATUSES:
+ return True
+
+ self.log.info('DMS Replication task (%s) will ignore unexpected '
+ 'status: %s', self.replication_task_arn, status)
+ return False
+
+ def get_dms_response(self):
+ self.log.info('Poking replication task %s', self.replication_task_arn)
+ return self.get_client().describe_replication_tasks(
+ Filters=[
+ {
+ 'Name': 'replication-task-arn',
+ 'Values': [
+ self.replication_task_arn,
+ ]
+ },
+ ],
+ )
+
+ def stop_reason_handling(self, stop_reason):
+ pass
+
+ def status_from_response(self, response):
+ replication_task = response['ReplicationTasks'][0]
+ if 'StopReason' in replication_task:
+ stop_reason = replication_task['StopReason']
+ self.stop_reason_handling(stop_reason)
+ self.log.info(
+ 'Replication task: %s stopped with '
+ 'StopReason: %s', self.replication_task_arn, stop_reason
+ )
+ return replication_task['Status']
+
+ def get_client(self):
+ return DMSHook(
+ aws_conn_id=self.aws_conn_id
+ )
diff --git a/airflow/contrib/sensors/dms_replication_task_ready_sensor.py b/airflow/contrib/sensors/dms_replication_task_ready_sensor.py
new file mode 100644
index 0000000000..535f80b8ef
--- /dev/null
+++ b/airflow/contrib/sensors/dms_replication_task_ready_sensor.py
@@ -0,0 +1,38 @@
+# -*- 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.contrib.sensors.dms_base_sensor import DMSBaseSensor
+from airflow.utils import apply_defaults
+
+
+class DMSReplicationTaskReadySensor(DMSBaseSensor):
+ """
+ Asks for the state of the step until it reaches a terminal state.
+ If it fails the sensor errors, failing the task.
+
+ :param job_flow_id: job_flow_idwhich contains the step check the state of
+ :type job_flow_id: string
+ :param step_id: step to check the state of
+ :type step_id: string
+ """
+
+ NON_TERMINAL_STATUSES = ['creating', 'starting', 'running']
+ FAILED_STATUSES = ['failed', 'stopped']
+ TARGET_STATUSES = ['ready']
+
+ @apply_defaults
+ def __init__(
+ self,
+ *args, **kwargs):
+ super(DMSReplicationTaskReadySensor, self).__init__(*args, **kwargs)
diff --git a/airflow/contrib/sensors/dms_replication_task_stopped_sensor.py b/airflow/contrib/sensors/dms_replication_task_stopped_sensor.py
new file mode 100644
index 0000000000..6ecf5353cb
--- /dev/null
+++ b/airflow/contrib/sensors/dms_replication_task_stopped_sensor.py
@@ -0,0 +1,50 @@
+# -*- 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.sensors.dms_base_sensor import DMSBaseSensor
+from airflow.utils import apply_defaults
+
+
+class DMSReplicationTaskStoppedSensor(DMSBaseSensor):
+ """
+ Asks for the state of the step until it reaches a terminal state.
+ If it fails the sensor errors, failing the task.
+
+ :param job_flow_id: job_flow_id which contains the step check the state of
+ :type job_flow_id: string
+ :param step_id: step to check the state of
+ :type step_id: string
+ """
+
+ NON_TERMINAL_STATUSES = ['creating', 'running', 'ready', 'starting']
+ FAILED_STATUSES = ['failed']
+ TARGET_STATUSES = ['stopped']
+ STOP_REASONS = ['Stop Reason FULL_LOAD_ONLY_FINISHED', 'Stop Reason NORMAL']
+
+ @apply_defaults
+ def __init__(
+ self,
+ stop_reasons=STOP_REASONS,
+ *args, **kwargs):
+ super(DMSReplicationTaskStoppedSensor, self).__init__(*args, **kwargs)
+ self.stop_reasons = stop_reasons
+
+ def stop_reason_handling(self, stop_reason):
+ if stop_reason not in self.stop_reasons:
+ raise AirflowException(
+ 'DMS replication task: %s failed to stop with acceptable '
+ 'reasons %s but instead %s',
+ self.replication_task_arn, self.stop_reasons, stop_reason
+ )
diff --git a/airflow/models.py b/airflow/models.py
index cef6efcc14..77ec368627 100755
--- a/airflow/models.py
+++ b/airflow/models.py
@@ -619,6 +619,7 @@ class Connection(Base, LoggingMixin):
('segment', 'Segment',),
('azure_data_lake', 'Azure Data Lake'),
('cassandra', 'Cassandra',),
+ ('dms', 'Database Migration Service',),
]
def __init__(
diff --git a/airflow/utils/db.py b/airflow/utils/db.py
index b5e0c49c61..a9eb8ddacf 100644
--- a/airflow/utils/db.py
+++ b/airflow/utils/db.py
@@ -264,6 +264,19 @@ def initdb(rbac=False):
]
}
'''))
+ merge_conn(
+ models.Connection(
+ conn_id='dms_default', conn_type='dms',
+ extra='''
+ { "ReplicationTaskIdentifier": "default_replication_task_name",
+ "SourceEndpointArn": "source-arn",
+ "TargetEndpointArn": "endpoint-arn",
+ "ReplicationInstanceArn": "instance-arn",
+ "MigrationType": "full-load",
+ "TableMappings": "my-table-mappings",
+ "ReplicationTaskSettings": "my-replication-task_settings"
+ }
+ '''))
merge_conn(
models.Connection(
conn_id='databricks_default', conn_type='databricks',
diff --git a/docs/code.rst b/docs/code.rst
index 3b5148470e..a07fedbf7e 100644
--- a/docs/code.rst
+++ b/docs/code.rst
@@ -141,6 +141,9 @@ Operators
.. autoclass:: airflow.contrib.operators.datastore_export_operator.DatastoreExportOperator
.. autoclass:: airflow.contrib.operators.datastore_import_operator.DatastoreImportOperator
.. autoclass:: airflow.contrib.operators.discord_webhook_operator.DiscordWebhookOperator
+.. autoclass:: airflow.contrib.operators.dms_create_replication_task_operator.DMSCreateReplicationTaskOperator
+.. autoclass:: airflow.contrib.operators.dms_delete_replication_task_operator.DMSDeleteReplicationTaskOperator
+.. autoclass:: airflow.contrib.operators.dms_start_replication_task_operator.DMSStartReplicationTaskOperator
.. autoclass:: airflow.contrib.operators.druid_operator.DruidOperator
.. autoclass:: airflow.contrib.operators.ecs_operator.ECSOperator
.. autoclass:: airflow.contrib.operators.emr_add_steps_operator.EmrAddStepsOperator
@@ -199,6 +202,9 @@ Sensors
.. autoclass:: airflow.contrib.sensors.bash_sensor.BashSensor
.. autoclass:: airflow.contrib.sensors.bigquery_sensor.BigQueryTableSensor
.. autoclass:: airflow.contrib.sensors.datadog_sensor.DatadogSensor
+.. autoclass:: airflow.contrib.sensors.dms_base_sensor.DMSBaseSensor
+.. autoclass:: airflow.contrib.sensors.dms_replication_task_ready_sensor.DMSReplicationTaskReadySensor
+.. autoclass:: airflow.contrib.sensors.dms_replication_task_stopped_sensor.DMSReplicationTaskStoppedSensor
.. autoclass:: airflow.contrib.sensors.emr_base_sensor.EmrBaseSensor
.. autoclass:: airflow.contrib.sensors.emr_job_flow_sensor.EmrJobFlowSensor
.. autoclass:: airflow.contrib.sensors.emr_step_sensor.EmrStepSensor
@@ -372,6 +378,7 @@ Community contributed hooks
.. autoclass:: airflow.contrib.hooks.datadog_hook.DatadogHook
.. autoclass:: airflow.contrib.hooks.datastore_hook.DatastoreHook
.. autoclass:: airflow.contrib.hooks.discord_webhook_hook.DiscordWebhookHook
+.. autoclass:: airflow.contrib.hooks.dms_hook.DMSHook
.. autoclass:: airflow.contrib.hooks.emr_hook.EmrHook
.. autoclass:: airflow.contrib.hooks.fs_hook.FSHook
.. autoclass:: airflow.contrib.hooks.ftp_hook.FTPHook
diff --git a/tests/contrib/hooks/test_dms_hook.py b/tests/contrib/hooks/test_dms_hook.py
new file mode 100644
index 0000000000..313458e589
--- /dev/null
+++ b/tests/contrib/hooks/test_dms_hook.py
@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+#
+# 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 unittest
+import boto3
+
+from airflow import configuration
+from airflow.contrib.hooks.dms_hook import DMSHook
+
+
+try:
+ from moto import mock_dms
+except ImportError:
+ mock_dms = None
+
+
+class TestDMSHook(unittest.TestCase):
+ @mock_dms
+ def setUp(self):
+ configuration.load_test_config()
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_get_conn_returns_a_boto3_connection(self):
+ hook = DMSHook(aws_conn_id='dms_default')
+ self.assertIsNotNone(hook.get_conn().describe_replication_tasks())
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_create_replication_task_uses_the_dms_config_to_create_a_task(self):
+ client = boto3.client('dms', region_name='us-east-1')
+ if len(client.describe_replication_tasks()['Clusters']):
+ raise ValueError('AWS not properly mocked')
+
+ hook = DMSHook(aws_conn_id='aws_default', dms_conn_id='dms_default')
+ task = hook.create_replication_task({'Name': 'test_cluster'})
+
+ self.assertEqual(client.describe_replication_tasks()['Clusters'][0]['Id'], task['JobFlowId'])
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_stop_replication_task(self):
+ hook = DMSHook(aws_conn_id='aws_default', dms_conn_id='dms_default')
+ replication_task_arn = 'test'
+ task = hook.stop_replication_task(replication_task_arn)
+
+ # self.assertEqual(client.describe_replication_tasks()['Clusters'][0]['Id'], task['JobFlowId'])
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_start_replication_task(self):
+ hook = DMSHook(aws_conn_id='aws_default', dms_conn_id='dms_default')
+ replication_task_arn = 'test'
+ start_replication_task_type = 'full_load'
+ task = hook.start_replication_task(replication_task_arn, start_replication_task_type)
+
+ # self.assertEqual(client.describe_replication_tasks()['Clusters'][0]['Id'], task['JobFlowId'])
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_delete_replication_task(self):
+ hook = DMSHook(aws_conn_id='aws_default', dms_conn_id='dms_default')
+ replication_task_arn = 'test'
+ task = hook.delete_replication_task(replication_task_arn)
+
+ # self.assertEqual(client.describe_replication_tasks()['Clusters'][0]['Id'], task['JobFlowId'])
+
+ @unittest.skipIf(mock_dms is None, 'mock_dms package not present')
+ @mock_dms
+ def test_describe_replication_task(self):
+ hook = DMSHook(aws_conn_id='aws_default', dms_conn_id='dms_default')
+ replication_task_arn = 'test'
+ task = hook.describe_replication_task(replication_task_arn)
+
+ # self.assertEqual(client.describe_replication_tasks()['Clusters'][0]['Id'], task['JobFlowId'])
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/tests/core.py b/tests/core.py
index c3adb3c500..b43621a903 100644
--- a/tests/core.py
+++ b/tests/core.py
@@ -7,7 +7,7 @@
# 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,
@@ -1054,6 +1054,7 @@ def test_cli_connections_list(self):
self.assertIn(['aws_default', 'aws'], conns)
self.assertIn(['beeline_default', 'beeline'], conns)
self.assertIn(['emr_default', 'emr'], conns)
+
self.assertIn(['mssql_default', 'mssql'], conns)
self.assertIn(['mysql_default', 'mysql'], conns)
self.assertIn(['postgres_default', 'postgres'], conns)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
With regards,
Apache Git Services