apache-airflow-providers-amazon 8.29.0rc1__py3-none-any.whl → 9.0.0__py3-none-any.whl

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

@@ -16,95 +16,16 @@
 # under the License.
 from __future__ import annotations
 
-from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from deprecated import deprecated
-
-from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.rds import RdsHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.providers.amazon.aws.utils.rds import RdsDbType
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
-from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
     from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 
 
-@deprecated(
-    reason=(
-        "This trigger is deprecated, please use the other RDS triggers "
-        "such as RdsDbDeletedTrigger, RdsDbStoppedTrigger or RdsDbAvailableTrigger"
-    ),
-    category=AirflowProviderDeprecationWarning,
-)
-class RdsDbInstanceTrigger(BaseTrigger):
-    """
-    Deprecated Trigger for RDS operations. Do not use.
-
-    :param waiter_name: Name of the waiter to use, for instance 'db_instance_available'
-        or 'db_instance_deleted'.
-    :param db_instance_identifier: The DB instance identifier for the DB instance to be polled.
-    :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.
-    :param region_name: AWS region where the DB is located, if different from the default one.
-    :param response: The response from the RdsHook, to be passed back to the operator.
-    """
-
-    def __init__(
-        self,
-        waiter_name: str,
-        db_instance_identifier: str,
-        waiter_delay: int,
-        waiter_max_attempts: int,
-        aws_conn_id: str | None,
-        region_name: str | None,
-        response: dict[str, Any],
-    ):
-        self.db_instance_identifier = db_instance_identifier
-        self.waiter_delay = waiter_delay
-        self.waiter_max_attempts = waiter_max_attempts
-        self.aws_conn_id = aws_conn_id
-        self.region_name = region_name
-        self.waiter_name = waiter_name
-        self.response = response
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        return (
-            # dynamically generate the fully qualified name of the class
-            self.__class__.__module__ + "." + self.__class__.__qualname__,
-            {
-                "db_instance_identifier": self.db_instance_identifier,
-                "waiter_delay": str(self.waiter_delay),
-                "waiter_max_attempts": str(self.waiter_max_attempts),
-                "aws_conn_id": self.aws_conn_id,
-                "region_name": self.region_name,
-                "waiter_name": self.waiter_name,
-                "response": self.response,
-            },
-        )
-
-    @cached_property
-    def hook(self) -> RdsHook:
-        return RdsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-
-    async def run(self):
-        async with self.hook.async_conn as client:
-            waiter = client.get_waiter(self.waiter_name)
-            await async_wait(
-                waiter=waiter,
-                waiter_delay=int(self.waiter_delay),
-                waiter_max_attempts=int(self.waiter_max_attempts),
-                args={"DBInstanceIdentifier": self.db_instance_identifier},
-                failure_message="Error checking DB Instance status",
-                status_message="DB instance status is",
-                status_args=["DBInstances[0].DBInstanceStatus"],
-            )
-        yield TriggerEvent({"status": "success", "response": self.response})
-
-
 _waiter_arg = {
     RdsDbType.INSTANCE.value: "DBInstanceIdentifier",
     RdsDbType.CLUSTER.value: "DBClusterIdentifier",

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

@@ -17,11 +17,9 @@
 from __future__ import annotations
 
 import asyncio
-import warnings
 from typing import TYPE_CHECKING, Any, AsyncIterator
 
-from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftAsyncHook, RedshiftHook
+from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
 from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
@@ -45,21 +43,10 @@ class RedshiftCreateClusterTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         cluster_identifier: str,
-        poll_interval: int | None = None,
-        max_attempt: int | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 15,
         waiter_max_attempts: int = 999999,
     ):
