apache-airflow-providers-amazon 9.1.0rc4__py3-none-any.whl → 9.2.0rc1__py3-none-any.whl

airflow/providers/amazon/aws/transfers/exasol_to_s3.py

@@ -19,8 +19,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/ftp_to_s3.py

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/gcs_to_s3.py

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from packaging.version import Version
 
@@ -123,9 +124,9 @@ class GCSToS3Operator(BaseOperator):
         self.s3_acl_policy = s3_acl_policy
         self.keep_directory_structure = keep_directory_structure
         try:
-            from airflow.providers.google import __version__
+            from airflow.providers.google import __version__ as _GOOGLE_PROVIDER_VERSION
 
-            if Version(__version__) >= Version("10.3.0"):
+            if Version(Version(_GOOGLE_PROVIDER_VERSION).base_version) >= Version("10.3.0"):
                 self.__is_match_glob_supported = True
             else:
                 self.__is_match_glob_supported = False

airflow/providers/amazon/aws/transfers/glacier_to_gcs.py

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import tempfile
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook

airflow/providers/amazon/aws/transfers/google_api_to_s3.py

@@ -21,7 +21,8 @@ from __future__ import annotations
 
 import json
 import sys
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.models.xcom import MAX_XCOM_SIZE, XCOM_RETURN_KEY
@@ -30,7 +31,6 @@ from airflow.providers.google.common.hooks.discovery_api import GoogleDiscoveryA
 
 if TYPE_CHECKING:
     from airflow.models import TaskInstance
-    from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic
     from airflow.utils.context import Context
 
 
@@ -174,9 +174,7 @@ class GoogleApiToS3Operator(BaseOperator):
             replace=self.s3_overwrite,
         )
 
-    def _update_google_api_endpoint_params_via_xcom(
-        self, task_instance: TaskInstance | TaskInstancePydantic
-    ) -> None:
+    def _update_google_api_endpoint_params_via_xcom(self, task_instance: TaskInstance) -> None:
         if self.google_api_endpoint_params_via_xcom:
             google_api_endpoint_params = task_instance.xcom_pull(
                 task_ids=self.google_api_endpoint_params_via_xcom_task_ids,
@@ -184,9 +182,7 @@ class GoogleApiToS3Operator(BaseOperator):
             )
             self.google_api_endpoint_params.update(google_api_endpoint_params)
 
-    def _expose_google_api_response_via_xcom(
-        self, task_instance: TaskInstance | TaskInstancePydantic, data: dict
-    ) -> None:
+    def _expose_google_api_response_via_xcom(self, task_instance: TaskInstance, data: dict) -> None:
         if sys.getsizeof(data) < MAX_XCOM_SIZE:
             task_instance.xcom_push(key=self.google_api_response_via_xcom or XCOM_RETURN_KEY, value=data)
         else:

airflow/providers/amazon/aws/transfers/hive_to_dynamodb.py

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Callable, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Callable
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook

airflow/providers/amazon/aws/transfers/imap_attachment_to_s3.py

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/local_to_s3.py

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/mongo_to_s3.py

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any, Iterable, Sequence, cast
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING, Any, cast
 
 from bson import json_util
 

airflow/providers/amazon/aws/transfers/redshift_to_s3.py

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Iterable, Mapping, Sequence
+from collections.abc import Iterable, Mapping, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator

airflow/providers/amazon/aws/transfers/s3_to_dynamodb.py

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Sequence, TypedDict
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, Literal, TypedDict
 
 from botocore.exceptions import ClientError, WaiterError
 

airflow/providers/amazon/aws/transfers/s3_to_ftp.py

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/s3_to_redshift.py

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable, Sequence
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator

airflow/providers/amazon/aws/transfers/s3_to_sftp.py

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
 from airflow.models import BaseOperator

airflow/providers/amazon/aws/transfers/s3_to_sql.py

@@ -16,9 +16,10 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Iterable, Sequence
 from functools import cached_property
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Callable, Iterable, Sequence
+from typing import TYPE_CHECKING, Callable
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.base import BaseHook

airflow/providers/amazon/aws/transfers/salesforce_to_s3.py

@@ -18,7 +18,8 @@ from __future__ import annotations
 
 import os
 import tempfile
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

airflow/providers/amazon/aws/transfers/sftp_to_s3.py

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
 from airflow.models import BaseOperator
@@ -49,6 +50,8 @@ class SFTPToS3Operator(BaseOperator):
         uploading the file to S3.
     :param use_temp_file: If True, copies file first to local,
         if False streams file from SFTP to S3.
