apache-airflow-providers-amazon 8.16.0__py3-none-any.whl → 8.17.0__py3-none-any.whl

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

@@ -26,6 +26,21 @@ from airflow.providers.amazon.aws.utils import trim_none_values
 
 if TYPE_CHECKING:
     from mypy_boto3_redshift_data import RedshiftDataAPIServiceClient  # noqa
+    from mypy_boto3_redshift_data.type_defs import DescribeStatementResponseTypeDef
+
+FINISHED_STATE = "FINISHED"
+FAILED_STATE = "FAILED"
+ABORTED_STATE = "ABORTED"
+FAILURE_STATES = {FAILED_STATE, ABORTED_STATE}
+RUNNING_STATES = {"PICKED", "STARTED", "SUBMITTED"}
+
+
+class RedshiftDataQueryFailedError(ValueError):
+    """Raise an error that redshift data query failed."""
+
+
+class RedshiftDataQueryAbortedError(ValueError):
+    """Raise an error that redshift data query was aborted."""
 
 
 class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
@@ -108,27 +123,40 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         return statement_id
 
-    def wait_for_results(self, statement_id, poll_interval):
+    def wait_for_results(self, statement_id: str, poll_interval: int) -> str:
         while True:
             self.log.info("Polling statement %s", statement_id)
-            resp = self.conn.describe_statement(
-                Id=statement_id,
-            )
-            status = resp["Status"]
-            if status == "FINISHED":
-                num_rows = resp.get("ResultRows")
-                if num_rows is not None:
-                    self.log.info("Processed %s rows", num_rows)
-                return status
-            elif status in ("FAILED", "ABORTED"):
-                raise ValueError(
-                    f"Statement {statement_id!r} terminated with status {status}. "
-                    f"Response details: {pformat(resp)}"
-                )
-            else:
-                self.log.info("Query %s", status)
+            is_finished = self.check_query_is_finished(statement_id)
+            if is_finished:
+                return FINISHED_STATE
+
             time.sleep(poll_interval)
 
+    def check_query_is_finished(self, statement_id: str) -> bool:
+        """Check whether query finished, raise exception is failed."""
+        resp = self.conn.describe_statement(Id=statement_id)
+        return self.parse_statement_resposne(resp)
+
+    def parse_statement_resposne(self, resp: DescribeStatementResponseTypeDef) -> bool:
+        """Parse the response of describe_statement."""
+        status = resp["Status"]
+        if status == FINISHED_STATE:
+            num_rows = resp.get("ResultRows")
+            if num_rows is not None:
+                self.log.info("Processed %s rows", num_rows)
+            return True
+        elif status in FAILURE_STATES:
+            exception_cls = (
+                RedshiftDataQueryFailedError if status == FAILED_STATE else RedshiftDataQueryAbortedError
+            )
+            raise exception_cls(
+                f"Statement {resp['Id']} terminated with status {status}. "
+                f"Response details: {pformat(resp)}"
+            )
+
+        self.log.info("Query status: %s", status)
+        return False
+
     def get_table_primary_key(
         self,
         table: str,
@@ -201,3 +229,24 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
                 break
 
         return pk_columns or None
+
+    async def is_still_running(self, statement_id: str) -> bool:
+        """Async function to check whether the query is still running.
+
+        :param statement_id: the UUID of the statement
+        """
+        async with self.async_conn as client:
+            desc = await client.describe_statement(Id=statement_id)
+            return desc["Status"] in RUNNING_STATES
+
+    async def check_query_is_finished_async(self, statement_id: str) -> bool:
+        """Async function to check statement is finished.
+
+        It takes statement_id, makes async connection to redshift data to get the query status
+        by statement_id and returns the query status.
+
+        :param statement_id: the UUID of the statement
+        """
+        async with self.async_conn as client:
+            resp = await client.describe_statement(Id=statement_id)
+            return self.parse_statement_resposne(resp)

airflow/providers/amazon/aws/hooks/redshift_sql.py

@@ -239,7 +239,7 @@ class RedshiftSQLHook(DbApiHook):
 
     def _get_identifier_from_hostname(self, hostname: str) -> str:
         parts = hostname.split(".")
-        if "amazonaws.com" in hostname and len(parts) == 6:
+        if hostname.endswith("amazonaws.com") and len(parts) == 6:
             return f"{parts[0]}.{parts[2]}"
         else:
             self.log.debug(

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

@@ -462,7 +462,9 @@ class S3Hook(AwsBaseHook):
         return prefixes
 
     @provide_bucket_name_async
-    async def get_file_metadata_async(self, client: AioBaseClient, bucket_name: str, key: str) -> list[Any]:
+    async def get_file_metadata_async(
+        self, client: AioBaseClient, bucket_name: str, key: str | None = None
+    ) -> list[Any]:
         """
         Get a list of files that a key matching a wildcard expression exists in a bucket asynchronously.
 
@@ -470,7 +472,7 @@ class S3Hook(AwsBaseHook):
         :param bucket_name: the name of the bucket
         :param key: the path to the key
         """
-        prefix = re.split(r"[\[*?]", key, 1)[0]
+        prefix = re.split(r"[\[\*\?]", key, 1)[0] if key else ""
         delimiter = ""
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter)
@@ -486,6 +488,7 @@ class S3Hook(AwsBaseHook):
         bucket_val: str,
         wildcard_match: bool,
         key: str,
+        use_regex: bool = False,
     ) -> bool:
         """
         Get a list of files that a key matching a wildcard expression or get the head object.
@@ -498,6 +501,7 @@ class S3Hook(AwsBaseHook):
         :param bucket_val: the name of the bucket
         :param key: S3 keys that will point to the file
         :param wildcard_match: the path to the key
+        :param use_regex: whether to use regex to check bucket
         """
         bucket_name, key = self.get_s3_bucket_key(bucket_val, key, "bucket_name", "bucket_key")
         if wildcard_match:
