apache-airflow-providers-amazon 9.9.0rc1__py3-none-any.whl → 9.9.1rc1__py3-none-any.whl

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

@@ -60,7 +60,6 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
     :param script_args: etl script arguments and AWS Glue arguments (templated)
     :param retry_limit: The maximum number of times to retry this job if it fails
     :param num_of_dpus: Number of AWS Glue DPUs to allocate to this Job.
-    :param region_name: aws region name (example: us-east-1)
     :param s3_bucket: S3 bucket where logs and local etl script will be uploaded
     :param iam_role_name: AWS IAM Role for Glue Job Execution. If set `iam_role_arn` must equal None.
     :param iam_role_arn: AWS IAM ARN for Glue Job Execution. If set `iam_role_name` must equal None.
@@ -78,7 +77,20 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         of limiting concurrency, Glue needs 5-10 seconds to clean up resources.
         Thus if status is returned immediately it might end up in case of more than 1 concurrent run.
         It is recommended to set this parameter to 10 when you are using concurrency=1.
-        For more information see: https://repost.aws/questions/QUaKgpLBMPSGWO0iq2Fob_bw/glue-run-concurrent-jobs#ANFpCL2fRnQRqgDFuIU_rpvA
+        For more information see:
+        https://repost.aws/questions/QUaKgpLBMPSGWO0iq2Fob_bw/glue-run-concurrent-jobs#ANFpCL2fRnQRqgDFuIU_rpvA
+    :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 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 or not 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 = GlueJobHook
@@ -122,9 +134,11 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         verbose: bool = False,
         replace_script_file: bool = False,
         update_config: bool = False,
-        job_poll_interval: int | float = 6,
         stop_job_run_on_kill: bool = False,
         sleep_before_return: int = 0,
+        job_poll_interval: int | float = 6,
+        waiter_delay: int = 60,
+        waiter_max_attempts: int = 75,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -152,6 +166,8 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
         self._job_run_id: str | None = None
         self.sleep_before_return: int = sleep_before_return
         self.s3_script_location: str | None = None
+        self.waiter_delay = waiter_delay
+        self.waiter_max_attempts = waiter_max_attempts
 
     @property
     def _hook_parameters(self):
@@ -231,7 +247,8 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
                     run_id=self._job_run_id,
                     verbose=self.verbose,
                     aws_conn_id=self.aws_conn_id,
-                    job_poll_interval=self.job_poll_interval,
+                    waiter_delay=self.waiter_delay,
+                    waiter_max_attempts=self.waiter_max_attempts,
                 ),
                 method_name="execute_complete",
             )
@@ -254,7 +271,7 @@ class GlueJobOperator(AwsBaseOperator[GlueJobHook]):
 
         if validated_event["status"] != "success":
             raise AirflowException(f"Error in glue job: {validated_event}")
-        return validated_event["value"]
+        return validated_event["run_id"]
 
     def on_kill(self):
         """Cancel the running AWS Glue Job."""
@@ -282,7 +299,6 @@ class GlueDataQualityOperator(AwsBaseOperator[GlueDataQualityHook]):
     :param description: A description of the data quality ruleset.
     :param update_rule_set: To update existing ruleset, Set this flag to True. (default: False)
     :param data_quality_ruleset_kwargs: Extra arguments for RuleSet.
-
     :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
