apache-airflow-providers-amazon 8.26.0rc2__py3-none-any.whl → 8.27.0rc1__py3-none-any.whl

airflow/providers/amazon/aws/operators/kinesis_analytics.py

@@ -0,0 +1,348 @@
+# 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, Any, Sequence
+
+from botocore.exceptions import ClientError
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.kinesis_analytics import KinesisAnalyticsV2Hook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.kinesis_analytics import (
+    KinesisAnalyticsV2ApplicationOperationCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class KinesisAnalyticsV2CreateApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Creates an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2CreateApplicationOperator`
+
+    :param application_name: The name of application. (templated)
+    :param runtime_environment: The runtime environment for the application. (templated)
+    :param service_execution_role: The IAM role used by the application to access services. (templated)
+    :param create_application_kwargs: Create application extra properties. (templated)
+    :param application_description: A summary description of the application. (templated)
+
+    :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
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "runtime_environment",
+        "service_execution_role",
+        "create_application_kwargs",
+        "application_description",
+    )
+    template_fields_renderers: dict = {
+        "create_application_kwargs": "json",
+    }
+
+    def __init__(
+        self,
+        application_name: str,
+        runtime_environment: str,
+        service_execution_role: str,
+        create_application_kwargs: dict[str, Any] | None = None,
+        application_description: str = "Managed Service for Apache Flink application created from Airflow",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.runtime_environment = runtime_environment
+        self.service_execution_role = service_execution_role
+        self.create_application_kwargs = create_application_kwargs or {}
+        self.application_description = application_description
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Creating AWS Managed Service for Apache Flink application %s.", self.application_name)
+        try:
+            response = self.hook.conn.create_application(
+                ApplicationName=self.application_name,
+                ApplicationDescription=self.application_description,
+                RuntimeEnvironment=self.runtime_environment,
+                ServiceExecutionRole=self.service_execution_role,
+                **self.create_application_kwargs,
+            )
+        except ClientError as error:
+            raise AirflowException(
+                f"AWS Managed Service for Apache Flink application creation failed: {error.response['Error']['Message']}"
+            )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application created successfully %s.",
+            self.application_name,
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}
+
+
+class KinesisAnalyticsV2StartApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Starts an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2StartApplicationOperator`
+
+    :param application_name: The name of application. (templated)
+    :param run_configuration: Application properties to start Apache Flink Job. (templated)
+
+    :param wait_for_completion: Whether to wait for job to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param deferrable: If True, the operator will wait asynchronously for the job to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :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
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "run_configuration",
+    )
+    template_fields_renderers: dict = {
+        "run_configuration": "json",
+    }
+
+    def __init__(
+        self,
+        application_name: str,
+        run_configuration: dict[str, Any] | None = None,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.run_configuration = run_configuration or {}
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        msg = "AWS Managed Service for Apache Flink application"
+
+        try:
+            self.log.info("Starting %s %s.", msg, self.application_name)
+            self.hook.conn.start_application(
+                ApplicationName=self.application_name, RunConfiguration=self.run_configuration
+            )
+        except ClientError as error:
+            raise AirflowException(
+                f"Failed to start {msg} {self.application_name}: {error.response['Error']['Message']}"
+            )
+
+        describe_response = self.hook.conn.describe_application(ApplicationName=self.application_name)
+
+        if self.deferrable:
+            self.log.info("Deferring for %s to start: %s.", msg, self.application_name)
+            self.defer(
+                trigger=KinesisAnalyticsV2ApplicationOperationCompleteTrigger(
+                    application_name=self.application_name,
+                    waiter_name="application_start_complete",
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for %s to start: %s.", msg, self.application_name)
+
+            self.hook.get_waiter("application_start_complete").wait(
+                ApplicationName=self.application_name,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        self.log.info("%s started successfully %s.", msg, self.application_name)
+
+        return {"ApplicationARN": describe_response["ApplicationDetail"]["ApplicationARN"]}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException(
+                "Error while starting AWS Managed Service for Apache Flink application: %s", event
+            )
+
+        response = self.hook.conn.describe_application(
+            ApplicationName=event["application_name"],
+        )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application %s started successfully.",
+            event["application_name"],
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}
+
+
+class KinesisAnalyticsV2StopApplicationOperator(AwsBaseOperator[KinesisAnalyticsV2Hook]):
+    """
+    Stop an AWS Managed Service for Apache Flink application.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KinesisAnalyticsV2StopApplicationOperator`
+
+    :param application_name: The name of your application. (templated)
+    :param force: Set to true to force the application to stop. If you set Force to true, Managed Service for
+        Apache Flink stops the application without taking a snapshot. (templated)
+
+    :param wait_for_completion: Whether to wait for job to stop. (default: True)
+    :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 20)
+    :param deferrable: If True, the operator will wait asynchronously for the job to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :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
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    aws_hook_class = KinesisAnalyticsV2Hook
+    ui_color = "#44b5e2"
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "application_name",
+        "force",
+    )
+
+    def __init__(
+        self,
+        application_name: str,
+        force: bool = False,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 20,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.application_name = application_name
+        self.force = force
+        self.wait_for_completion = wait_for_completion
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
+        self.deferrable = deferrable
+
+    def execute(self, context: Context) -> dict[str, Any]:
+        msg = "AWS Managed Service for Apache Flink application"
+
+        try:
+            self.log.info("Stopping %s %s.", msg, self.application_name)
+
+            self.hook.conn.stop_application(ApplicationName=self.application_name, Force=self.force)
+        except ClientError as error:
+            raise AirflowException(
+                f"Failed to stop {msg} {self.application_name}: {error.response['Error']['Message']}"
+            )
+
+        describe_response = self.hook.conn.describe_application(ApplicationName=self.application_name)
+
+        if self.deferrable:
+            self.log.info("Deferring for %s to stop: %s.", msg, self.application_name)
+            self.defer(
+                trigger=KinesisAnalyticsV2ApplicationOperationCompleteTrigger(
+                    application_name=self.application_name,
+                    waiter_name="application_stop_complete",
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
+                    region_name=self.region_name,
+                    verify=self.verify,
+                    botocore_config=self.botocore_config,
+                ),
+                method_name="execute_complete",
+            )
+        if self.wait_for_completion:
+            self.log.info("Waiting for %s to stop: %s.", msg, self.application_name)
+
+            self.hook.get_waiter("application_stop_complete").wait(
+                ApplicationName=self.application_name,
+                WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+            )
+
+        self.log.info("%s stopped successfully %s.", msg, self.application_name)
+
+        return {"ApplicationARN": describe_response["ApplicationDetail"]["ApplicationARN"]}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+        event = validate_execute_complete_event(event)
+
+        if event["status"] != "success":
+            raise AirflowException("Error while stopping AWS Managed Service for Apache Flink application")
+
+        response = self.hook.conn.describe_application(
+            ApplicationName=event["application_name"],
+        )
+
+        self.log.info(
+            "AWS Managed Service for Apache Flink application %s stopped successfully.",
+            event["application_name"],
+        )
+
+        return {"ApplicationARN": response["ApplicationDetail"]["ApplicationARN"]}

airflow/providers/amazon/aws/operators/rds.py

@@ -36,6 +36,7 @@ from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
 from airflow.providers.amazon.aws.utils.tags import format_tags
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
+from airflow.utils.helpers import prune_dict
 
 if TYPE_CHECKING:
     from mypy_boto3_rds.type_defs import TagTypeDef
@@ -782,7 +783,7 @@ class RdsStartDbOperator(RdsBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     response=start_db_response,
-                    db_type=RdsDbType.INSTANCE,
+                    db_type=self.db_type,
                 ),
                 method_name="execute_complete",
             )
@@ -881,12 +882,25 @@ class RdsStopDbOperator(RdsBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                     region_name=self.region_name,
                     response=stop_db_response,
-                    db_type=RdsDbType.INSTANCE,
+                    db_type=self.db_type,
                 ),
                 method_name="execute_complete",
             )
         elif self.wait_for_completion:
-            self._wait_until_db_stopped()
+            waiter = self.hook.get_waiter(f"db_{self.db_type.value}_stopped")
+            waiter_key = (
+                "DBInstanceIdentifier" if self.db_type == RdsDbType.INSTANCE else "DBClusterIdentifier"
+            )
+            kwargs = {waiter_key: self.db_identifier}
+            waiter.wait(
+                WaiterConfig=prune_dict(
+                    {
+                        "Delay": self.waiter_delay,
+                        "MaxAttempts": self.waiter_max_attempts,
+                    }
+                ),
+                **kwargs,
+            )
         return json.dumps(stop_db_response, default=str)
 
     def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
@@ -915,23 +929,6 @@ class RdsStopDbOperator(RdsBaseOperator):
             response = self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.db_identifier)
         return response
 
-    def _wait_until_db_stopped(self):
-        self.log.info("Waiting for DB %s to reach 'stopped' state", self.db_type.value)
-        if self.db_type == RdsDbType.INSTANCE:
-            self.hook.wait_for_db_instance_state(
-                self.db_identifier,
-                target_state="stopped",
-                check_interval=self.waiter_delay,
-                max_attempts=self.waiter_max_attempts,
-            )
-        else:
-            self.hook.wait_for_db_cluster_state(
-                self.db_identifier,
-                target_state="stopped",
-                check_interval=self.waiter_delay,
-                max_attempts=self.waiter_max_attempts,
-            )
-
 
 __all__ = [
     "RdsCreateDbSnapshotOperator",

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

@@ -256,8 +256,6 @@ class RedshiftCreateClusterOperator(BaseOperator):
             params["ClusterVersion"] = self.cluster_version
         if self.allow_version_upgrade:
             params["AllowVersionUpgrade"] = self.allow_version_upgrade
-        if self.publicly_accessible:
-            params["PubliclyAccessible"] = self.publicly_accessible
         if self.encrypted:
             params["Encrypted"] = self.encrypted
         if self.hsm_client_certificate_identifier:
@@ -287,6 +285,10 @@ class RedshiftCreateClusterOperator(BaseOperator):
         if self.default_iam_role_arn:
             params["DefaultIamRoleArn"] = self.default_iam_role_arn
 
+        # PubliclyAccessible is True by default on Redshift side, hence, we should always set it regardless
+        # of its value
+        params["PubliclyAccessible"] = self.publicly_accessible
+
         cluster = redshift_hook.create_cluster(
             self.cluster_identifier,
             self.node_type,

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

@@ -324,8 +324,7 @@ class S3CopyObjectOperator(BaseOperator):
         )
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         dest_bucket_name, dest_bucket_key = S3Hook.get_s3_bucket_key(
@@ -439,8 +438,7 @@ class S3CreateObjectOperator(BaseOperator):
             s3_hook.load_bytes(self.data, s3_key, s3_bucket, self.replace, self.encrypt, self.acl_policy)
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         bucket, key = S3Hook.get_s3_bucket_key(self.s3_bucket, self.s3_key, "dest_bucket", "dest_key")
@@ -546,13 +544,12 @@ class S3DeleteObjectsOperator(BaseOperator):
 
     def get_openlineage_facets_on_complete(self, task_instance):
         """Implement _on_complete because object keys are resolved in execute()."""
-        from openlineage.client.facet import (
+        from airflow.providers.common.compat.openlineage.facet import (
+            Dataset,
             LifecycleStateChange,
             LifecycleStateChangeDatasetFacet,
-            LifecycleStateChangeDatasetFacetPreviousIdentifier,
+            PreviousIdentifier,
         )
-        from openlineage.client.run import Dataset
-
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         if not self._keys:
@@ -570,7 +567,7 @@ class S3DeleteObjectsOperator(BaseOperator):
                 facets={
                     "lifecycleStateChange": LifecycleStateChangeDatasetFacet(
                         lifecycleStateChange=LifecycleStateChange.DROP.value,
-                        previousIdentifier=LifecycleStateChangeDatasetFacetPreviousIdentifier(
+                        previousIdentifier=PreviousIdentifier(
                             namespace=bucket_url,
                             name=key,
                         ),
@@ -725,8 +722,7 @@ class S3FileTransformOperator(BaseOperator):
             self.log.info("Upload successful")
 
     def get_openlineage_facets_on_start(self):
-        from openlineage.client.run import Dataset
-
+        from airflow.providers.common.compat.openlineage.facet import Dataset
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         dest_bucket_name, dest_bucket_key = S3Hook.get_s3_bucket_key(

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

@@ -46,8 +46,7 @@ from airflow.utils.helpers import prune_dict
 from airflow.utils.json import AirflowJsonEncoder
 
 if TYPE_CHECKING:
-    from openlineage.client.run import Dataset
-
+    from airflow.providers.common.compat.openlineage.facet import Dataset
     from airflow.providers.openlineage.extractors.base import OperatorLineage
     from airflow.utils.context import Context
 
@@ -208,7 +207,7 @@ class SageMakerBaseOperator(BaseOperator):
 
     @staticmethod
     def path_to_s3_dataset(path) -> Dataset:
-        from openlineage.client.run import Dataset
+        from airflow.providers.common.compat.openlineage.facet import Dataset
 
         path = path.replace("s3://", "")
         split_path = path.split("/")

airflow/providers/amazon/aws/operators/step_function.py

@@ -48,6 +48,8 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
 
     :param state_machine_arn: ARN of the Step Function State Machine
     :param name: The name of the execution.
+    :param is_redrive_execution: Restarts unsuccessful executions of Standard workflows that did not
+            complete successfully in the last 14 days.
     :param state_machine_input: JSON data input to pass to the State Machine
     :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
@@ -73,7 +75,9 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
     """
 
     aws_hook_class = StepFunctionHook
-    template_fields: Sequence[str] = aws_template_fields("state_machine_arn", "name", "input")
+    template_fields: Sequence[str] = aws_template_fields(
+        "state_machine_arn", "name", "input", "is_redrive_execution"
+    )
     ui_color = "#f9c915"
     operator_extra_links = (StateMachineDetailsLink(), StateMachineExecutionsDetailsLink())
 
@@ -82,6 +86,7 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
         *,
         state_machine_arn: str,
         name: str | None = None,
+        is_redrive_execution: bool = False,
         state_machine_input: dict | str | None = None,
         waiter_max_attempts: int = 30,
         waiter_delay: int = 60,
@@ -91,6 +96,7 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
         super().__init__(**kwargs)
         self.state_machine_arn = state_machine_arn
         self.name = name
+        self.is_redrive_execution = is_redrive_execution
         self.input = state_machine_input
         self.waiter_delay = waiter_delay
         self.waiter_max_attempts = waiter_max_attempts
@@ -105,7 +111,11 @@ class StepFunctionStartExecutionOperator(AwsBaseOperator[StepFunctionHook]):
             state_machine_arn=self.state_machine_arn,
         )
 
-        if not (execution_arn := self.hook.start_execution(self.state_machine_arn, self.name, self.input)):
+        if not (
+            execution_arn := self.hook.start_execution(
+                self.state_machine_arn, self.name, self.input, self.is_redrive_execution
+            )
+        ):
             raise AirflowException(f"Failed to start State Machine execution for: {self.state_machine_arn}")
 
         StateMachineExecutionsDetailsLink.persist(