-        if poll_interval is not None or max_attempt is not None:
-            warnings.warn(
-                "please use waiter_delay instead of poll_interval "
-                "and waiter_max_attempts instead of max_attempt.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = poll_interval or waiter_delay
-            waiter_max_attempts = max_attempt or waiter_max_attempts
         super().__init__(
             serialized_fields={"cluster_identifier": cluster_identifier},
             waiter_name="cluster_available",
@@ -93,21 +80,10 @@ class RedshiftPauseClusterTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         cluster_identifier: str,
-        poll_interval: int | None = None,
-        max_attempts: int | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 15,
         waiter_max_attempts: int = 999999,
     ):
-        if poll_interval is not None or max_attempts is not None:
-            warnings.warn(
-                "please use waiter_delay instead of poll_interval "
-                "and waiter_max_attempts instead of max_attempt.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = poll_interval or waiter_delay
-            waiter_max_attempts = max_attempts or waiter_max_attempts
         super().__init__(
             serialized_fields={"cluster_identifier": cluster_identifier},
             waiter_name="cluster_paused",
@@ -141,21 +117,10 @@ class RedshiftCreateClusterSnapshotTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         cluster_identifier: str,
-        poll_interval: int | None = None,
-        max_attempts: int | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 15,
         waiter_max_attempts: int = 999999,
     ):
-        if poll_interval is not None or max_attempts is not None:
-            warnings.warn(
-                "please use waiter_delay instead of poll_interval "
-                "and waiter_max_attempts instead of max_attempt.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = poll_interval or waiter_delay
-            waiter_max_attempts = max_attempts or waiter_max_attempts
         super().__init__(
             serialized_fields={"cluster_identifier": cluster_identifier},
             waiter_name="snapshot_available",
@@ -189,21 +154,10 @@ class RedshiftResumeClusterTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         cluster_identifier: str,
-        poll_interval: int | None = None,
-        max_attempts: int | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 15,
         waiter_max_attempts: int = 999999,
     ):
-        if poll_interval is not None or max_attempts is not None:
-            warnings.warn(
-                "please use waiter_delay instead of poll_interval "
-                "and waiter_max_attempts instead of max_attempt.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = poll_interval or waiter_delay
-            waiter_max_attempts = max_attempts or waiter_max_attempts
         super().__init__(
             serialized_fields={"cluster_identifier": cluster_identifier},
             waiter_name="cluster_resumed",
@@ -234,21 +188,10 @@ class RedshiftDeleteClusterTrigger(AwsBaseWaiterTrigger):
     def __init__(
         self,
         cluster_identifier: str,
-        poll_interval: int | None = None,
-        max_attempts: int | None = None,
         aws_conn_id: str | None = "aws_default",
         waiter_delay: int = 30,
         waiter_max_attempts: int = 30,
     ):
-        if poll_interval is not None or max_attempts is not None:
-            warnings.warn(
-                "please use waiter_delay instead of poll_interval "
-                "and waiter_max_attempts instead of max_attempt.",
-                AirflowProviderDeprecationWarning,
-                stacklevel=2,
-            )
-            waiter_delay = poll_interval or waiter_delay
-            waiter_max_attempts = max_attempts or waiter_max_attempts
         super().__init__(
             serialized_fields={"cluster_identifier": cluster_identifier},
             waiter_name="cluster_deleted",
@@ -304,13 +247,11 @@ class RedshiftClusterTrigger(BaseTrigger):
     async def run(self) -> AsyncIterator[TriggerEvent]:
         """Run async until the cluster status matches the target status."""
         try:
-            hook = RedshiftAsyncHook(aws_conn_id=self.aws_conn_id)
+            hook = RedshiftHook(aws_conn_id=self.aws_conn_id)
             while True:
-                res = await hook.cluster_status(self.cluster_identifier)
-                if (res["status"] == "success" and res["cluster_state"] == self.target_status) or res[
-                    "status"
-                ] == "error":
-                    yield TriggerEvent(res)
+                status = await hook.cluster_status_async(self.cluster_identifier)
+                if status == self.target_status:
+                    yield TriggerEvent({"status": "success", "message": "target state met"})
                     return
                 await asyncio.sleep(self.poke_interval)
         except Exception as e:

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

@@ -18,17 +18,15 @@
 from __future__ import annotations
 
 import asyncio
-import time
 from collections import Counter
 from enum import IntEnum
 from functools import cached_property
 from typing import Any, AsyncIterator
 
 from botocore.exceptions import WaiterError
-from deprecated import deprecated
 
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
-from airflow.providers.amazon.aws.hooks.sagemaker import LogState, SageMakerHook
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
@@ -198,92 +196,3 @@ class SageMakerPipelineTrigger(BaseTrigger):
                     await asyncio.sleep(int(self.waiter_delay))
 
             raise AirflowException("Waiter error: max attempts reached")
-
-
-@deprecated(
-    reason=(
-        "`airflow.providers.amazon.aws.triggers.sagemaker.SageMakerTrainingPrintLogTrigger` "
-        "has been deprecated and will be removed in future. Please use ``SageMakerTrigger`` instead."
-    ),
-    category=AirflowProviderDeprecationWarning,
-)
-class SageMakerTrainingPrintLogTrigger(BaseTrigger):
-    """
-    SageMakerTrainingPrintLogTrigger is fired as deferred class with params to run the task in triggerer.
-
-    :param job_name: name of the job to check status
-    :param poke_interval:  polling period in seconds to check for the status
-    :param aws_conn_id: AWS connection ID for sagemaker
-    """
-
-    def __init__(
-        self,
-        job_name: str,
-        poke_interval: float,
-        aws_conn_id: str | None = "aws_default",
-    ):
-        super().__init__()
-        self.job_name = job_name
-        self.poke_interval = poke_interval
-        self.aws_conn_id = aws_conn_id
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serialize SageMakerTrainingPrintLogTrigger arguments and classpath."""
-        return (
-            "airflow.providers.amazon.aws.triggers.sagemaker.SageMakerTrainingPrintLogTrigger",
-            {
-                "poke_interval": self.poke_interval,
-                "aws_conn_id": self.aws_conn_id,
-                "job_name": self.job_name,
-            },
-        )
-
-    @cached_property
-    def hook(self) -> SageMakerHook:
-        return SageMakerHook(aws_conn_id=self.aws_conn_id)
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        """Make async connection to sagemaker async hook and gets job status for a job submitted by the operator."""
-        stream_names: list[str] = []  # The list of log streams
-        positions: dict[str, Any] = {}  # The current position in each stream, map of stream name -> position
-
-        last_description = await self.hook.describe_training_job_async(self.job_name)
-        instance_count = last_description["ResourceConfig"]["InstanceCount"]
-        status = last_description["TrainingJobStatus"]
-        job_already_completed = status not in self.hook.non_terminal_states
-        state = LogState.COMPLETE if job_already_completed else LogState.TAILING
-        last_describe_job_call = time.time()
-        try:
-            while True:
-                (
-                    state,
-                    last_description,
-                    last_describe_job_call,
-                ) = await self.hook.describe_training_job_with_log_async(
-                    self.job_name,
-                    positions,
-                    stream_names,
-                    instance_count,
-                    state,
-                    last_description,
-                    last_describe_job_call,
-                )
-                status = last_description["TrainingJobStatus"]
-                if status in self.hook.non_terminal_states:
-                    await asyncio.sleep(self.poke_interval)
-                elif status in self.hook.failed_states:
-                    reason = last_description.get("FailureReason", "(No reason provided)")
-                    error_message = f"SageMaker job failed because {reason}"
-                    yield TriggerEvent({"status": "error", "message": error_message})
-                    return
-                else:
-                    billable_seconds = SageMakerHook.count_billable_seconds(
-                        training_start_time=last_description["TrainingStartTime"],
-                        training_end_time=last_description["TrainingEndTime"],
-                        instance_count=instance_count,
-                    )
-                    self.log.info("Billable seconds: %d", billable_seconds)
-                    yield TriggerEvent({"status": "success", "message": last_description})
-                    return
-        except Exception as e:
-            yield TriggerEvent({"status": "error", "message": str(e)})

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

@@ -0,0 +1,106 @@
+# 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 importlib.util import find_spec
+
+
+def _get_asset_compat_hook_lineage_collector():
+    from airflow.lineage.hook import get_hook_lineage_collector
+
+    collector = get_hook_lineage_collector()
+
+    if all(
+        getattr(collector, asset_method_name, None)
+        for asset_method_name in ("add_input_asset", "add_output_asset", "collected_assets")
+    ):
+        return collector
+
+    # dataset is renamed as asset in Airflow 3.0
+
+    from functools import wraps
+
+    from airflow.lineage.hook import DatasetLineageInfo, HookLineage
+
+    DatasetLineageInfo.asset = DatasetLineageInfo.dataset
+
+    def rename_dataset_kwargs_as_assets_kwargs(function):
+        @wraps(function)
+        def wrapper(*args, **kwargs):
+            if "asset_kwargs" in kwargs:
+                kwargs["dataset_kwargs"] = kwargs.pop("asset_kwargs")
+
+            if "asset_extra" in kwargs:
+                kwargs["dataset_extra"] = kwargs.pop("asset_extra")
+
+            return function(*args, **kwargs)
+
+        return wrapper
+
+    collector.create_asset = rename_dataset_kwargs_as_assets_kwargs(collector.create_dataset)
+    collector.add_input_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_input_dataset)
+    collector.add_output_asset = rename_dataset_kwargs_as_assets_kwargs(collector.add_output_dataset)
+
+    def collected_assets_compat(collector) -> HookLineage:
+        """Get the collected hook lineage information."""
+        lineage = collector.collected_datasets
+        return HookLineage(
+            [
+                DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context)
+                for item in lineage.inputs
+            ],
+            [
+                DatasetLineageInfo(dataset=item.dataset, count=item.count, context=item.context)
+                for item in lineage.outputs
+            ],
+        )
+
+    setattr(
+        collector.__class__,
+        "collected_assets",
+        property(lambda collector: collected_assets_compat(collector)),
+    )
+
+    return collector
+
+
+def get_hook_lineage_collector():
+    # HookLineageCollector added in 2.10
+    try:
+        if find_spec("airflow.assets"):
+            # Dataset has been renamed as Asset in 3.0
+            from airflow.lineage.hook import get_hook_lineage_collector
+
+            return get_hook_lineage_collector()
+
+        return _get_asset_compat_hook_lineage_collector()
+    except ImportError:
+
+        class NoOpCollector:
+            """
+            NoOpCollector is a hook lineage collector that does nothing.
+
+            It is used when you want to disable lineage collection.
+            """
+
+            def add_input_asset(self, *_, **__):
+                pass
+
+            def add_output_asset(self, *_, **__):
+                pass
+
+        return NoOpCollector()