o-nikolas commented on code in PR #43988:
URL: https://github.com/apache/airflow/pull/43988#discussion_r1842705282
##########
providers/src/airflow/providers/amazon/aws/operators/dms.py:
##########
@@ -277,3 +288,551 @@ def execute(self, context: Context):
"""Stop AWS DMS replication task from Airflow."""
self.hook.stop_replication_task(replication_task_arn=self.replication_task_arn)
self.log.info("DMS replication task(%s) is stopping.",
self.replication_task_arn)
+
+
+class DmsDescribeReplicationConfigsOperator(AwsBaseOperator[DmsHook]):
+ """
+ Describes AWS DMS Serverless replication configurations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DmsDescribeReplicationConfigsOperator`
+
+ :param describe_config_filter: Filters block for filtering results.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used.
If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ """
+
+ aws_hook_class = DmsHook
+ template_fields: Sequence[str] = aws_template_fields("filter")
+ template_fields_renderers = {"filter": "json"}
+
+ def __init__(
+ self,
+ *,
+ filter: list[dict] | None = None,
+ aws_conn_id: str | None = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+ self.filter = filter
+
+ def execute(self, context: Context) -> list:
+ """
+ Describe AWS DMS replication configurations.
+
+ :return: List of replication configurations
+ """
+ return self.hook.describe_replication_configs(filters=self.filter)
+
+
+class DmsCreateReplicationConfigOperator(AwsBaseOperator[DmsHook]):
+ """
+
+ Creates an AWS DMS Serverless replication configuration.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DmsCreateReplicationConfigOperator`
+
+ :param replication_config_id: Unique identifier used to create a
ReplicationConfigArn.
+ :param source_endpoint_arn: ARN of the source endpoint
+ :param target_endpoint_arn: ARN of the target endpoint
+ :param compute_config: Parameters for provisioning an DMS Serverless
replication.
+ :param replication_type: type of DMS Serverless replication
+ :param table_mappings: JSON table mappings
+ :param tags: Key-value tag pairs
+ :param additional_config_kwargs: Additional configuration parameters for
DMS Serverless replication. Passed directly to the API
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used.
If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ """
+
+ aws_hook_class = DmsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "replication_config_id",
+ "source_endpoint_arn",
+ "target_endpoint_arn",
+ "compute_config",
+ "replication_type",
+ "table_mappings",
+ )
+
+ template_fields_renderers = {"compute_config": "json", "tableMappings":
"json"}
+
+ def __init__(
+ self,
+ *,
+ replication_config_id: str,
+ source_endpoint_arn: str,
+ target_endpoint_arn: str,
+ compute_config: dict[str, Any],
+ replication_type: str,
+ table_mappings: str,
+ additional_config_kwargs: dict | None = None,
+ aws_conn_id: str | None = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(
+ aws_conn_id=aws_conn_id,
+ **kwargs,
+ )
+
+ self.replication_config_id = replication_config_id
+ self.source_endpoint_arn = source_endpoint_arn
+ self.target_endpoint_arn = target_endpoint_arn
+ self.compute_config = compute_config
+ self.replication_type = replication_type
+ self.table_mappings = table_mappings
+ self.additional_config_kwargs = additional_config_kwargs or {}
+
+ def execute(self, context: Context) -> str:
+ resp = self.hook.create_replication_config(
+ replication_config_id=self.replication_config_id,
+ source_endpoint_arn=self.source_endpoint_arn,
+ target_endpoint_arn=self.target_endpoint_arn,
+ compute_config=self.compute_config,
+ replication_type=self.replication_type,
+ table_mappings=self.table_mappings,
+ additional_config_kwargs=self.additional_config_kwargs,
+ )
+
+ self.log.info("DMS replication config(%s) has been created.",
self.replication_config_id)
+ return resp
+
+
+class DmsDeleteReplicationConfigOperator(AwsBaseOperator[DmsHook]):
+ """
+
Review Comment:
```suggestion
```
##########
providers/src/airflow/providers/amazon/aws/hooks/dms.py:
##########
@@ -219,3 +224,158 @@ def wait_for_task_status(self, replication_task_arn: str,
status: DmsTaskWaiterS
],
WithoutSettings=True,
)
+
Review Comment:
I think some below are quite simple, but others are decent number of lines
that's nice to abstract out of the Operators (start and create for example). So
I'm on the fence here, I could see both ways, but might lean on keeping it how
it is since you've already spent time writing and testing the code as it is?
##########
providers/src/airflow/providers/amazon/aws/triggers/dms.py:
##########
@@ -0,0 +1,221 @@
+# 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
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.hooks.dms import DmsHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class DmsReplicationTerminalStatusTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when an AWS DMS Serverless replication is in a terminal state.
+
+ :param replication_config_arn: The ARN of the replication config.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ replication_config_arn: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ aws_conn_id: str | None = "aws_default",
+ ) -> None:
+ super().__init__(
+ serialized_fields={"replication_config_arn":
replication_config_arn},
+ waiter_name="replication_terminal_status",
+ waiter_delay=waiter_delay,
+ waiter_args={"Filters": [{"Name": "replication-config-arn",
"Values": [replication_config_arn]}]},
+ waiter_max_attempts=waiter_max_attempts,
+ failure_message="Replication failed to reach terminal status.",
+ status_message="Status replication is",
+ status_queries=["Replications[0].Status"],
+ return_key="replication_config_arn",
+ return_value=replication_config_arn,
+ aws_conn_id=aws_conn_id,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return DmsHook(
+ self.aws_conn_id,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+
+class DmsReplicationConfigDeletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when an AWS DMS Serverless replication config is deleted.
+
+ :param replication_config_arn: The ARN of the replication config.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ replication_config_arn: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ aws_conn_id: str | None = "aws_default",
+ ) -> None:
+ super().__init__(
+ serialized_fields={"replication_config_arn":
replication_config_arn},
+ waiter_name="replication_config_deleted",
+ waiter_delay=waiter_delay,
+ waiter_args={"Filters": [{"Name": "replication-config-arn",
"Values": [replication_config_arn]}]},
+ waiter_max_attempts=waiter_max_attempts,
+ failure_message="Replication config failed to be deleted.",
+ status_message="Status replication config is",
+ status_queries=["ReplicationConfigs[0].Status"],
+ return_key="replication_config_arn",
+ return_value=replication_config_arn,
+ aws_conn_id=aws_conn_id,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return DmsHook(
+ self.aws_conn_id,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+
+class DmsReplicationCompleteTrigger(AwsBaseWaiterTrigger):
+ """
+ Trigger when an AWS DMS Serverless replication completes.
+
+ :param replication_config_arn: The ARN of the replication config.
+ :param waiter_delay: The amount of time in seconds to wait between
attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ """
+
+ def __init__(
+ self,
+ replication_config_arn: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ aws_conn_id: str | None = "aws_default",
+ ) -> None:
+ super().__init__(
+ serialized_fields={"replication_config_arn":
replication_config_arn},
+ waiter_name="replication_complete",
+ waiter_delay=waiter_delay,
+ waiter_args={"Filters": [{"Name": "replication-config-arn",
"Values": [replication_config_arn]}]},
+ waiter_max_attempts=waiter_max_attempts,
+ failure_message="Replication failed to reach terminal status.",
Review Comment:
```suggestion
failure_message="Replication failed to complete.",
```
Otherwise it's the same message as `DmsReplicationTerminalStatusTrigger`
uses.
##########
providers/src/airflow/providers/amazon/aws/operators/dms.py:
##########
@@ -277,3 +288,551 @@ def execute(self, context: Context):
"""Stop AWS DMS replication task from Airflow."""
self.hook.stop_replication_task(replication_task_arn=self.replication_task_arn)
self.log.info("DMS replication task(%s) is stopping.",
self.replication_task_arn)
+
+
+class DmsDescribeReplicationConfigsOperator(AwsBaseOperator[DmsHook]):
+ """
+ Describes AWS DMS Serverless replication configurations.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DmsDescribeReplicationConfigsOperator`
+
+ :param describe_config_filter: Filters block for filtering results.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used.
If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ """
+
+ aws_hook_class = DmsHook
+ template_fields: Sequence[str] = aws_template_fields("filter")
+ template_fields_renderers = {"filter": "json"}
+
+ def __init__(
+ self,
+ *,
+ filter: list[dict] | None = None,
+ aws_conn_id: str | None = "aws_default",
+ **kwargs,
+ ):
+ super().__init__(aws_conn_id=aws_conn_id, **kwargs)
+ self.filter = filter
+
+ def execute(self, context: Context) -> list:
+ """
+ Describe AWS DMS replication configurations.
+
+ :return: List of replication configurations
+ """
+ return self.hook.describe_replication_configs(filters=self.filter)
+
+
+class DmsCreateReplicationConfigOperator(AwsBaseOperator[DmsHook]):
+ """
+
Review Comment:
```suggestion
```
##########
docs/apache-airflow-providers-amazon/operators/dms.rst:
##########
@@ -114,6 +114,69 @@ To delete a replication task you can use
:start-after: [START howto_operator_dms_delete_task]
:end-before: [END howto_operator_dms_delete_task]
+
+Create a serverless replication config
+======================================
+
+To create a serverless replication config use
+:class:`~airflow.providers.amazon.aws.operators.dms.DmsCreateReplicationConfigOperator`.
+
+.. exampleinclude::
/../../tests/system/providers/amazon/aws/example_dms_serverless.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_dms_create_replication_config]
+ :end-before: [END howto_operator_dms_create_replication_config]
+
+Describe a serverless replication config
+========================================
+
+To describe a serverless replication config use
+:class:`~airflow.providers.amazon.aws.operators.dms.DmsDescribeReplicationConfigsOperator`.
+
+.. exampleinclude::
/../../tests/system/providers/amazon/aws/example_dms_serverless.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_dms_describe_replication_config]
+ :end-before: [END howto_operator_dms_describe_replication_config]
+
+Start a serverless replication
+==============================
+
+To start a serverless replication use
+:class:`~airflow.providers.amazon.aws.operators.dms.DmsStartReplicationOperator`.
+
+.. exampleinclude::
/../../tests/system/providers/amazon/aws/example_dms_serverless.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_dms_serverless_start_replication]
+ :end-before: [END howto_operator_dms_serverless_start_replication]
+
+Get the status of a serverless replication
+==========================================
+
+To get the status of a serverless replication use
+:class:`~airflow.providers.amazon.aws.operators.dms.DmsDescribeReplicationsOperator`.
+
+.. exampleinclude::
/../../tests/system/providers/amazon/aws/example_dms_serverless.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_dms_serverless_describe_replication]
+ :end-before: [END howto_operator_dms_serverless_describe_replication]
+
+Delete a serverless replication configuration
+=============================================
+
+To delete a serverless replication config use
+:class:`~airflow.providers.amazon.aws.operators.dms.DmsDescribeReplicationsOperator`.
Review Comment:
Copy/paste error?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
