Source code for airflow.providers.amazon.aws.hooks.dms

#
# 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 json
from enum import Enum
from typing import Optional

from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook


[docs]class DmsTaskWaiterStatus(str, Enum): """Available AWS DMS Task Waiter statuses."""
[docs] DELETED = 'deleted'
[docs] READY = 'ready'
[docs] RUNNING = 'running'
[docs] STOPPED = 'stopped'
[docs]class DmsHook(AwsBaseHook): """Interact with AWS Database Migration Service.""" def __init__( self, *args, **kwargs, ): kwargs['client_type'] = 'dms' super().__init__(*args, **kwargs)
[docs] def describe_replication_tasks(self, **kwargs): """ Describe replication tasks :return: Marker and list of replication tasks :rtype: (Optional[str], list) """ dms_client = self.get_conn() response = dms_client.describe_replication_tasks(**kwargs) return response.get('Marker'), response.get('ReplicationTasks', [])
[docs] def find_replication_tasks_by_arn( self, replication_task_arn: str, without_settings: Optional[bool] = False ): """ Find and describe replication tasks by task ARN :param replication_task_arn: Replication task arn :type replication_task_arn: str :param without_settings: Indicates whether to return task information with settings. :type without_settings: Optional[bool] :return: list of replication tasks that match the ARN """ _, tasks = self.describe_replication_tasks( Filters=[ { 'Name': 'replication-task-arn', 'Values': [replication_task_arn], } ], WithoutSettings=without_settings, ) return tasks
[docs] def get_task_status(self, replication_task_arn: str) -> Optional[str]: """ Retrieve task status. :param replication_task_arn: Replication task ARN :type replication_task_arn: str :return: Current task status """ replication_tasks = self.find_replication_tasks_by_arn( replication_task_arn=replication_task_arn, without_settings=True, ) if len(replication_tasks) == 1: status = replication_tasks[0]['Status'] self.log.info('Replication task with ARN(%s) has status "%s".', replication_task_arn, status) return status else: self.log.info('Replication task with ARN(%s) is not found.', replication_task_arn) return None
[docs] def create_replication_task( self, replication_task_id: str, source_endpoint_arn: str, target_endpoint_arn: str, replication_instance_arn: str, migration_type: str, table_mappings: dict, **kwargs, ) -> str: """ Create DMS replication task :param replication_task_id: Replication task id :type replication_task_id: str :param source_endpoint_arn: Source endpoint ARN :type source_endpoint_arn: str :param target_endpoint_arn: Target endpoint ARN :type target_endpoint_arn: str :param replication_instance_arn: Replication instance ARN :type replication_instance_arn: str :param table_mappings: Table mappings :type table_mappings: dict :param migration_type: Migration type ('full-load'|'cdc'|'full-load-and-cdc'), full-load by default. :type migration_type: str :return: Replication task ARN """ dms_client = self.get_conn() create_task_response = dms_client.create_replication_task( ReplicationTaskIdentifier=replication_task_id, SourceEndpointArn=source_endpoint_arn, TargetEndpointArn=target_endpoint_arn, ReplicationInstanceArn=replication_instance_arn, MigrationType=migration_type, TableMappings=json.dumps(table_mappings), **kwargs, ) replication_task_arn = create_task_response['ReplicationTask']['ReplicationTaskArn'] self.wait_for_task_status(replication_task_arn, DmsTaskWaiterStatus.READY) return replication_task_arn
[docs] def start_replication_task( self, replication_task_arn: str, start_replication_task_type: str, **kwargs, ): """ Starts replication task. :param replication_task_arn: Replication task ARN :type replication_task_arn: str :param start_replication_task_type: Replication task start type ('start-replication'|'resume-processing'|'reload-target') :type start_replication_task_type: str """ dms_client = self.get_conn() dms_client.start_replication_task( ReplicationTaskArn=replication_task_arn, StartReplicationTaskType=start_replication_task_type, **kwargs,
)
[docs] def stop_replication_task(self, replication_task_arn): """ Stops replication task. :param replication_task_arn: Replication task ARN :type replication_task_arn: str """ dms_client = self.get_conn() dms_client.stop_replication_task(ReplicationTaskArn=replication_task_arn)
[docs] def delete_replication_task(self, replication_task_arn): """ Starts replication task deletion and waits for it to be deleted :param replication_task_arn: Replication task ARN :type replication_task_arn: str """ dms_client = self.get_conn() dms_client.delete_replication_task(ReplicationTaskArn=replication_task_arn) self.wait_for_task_status(replication_task_arn, DmsTaskWaiterStatus.DELETED)
[docs] def wait_for_task_status(self, replication_task_arn: str, status: DmsTaskWaiterStatus): """ Waits for replication task to reach status. Supported statuses: deleted, ready, running, stopped. :param status: Status to wait for :type status: DmsTaskWaiterStatus :param replication_task_arn: Replication task ARN :type replication_task_arn: str """ if not isinstance(status, DmsTaskWaiterStatus): raise TypeError('Status must be an instance of DmsTaskWaiterStatus') dms_client = self.get_conn() waiter = dms_client.get_waiter(f'replication_task_{status}') waiter.wait( Filters=[ { 'Name': 'replication-task-arn', 'Values': [ replication_task_arn, ], }, ], WithoutSettings=True,
)

Was this entry helpful?