@@ -378,7 +394,6 @@ class GlueDataQualityRuleSetEvaluationRunOperator(AwsBaseOperator[GlueDataQualit
     :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
@@ -543,7 +558,6 @@ class GlueDataQualityRuleRecommendationRunOperator(AwsBaseOperator[GlueDataQuali
     :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

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

@@ -159,7 +159,7 @@ class RedshiftDataOperator(AwsBaseOperator[RedshiftDataHook]):
         self.statement_id: str = query_execution_output.statement_id
 
         if query_execution_output.session_id:
-            self.xcom_push(context, key="session_id", value=query_execution_output.session_id)
+            context["ti"].xcom_push(key="session_id", value=query_execution_output.session_id)
 
         if self.deferrable and self.wait_for_completion:
             is_finished: bool = self.hook.check_query_is_finished(self.statement_id)

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

@@ -20,8 +20,8 @@ import datetime
 import json
 import time
 import urllib
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, Callable, ClassVar
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from botocore.exceptions import ClientError
 

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

@@ -24,7 +24,6 @@ from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
     SageMakerNotebookHook,
 )
@@ -34,6 +33,7 @@ from airflow.providers.amazon.aws.links.sagemaker_unified_studio import (
 from airflow.providers.amazon.aws.triggers.sagemaker_unified_studio import (
     SageMakerNotebookJobTrigger,
 )
+from airflow.providers.amazon.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/amazon/aws/sensors/base_aws.py

@@ -25,7 +25,7 @@ from airflow.providers.amazon.aws.utils.mixins import (
     AwsHookType,
     aws_template_fields,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.version_compat import BaseSensorOperator
 from airflow.utils.types import NOTSET, ArgNotSet
 
 

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

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
@@ -28,16 +27,16 @@ from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.triggers.glue import (
     GlueDataQualityRuleRecommendationRunCompleteTrigger,
     GlueDataQualityRuleSetEvaluationRunCompleteTrigger,
+    GlueJobCompleteTrigger,
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
-from airflow.sensors.base import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class GlueJobSensor(BaseSensorOperator):
+class GlueJobSensor(AwsBaseSensor[GlueJobHook]):
     """
     Waits for an AWS Glue Job to reach any of the status below.
 
@@ -50,9 +49,29 @@ class GlueJobSensor(BaseSensorOperator):
     :param job_name: The AWS Glue Job unique name
     :param run_id: The AWS Glue current running job identifier
     :param verbose: If True, more Glue Job Run logs show in the Airflow Task Logs.  (default: False)
+    :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore
+        module to be installed.
+        (default: False, but can be overridden in config file by setting default_deferrable to True)
+    :param poke_interval: Polling period in seconds to check for the status of the job. (default: 120)
+    :param max_retries: Number of times before returning the current state. (default: 60)
+
+    :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
     """
 
-    template_fields: Sequence[str] = ("job_name", "run_id")
+    SUCCESS_STATES = ("SUCCEEDED",)
+    FAILURE_STATES = ("FAILED", "STOPPED", "TIMEOUT")
+
+    aws_hook_class = GlueJobHook
+    template_fields: Sequence[str] = aws_template_fields("job_name", "run_id")
 
     def __init__(
         self,
@@ -60,6 +79,9 @@ class GlueJobSensor(BaseSensorOperator):
         job_name: str,
         run_id: str,
         verbose: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        poke_interval: int = 120,
+        max_retries: int = 60,
         aws_conn_id: str | None = "aws_default",
         **kwargs,
     ):
@@ -67,24 +89,46 @@ class GlueJobSensor(BaseSensorOperator):
         self.job_name = job_name
         self.run_id = run_id
         self.verbose = verbose
+        self.deferrable = deferrable
+        self.poke_interval = poke_interval
+        self.max_retries = max_retries
         self.aws_conn_id = aws_conn_id
-        self.success_states: list[str] = ["SUCCEEDED"]
-        self.errored_states: list[str] = ["FAILED", "STOPPED", "TIMEOUT"]
         self.next_log_tokens = GlueJobHook.LogContinuationTokens()
 
-    @cached_property
-    def hook(self):
-        return GlueJobHook(aws_conn_id=self.aws_conn_id)
+    def execute(self, context: Context) -> Any:
+        if self.deferrable:
+            self.defer(
+                trigger=GlueJobCompleteTrigger(
+                    job_name=self.job_name,
+                    run_id=self.run_id,
+                    verbose=self.verbose,
+                    aws_conn_id=self.aws_conn_id,
+                    waiter_delay=int(self.poke_interval),
+                    waiter_max_attempts=self.max_retries,
+                ),
+                method_name="execute_complete",
+            )
+        else:
+            super().execute(context=context)
 
-    def poke(self, context: Context):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+        validated_event = validate_execute_complete_event(event)
+
+        if validated_event["status"] != "success":
+            message = f"Error: AWS Glue Job: {validated_event}"
+            raise AirflowException(message)
+
+        self.log.info("AWS Glue Job completed.")
+
+    def poke(self, context: Context) -> bool:
         self.log.info("Poking for job run status :for Glue Job %s and ID %s", self.job_name, self.run_id)
         job_state = self.hook.get_job_state(job_name=self.job_name, run_id=self.run_id)
 
         try:
-            if job_state in self.success_states:
+            if job_state in self.SUCCESS_STATES:
                 self.log.info("Exiting Job %s Run State: %s", self.run_id, job_state)
                 return True
-            if job_state in self.errored_states:
+            if job_state in self.FAILURE_STATES:
                 job_error_message = "Exiting Job %s Run State: %s", self.run_id, job_state
                 self.log.info(job_error_message)
                 raise AirflowException(job_error_message)

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

@@ -21,9 +21,9 @@ import fnmatch
 import inspect
 import os
 import re
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Callable, cast
+from typing import TYPE_CHECKING, Any, cast
 
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event

airflow/providers/amazon/aws/sensors/sagemaker_unified_studio.py

@@ -25,7 +25,7 @@ from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker_unified_studio import (
     SageMakerNotebookHook,
 )
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.amazon.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -22,8 +22,8 @@ import tempfile
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 
 try:
     from airflow.providers.microsoft.azure.hooks.wasb import WasbHook

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

@@ -21,8 +21,8 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.utils.types import NOTSET, ArgNotSet
 
 

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

@@ -20,13 +20,13 @@ from __future__ import annotations
 
 import json
 import os
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from copy import copy
 from datetime import datetime
 from decimal import Decimal
 from functools import cached_property
 from tempfile import NamedTemporaryFile
-from typing import IO, TYPE_CHECKING, Any, Callable
+from typing import IO, TYPE_CHECKING, Any
 from uuid import uuid4
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

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

@@ -23,8 +23,8 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.exasol.hooks.exasol import ExasolHook
 
 if TYPE_CHECKING:

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

@@ -21,8 +21,8 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.ftp.hooks.ftp import FTPHook
 
 if TYPE_CHECKING:

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

@@ -26,8 +26,8 @@ from typing import TYPE_CHECKING
 from packaging.version import Version
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:

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

@@ -21,8 +21,8 @@ import tempfile
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 
 if TYPE_CHECKING:

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

@@ -24,9 +24,9 @@ import sys
 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
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.google.common.hooks.discovery_api import GoogleDiscoveryApiHook
 
 if TYPE_CHECKING:

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

@@ -20,11 +20,11 @@
 from __future__ import annotations
 
 import json
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Callable
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.apache.hive.hooks.hive import HiveServer2Hook
 
 if TYPE_CHECKING:

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

@@ -22,8 +22,8 @@ from __future__ import annotations
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.http.hooks.http import HttpHook
 
 if TYPE_CHECKING:

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

@@ -22,8 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.imap.hooks.imap import ImapHook
 
 if TYPE_CHECKING:

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

@@ -20,8 +20,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -23,8 +23,8 @@ from typing import TYPE_CHECKING, Any, cast
 
 from bson import json_util
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.mongo.hooks.mongo import MongoHook
 
 if TYPE_CHECKING:

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

@@ -24,11 +24,11 @@ from collections.abc import Iterable, Mapping, Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:

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

@@ -23,8 +23,8 @@ from typing import TYPE_CHECKING, Any, Literal, TypedDict
 from botocore.exceptions import ClientError, WaiterError
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
+from airflow.providers.amazon.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,8 +21,8 @@ from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.ftp.hooks.ftp import FTPHook
 
 if TYPE_CHECKING:

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

@@ -20,11 +20,11 @@ from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.redshift_data import RedshiftDataHook
 from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.utils.types import NOTSET, ArgNotSet
 
 if TYPE_CHECKING:

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

@@ -22,8 +22,8 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
 if TYPE_CHECKING:

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

@@ -16,15 +16,19 @@
 # under the License.
 from __future__ import annotations
 
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from functools import cached_property
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

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

@@ -21,8 +21,8 @@ import tempfile
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.salesforce.hooks.salesforce import SalesforceHook
 
 if TYPE_CHECKING:

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

@@ -22,8 +22,8 @@ from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
 from urllib.parse import urlsplit
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
 from airflow.providers.ssh.hooks.ssh import SSHHook
 
 if TYPE_CHECKING:

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

@@ -22,14 +22,16 @@ import gzip
 import io
 from collections import namedtuple
 from collections.abc import Iterable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any, cast
-
-from typing_extensions import Literal
+from typing import TYPE_CHECKING, Any, Literal, cast
 
 from airflow.exceptions import AirflowException
-from airflow.hooks.base import BaseHook
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
+from airflow.providers.amazon.version_compat import BaseOperator
+
+try:
+    from airflow.sdk import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
 
 if TYPE_CHECKING:
     import pandas as pd

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

@@ -88,7 +88,6 @@ class AwsBaseWaiterTrigger(BaseTrigger):
         super().__init__()
         # parameters that should be hardcoded in the child's implem
         self.serialized_fields = serialized_fields
-
         self.waiter_name = waiter_name
         self.waiter_args = waiter_args
         self.failure_message = failure_message