@@ -505,6 +509,11 @@ class S3Hook(AwsBaseHook):
             key_matches = [k for k in keys if fnmatch.fnmatch(k["Key"], key)]
             if not key_matches:
                 return False
+        elif use_regex:
+            keys = await self.get_file_metadata_async(client, bucket_name)
+            key_matches = [k for k in keys if re.match(pattern=key, string=k["Key"])]
+            if not key_matches:
+                return False
         else:
             obj = await self.get_head_object_async(client, key, bucket_name)
             if obj is None:
@@ -518,6 +527,7 @@ class S3Hook(AwsBaseHook):
         bucket: str,
         bucket_keys: str | list[str],
         wildcard_match: bool,
+        use_regex: bool = False,
     ) -> bool:
         """
         Get a list of files that a key matching a wildcard expression or get the head object.
@@ -530,14 +540,18 @@ class S3Hook(AwsBaseHook):
         :param bucket: the name of the bucket
         :param bucket_keys: S3 keys that will point to the file
         :param wildcard_match: the path to the key
+        :param use_regex: whether to use regex to check bucket
         """
         if isinstance(bucket_keys, list):
             return all(
                 await asyncio.gather(
-                    *(self._check_key_async(client, bucket, wildcard_match, key) for key in bucket_keys)
+                    *(
+                        self._check_key_async(client, bucket, wildcard_match, key, use_regex)
+                        for key in bucket_keys
+                    )
                 )
             )
-        return await self._check_key_async(client, bucket, wildcard_match, bucket_keys)
+        return await self._check_key_async(client, bucket, wildcard_match, bucket_keys, use_regex)
 
     async def check_for_prefix_async(
         self, client: AioBaseClient, prefix: str, delimiter: str, bucket_name: str | None = None

airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -98,13 +98,13 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None):
         super().set_context(ti)
-        _json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer")
+        _json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer", fallback=None)
         self.handler = watchtower.CloudWatchLogHandler(
             log_group_name=self.log_group,
             log_stream_name=self._render_filename(ti, ti.try_number),
             use_queues=not getattr(ti, "is_trigger_log_context", False),
             boto3_client=self.hook.get_conn(),
-            json_serialize_default=_json_serialize,
+            json_serialize_default=_json_serialize or json_serialize_legacy,
         )
 
     def close(self):

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

@@ -230,7 +230,7 @@ class BatchOperator(BaseOperator):
             region_name=self.region_name,
         )
 
