Skip to content

Commit

Permalink
Use AirflowProviderDeprecationWarning in the deprecated decorator in …
Browse files Browse the repository at this point in the history
…Amazon provider (#34488)
  • Loading branch information
hussein-awala committed Sep 19, 2023
1 parent 21e31f6 commit 9924e3e
Show file tree
Hide file tree
Showing 11 changed files with 23 additions and 21 deletions.
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/batch.py
Expand Up @@ -23,7 +23,7 @@
from deprecated import deprecated

from airflow.configuration import conf
from airflow.exceptions import AirflowException, AirflowSkipException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowSkipException
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
from airflow.providers.amazon.aws.triggers.batch import BatchJobTrigger
from airflow.sensors.base import BaseSensorOperator
Expand Down Expand Up @@ -124,7 +124,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
job_id = event["job_id"]
self.log.info("Batch Job %s complete", job_id)

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> BatchClientHook:
"""Create and return a BatchClientHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/dms.py
Expand Up @@ -22,7 +22,7 @@

from deprecated import deprecated

from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.dms import DmsHook
from airflow.sensors.base import BaseSensorOperator

Expand Down Expand Up @@ -62,7 +62,7 @@ def __init__(
self.target_statuses: Iterable[str] = target_statuses or []
self.termination_statuses: Iterable[str] = termination_statuses or []

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> DmsHook:
"""Get DmsHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/emr.py
Expand Up @@ -24,7 +24,7 @@
from deprecated import deprecated

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
from airflow.providers.amazon.aws.links.emr import EmrClusterLink, EmrLogsLink, get_log_uri
from airflow.providers.amazon.aws.triggers.emr import (
Expand Down Expand Up @@ -60,7 +60,7 @@ def __init__(self, *, aws_conn_id: str = "aws_default", **kwargs):
self.target_states: Iterable[str] = [] # will be set in subclasses
self.failed_states: Iterable[str] = [] # will be set in subclasses

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> EmrHook:
return self.hook

Expand Down
Expand Up @@ -24,7 +24,7 @@
from deprecated import deprecated

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.glue_catalog import GlueCatalogHook
from airflow.providers.amazon.aws.triggers.glue import GlueCatalogPartitionTrigger
from airflow.sensors.base import BaseSensorOperator
Expand Down Expand Up @@ -115,7 +115,7 @@ def execute_complete(self, context: Context, event: dict | None = None) -> None:
raise AirflowException(f"Trigger error: event is {event}")
self.log.info("Partition exists in the Glue Catalog")

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> GlueCatalogHook:
"""Get the GlueCatalogHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/glue_crawler.py
Expand Up @@ -22,7 +22,7 @@

from deprecated import deprecated

from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.glue_crawler import GlueCrawlerHook
from airflow.sensors.base import BaseSensorOperator

Expand Down Expand Up @@ -67,7 +67,7 @@ def poke(self, context: Context):
else:
return False

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> GlueCrawlerHook:
"""Return a new or pre-existing GlueCrawlerHook."""
return self.hook
Expand Down
3 changes: 2 additions & 1 deletion airflow/providers/amazon/aws/sensors/redshift_cluster.py
Expand Up @@ -21,6 +21,7 @@

from deprecated import deprecated

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook
from airflow.sensors.base import BaseSensorOperator

Expand Down Expand Up @@ -65,7 +66,7 @@ def poke(self, context: Context):
)
return current_status == self.target_status

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> RedshiftHook:
"""Create and return a RedshiftHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/s3.py
Expand Up @@ -31,7 +31,7 @@
if TYPE_CHECKING:
from airflow.utils.context import Context

from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.amazon.aws.triggers.s3 import S3KeysUnchangedTrigger, S3KeyTrigger
from airflow.sensors.base import BaseSensorOperator, poke_mode_only
Expand Down Expand Up @@ -179,7 +179,7 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> bool | No
raise AirflowException(event["message"])
return None

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> S3Hook:
"""Create and return an S3Hook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/sagemaker.py
Expand Up @@ -22,7 +22,7 @@

from deprecated import deprecated

from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.sagemaker import LogState, SageMakerHook
from airflow.sensors.base import BaseSensorOperator

Expand All @@ -45,7 +45,7 @@ def __init__(self, *, aws_conn_id: str = "aws_default", resource_type: str = "jo
self.aws_conn_id = aws_conn_id
self.resource_type = resource_type # only used for logs, to say what kind of resource we are sensing

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> SageMakerHook:
"""Get SageMakerHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/sqs.py
Expand Up @@ -25,7 +25,7 @@
from typing_extensions import Literal

from airflow.configuration import conf
from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.sqs import SqsHook
from airflow.providers.amazon.aws.triggers.sqs import SqsSensorTrigger
from airflow.providers.amazon.aws.utils.sqs import process_response
Expand Down Expand Up @@ -210,7 +210,7 @@ def poke(self, context: Context):
else:
return False

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> SqsHook:
"""Create and return an SqsHook."""
return self.hook
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/amazon/aws/sensors/step_function.py
Expand Up @@ -22,7 +22,7 @@

from deprecated import deprecated

from airflow.exceptions import AirflowException
from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook
from airflow.sensors.base import BaseSensorOperator

Expand Down Expand Up @@ -85,7 +85,7 @@ def poke(self, context: Context):
self.xcom_push(context, "output", output)
return True

@deprecated(reason="use `hook` property instead.")
@deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
def get_hook(self) -> StepFunctionHook:
"""Create and return a StepFunctionHook."""
return self.hook
Expand Down
5 changes: 3 additions & 2 deletions airflow/providers/amazon/aws/triggers/batch.py
Expand Up @@ -24,6 +24,7 @@
from botocore.exceptions import WaiterError
from deprecated import deprecated

from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.providers.amazon.aws.hooks.batch_client import BatchClientHook
from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
from airflow.triggers.base import BaseTrigger, TriggerEvent
Expand All @@ -32,7 +33,7 @@
from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook


@deprecated(reason="use BatchJobTrigger instead")
@deprecated(reason="use BatchJobTrigger instead", category=AirflowProviderDeprecationWarning)
class BatchOperatorTrigger(BaseTrigger):
"""
Asynchronously poll the boto3 API and wait for the Batch job to be in the `SUCCEEDED` state.
Expand Down Expand Up @@ -109,7 +110,7 @@ async def run(self):
yield TriggerEvent({"status": "failure", "message": "Job Failed - max attempts reached."})


@deprecated(reason="use BatchJobTrigger instead")
@deprecated(reason="use BatchJobTrigger instead", category=AirflowProviderDeprecationWarning)
class BatchSensorTrigger(BaseTrigger):
"""
Checks for the status of a submitted job_id to AWS Batch until it reaches a failure or a success state.
Expand Down

0 comments on commit 9924e3e

Please sign in to comment.