+    :param fail_on_file_not_exist: If True, operator fails when file does not exist,
+        if False, operator will not fail and skips transfer. Default is True.
     """
 
     template_fields: Sequence[str] = ("s3_key", "sftp_path", "s3_bucket")
@@ -62,6 +65,7 @@ class SFTPToS3Operator(BaseOperator):
         sftp_conn_id: str = "ssh_default",
         s3_conn_id: str = "aws_default",
         use_temp_file: bool = True,
+        fail_on_file_not_exist: bool = True,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -71,6 +75,7 @@ class SFTPToS3Operator(BaseOperator):
         self.s3_key = s3_key
         self.s3_conn_id = s3_conn_id
         self.use_temp_file = use_temp_file
+        self.fail_on_file_not_exist = fail_on_file_not_exist
 
     @staticmethod
     def get_s3_key(s3_key: str) -> str:
@@ -85,6 +90,14 @@ class SFTPToS3Operator(BaseOperator):
 
         sftp_client = ssh_hook.get_conn().open_sftp()
 
+        try:
+            sftp_client.stat(self.sftp_path)
+        except FileNotFoundError:
+            if self.fail_on_file_not_exist:
+                raise
+            self.log.info("File %s not found on SFTP server. Skipping transfer.", self.sftp_path)
+            return
+
         if self.use_temp_file:
             with NamedTemporaryFile("w") as f:
                 sftp_client.get(self.sftp_path, f.name)

airflow/providers/amazon/aws/transfers/sql_to_s3.py

@@ -19,8 +19,9 @@ from __future__ import annotations
 
 import enum
 from collections import namedtuple
+from collections.abc import Iterable, Mapping, Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Any, Iterable, Mapping, Sequence, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from typing_extensions import Literal
 

airflow/providers/amazon/aws/triggers/base.py

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, AsyncIterator
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
 from airflow.triggers.base import BaseTrigger, TriggerEvent

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 complete.",
+            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 DmsReplicationStoppedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an AWS DMS Serverless replication is stopped.
+
+    :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_stopped",
+            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 stop.",
+            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 DmsReplicationDeprovisionedTrigger(AwsBaseWaiterTrigger):
+    """
+    Trigger when an AWS DMS Serverless replication is de-provisioned.
+
+    :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_deprovisioned",
+            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 deprovision.",
+            status_message="Status replication is",
+            status_queries=["Replications[0].ProvisionData.ProvisionState"],
+            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,
+        )

airflow/providers/amazon/aws/triggers/glue.py

@@ -18,8 +18,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncIterator
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, AsyncIterator
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook

airflow/providers/amazon/aws/triggers/redshift_cluster.py

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import asyncio
-from typing import TYPE_CHECKING, Any, AsyncIterator
+from collections.abc import AsyncIterator
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger

airflow/providers/amazon/aws/triggers/redshift_data.py

@@ -18,8 +18,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncIterator
 from functools import cached_property
-from typing import Any, AsyncIterator
+from typing import Any
 
 from airflow.providers.amazon.aws.hooks.redshift_data import (
     ABORTED_STATE,

airflow/providers/amazon/aws/triggers/s3.py

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import AsyncIterator
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, AsyncIterator
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.triggers.base import BaseTrigger, TriggerEvent

airflow/providers/amazon/aws/triggers/sagemaker.py

@@ -19,9 +19,10 @@ from __future__ import annotations
 
 import asyncio
 from collections import Counter
+from collections.abc import AsyncIterator
 from enum import IntEnum
 from functools import cached_property
-from typing import Any, AsyncIterator
+from typing import Any
 
 from botocore.exceptions import WaiterError
 

airflow/providers/amazon/aws/triggers/sqs.py

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import asyncio
-from typing import TYPE_CHECKING, Any, AsyncIterator, Collection
+from collections.abc import AsyncIterator, Collection
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook

airflow/providers/amazon/aws/utils/__init__.py

@@ -31,21 +31,7 @@ log = logging.getLogger(__name__)
 
 
 def trim_none_values(obj: dict):
-    from packaging.version import Version
-
-    from airflow.version import version
-
-    if Version(version) < Version("2.7"):
-        # before version 2.7, the behavior is not the same.
-        # Empty dict and lists are removed from the given dict.
-        return {key: val for key, val in obj.items() if val is not None}
-    else:
-        # once airflow 2.6 rolls out of compatibility support for provider packages,
-        # we can replace usages of this method with the core one in our code,
-        # and uncomment this warning for users who may use it.
-        # warnings.warn("use airflow.utils.helpers.prune_dict() instead",
-        #     AirflowProviderDeprecationWarning, stacklevel=2)
-        return prune_dict(obj)
+    return prune_dict(obj)
 
 
 def datetime_to_epoch(date_time: datetime) -> int:

airflow/providers/amazon/aws/utils/task_log_fetcher.py

@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Generator
 from datetime import datetime, timedelta, timezone
 from threading import Event, Thread
-from typing import TYPE_CHECKING, Generator
+from typing import TYPE_CHECKING
 
 from botocore.exceptions import ClientError, ConnectionClosedError
 

airflow/providers/amazon/aws/utils/waiter.py

@@ -19,6 +19,7 @@ from __future__ import annotations
 
 import logging
 import time
+from enum import Enum
 from typing import Callable
 
 from airflow.exceptions import AirflowException
@@ -83,3 +84,22 @@ def get_state(response, keys) -> str:
         if value is not None:
             value = value.get(key, None)
     return value
+
+
+class WaitPolicy(str, Enum):
+    """
+    Used to control the waiting behaviour within EMRClusterJobFlowOperator.
+
+    Choices:
+    - WAIT_FOR_COMPLETION - Will wait for the cluster to report "Running" state
+    - WAIT_FOR_STEPS_COMPLETION - Will wait for the cluster to report "Terminated" state
+    """
+
+    WAIT_FOR_COMPLETION = "wait_for_completion"
+    WAIT_FOR_STEPS_COMPLETION = "wait_for_steps_completion"
+
+
+WAITER_POLICY_NAME_MAPPING: dict[WaitPolicy, str] = {
+    WaitPolicy.WAIT_FOR_COMPLETION: "job_flow_waiting",
+    WaitPolicy.WAIT_FOR_STEPS_COMPLETION: "job_flow_terminated",
+}