-    def execute(self, context: Context):
+    def execute(self, context: Context) -> str | None:
         """Submit and monitor an AWS Batch job.
 
         :raises: AirflowException
@@ -238,28 +238,46 @@ class BatchOperator(BaseOperator):
         self.submit_job(context)
 
         if self.deferrable:
-            self.defer(
-                timeout=self.execution_timeout,
-                trigger=BatchJobTrigger(
-                    job_id=self.job_id,
-                    waiter_max_attempts=self.max_retries,
-                    aws_conn_id=self.aws_conn_id,
-                    region_name=self.region_name,
-                    waiter_delay=self.poll_interval,
-                ),
-                method_name="execute_complete",
-            )
+            if not self.job_id:
+                raise AirflowException("AWS Batch job - job_id was not found")
+
+            job = self.hook.get_job_description(self.job_id)
+            job_status = job.get("status")
+            if job_status == self.hook.SUCCESS_STATE:
+                self.log.info("Job completed.")
+                return self.job_id
+            elif job_status == self.hook.FAILURE_STATE:
+                raise AirflowException(f"Error while running job: {self.job_id} is in {job_status} state")
+            elif job_status in self.hook.INTERMEDIATE_STATES:
+                self.defer(
+                    timeout=self.execution_timeout,
+                    trigger=BatchJobTrigger(
+                        job_id=self.job_id,
+                        waiter_max_attempts=self.max_retries,
+                        aws_conn_id=self.aws_conn_id,
+                        region_name=self.region_name,
+                        waiter_delay=self.poll_interval,
+                    ),
+                    method_name="execute_complete",
+                )
+
+            raise AirflowException(f"Unexpected status: {job_status}")
 
         if self.wait_for_completion:
             self.monitor_job(context)
 
         return self.job_id
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:
+        if event is None:
+            err_msg = "Trigger error: event is None"
+            self.log.info(err_msg)
+            raise AirflowException(err_msg)
+
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info("Job completed.")
+
+        self.log.info("Job completed.")
         return event["job_id"]
 
     def on_kill(self):

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

@@ -15,66 +15,79 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module contains CloudFormation create/delete stack operators."""
+"""This module contains AWS CloudFormation create/delete stack operators."""
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
 
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.cloud_formation import CloudFormationHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class CloudFormationCreateStackOperator(BaseOperator):
+class CloudFormationCreateStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that creates a CloudFormation stack.
+    An operator that creates a AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:CloudFormationCreateStackOperator`
 
     :param stack_name: stack name (templated)
-    :param cloudformation_parameters: parameters to be passed to CloudFormation.
-    :param aws_conn_id: aws connection to uses
+    :param cloudformation_parameters: parameters to be passed to AWS CloudFormation.
+    :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
     """
 
-    template_fields: Sequence[str] = ("stack_name", "cloudformation_parameters")
-    template_ext: Sequence[str] = ()
+    aws_hook_class = CloudFormationHook
+    template_fields: Sequence[str] = aws_template_fields("stack_name", "cloudformation_parameters")
     ui_color = "#6b9659"
 
-    def __init__(
-        self, *, stack_name: str, cloudformation_parameters: dict, aws_conn_id: str = "aws_default", **kwargs
-    ):
+    def __init__(self, *, stack_name: str, cloudformation_parameters: dict, **kwargs):
         super().__init__(**kwargs)
         self.stack_name = stack_name
         self.cloudformation_parameters = cloudformation_parameters
-        self.aws_conn_id = aws_conn_id
 
     def execute(self, context: Context):
         self.log.info("CloudFormation parameters: %s", self.cloudformation_parameters)
-
-        cloudformation_hook = CloudFormationHook(aws_conn_id=self.aws_conn_id)
-        cloudformation_hook.create_stack(self.stack_name, self.cloudformation_parameters)
+        self.hook.create_stack(self.stack_name, self.cloudformation_parameters)
 
 
-class CloudFormationDeleteStackOperator(BaseOperator):
+class CloudFormationDeleteStackOperator(AwsBaseOperator[CloudFormationHook]):
     """
-    An operator that deletes a CloudFormation stack.
-
-    :param stack_name: stack name (templated)
-    :param cloudformation_parameters: parameters to be passed to CloudFormation.
+    An operator that deletes a AWS CloudFormation stack.
 
     .. seealso::
         For more information on how to use this operator, take a look at the guide:
         :ref:`howto/operator:CloudFormationDeleteStackOperator`
 
-    :param aws_conn_id: aws connection to uses
+    :param stack_name: stack name (templated)
+    :param cloudformation_parameters: parameters to be passed to CloudFormation.
+    :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
     """
 
-    template_fields: Sequence[str] = ("stack_name",)
-    template_ext: Sequence[str] = ()
+    aws_hook_class = CloudFormationHook
+    template_fields: Sequence[str] = aws_template_fields("stack_name")
     ui_color = "#1d472b"
     ui_fgcolor = "#FFF"
 
@@ -93,6 +106,4 @@ class CloudFormationDeleteStackOperator(BaseOperator):
 
     def execute(self, context: Context):
         self.log.info("CloudFormation Parameters: %s", self.cloudformation_parameters)
-
-        cloudformation_hook = CloudFormationHook(aws_conn_id=self.aws_conn_id)
-        cloudformation_hook.delete_stack(self.stack_name, self.cloudformation_parameters)
+        self.hook.delete_stack(self.stack_name, self.cloudformation_parameters)

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

@@ -19,20 +19,20 @@ from __future__ import annotations
 
 import logging
 import random
-from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from deprecated.classic import deprecated
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning, AirflowTaskTimeout
-from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
 
-class DataSyncOperator(BaseOperator):
+class DataSyncOperator(AwsBaseOperator[DataSyncHook]):
     """Find, Create, Update, Execute and Delete AWS DataSync Tasks.
 
     If ``do_xcom_push`` is True, then the DataSync TaskArn and TaskExecutionArn
@@ -46,7 +46,6 @@ class DataSyncOperator(BaseOperator):
         environment. The default behavior is to create a new Task if there are 0, or
         execute the Task if there was 1 Task, or fail if there were many Tasks.
 
-    :param aws_conn_id: AWS connection to use.
     :param wait_interval_seconds: Time to wait between two
         consecutive calls to check TaskExecution status.
     :param max_iterations: Maximum number of
@@ -91,6 +90,16 @@ class DataSyncOperator(BaseOperator):
         ``boto3.start_task_execution(TaskArn=task_arn, **task_execution_kwargs)``
     :param  delete_task_after_execution: If True then the TaskArn which was executed
         will be deleted from AWS DataSync on successful completion.
+    :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
     :raises AirflowException: If ``task_arn`` was not specified, or if
         either ``source_location_uri`` or ``destination_location_uri`` were
         not specified.
@@ -100,7 +109,8 @@ class DataSyncOperator(BaseOperator):
     :raises AirflowException: If Task creation, update, execution or delete fails.
     """
 
-    template_fields: Sequence[str] = (
+    aws_hook_class = DataSyncHook
+    template_fields: Sequence[str] = aws_template_fields(
         "task_arn",
         "source_location_uri",
         "destination_location_uri",
@@ -122,7 +132,6 @@ class DataSyncOperator(BaseOperator):
     def __init__(
         self,
         *,
-        aws_conn_id: str = "aws_default",
         wait_interval_seconds: int = 30,
         max_iterations: int = 60,
         wait_for_completion: bool = True,
@@ -142,7 +151,6 @@ class DataSyncOperator(BaseOperator):
         super().__init__(**kwargs)
 
         # Assignments
-        self.aws_conn_id = aws_conn_id
         self.wait_interval_seconds = wait_interval_seconds
         self.max_iterations = max_iterations
         self.wait_for_completion = wait_for_completion
@@ -185,16 +193,9 @@ class DataSyncOperator(BaseOperator):
         self.destination_location_arn: str | None = None
         self.task_execution_arn: str | None = None
 
-    @cached_property
-    def hook(self) -> DataSyncHook:
-        """Create and return DataSyncHook.
-
-        :return DataSyncHook: An DataSyncHook instance.
-        """
-        return DataSyncHook(
-            aws_conn_id=self.aws_conn_id,
-            wait_interval_seconds=self.wait_interval_seconds,
-        )
+    @property
+    def _hook_parameters(self) -> dict[str, Any]:
+        return {**super()._hook_parameters, "wait_interval_seconds": self.wait_interval_seconds}
 
     @deprecated(reason="use `hook` property instead.", category=AirflowProviderDeprecationWarning)
     def get_hook(self) -> DataSyncHook: