apache-airflow-providers-amazon 8.25.0rc1__py3-none-any.whl → 8.26.0__py3-none-any.whl

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

@@ -27,7 +27,8 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseAsyncHook, AwsBas
 
 
 class RedshiftHook(AwsBaseHook):
-    """Interact with Amazon Redshift.
+    """
+    Interact with Amazon Redshift.
 
     This is a thin wrapper around
     :external+boto3:py:class:`boto3.client("redshift") <Redshift.Client>`.
@@ -53,7 +54,8 @@ class RedshiftHook(AwsBaseHook):
         master_user_password: str,
         params: dict[str, Any],
     ) -> dict[str, Any]:
-        """Create a new cluster with the specified parameters.
+        """
+        Create a new cluster with the specified parameters.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.create_cluster`
@@ -80,7 +82,8 @@ class RedshiftHook(AwsBaseHook):
 
     # TODO: Wrap create_cluster_snapshot
     def cluster_status(self, cluster_identifier: str) -> str:
-        """Get status of a cluster.
+        """
+        Get status of a cluster.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.describe_clusters`
@@ -101,7 +104,8 @@ class RedshiftHook(AwsBaseHook):
         skip_final_cluster_snapshot: bool = True,
         final_cluster_snapshot_identifier: str | None = None,
     ):
-        """Delete a cluster and optionally create a snapshot.
+        """
+        Delete a cluster and optionally create a snapshot.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.delete_cluster`
@@ -120,7 +124,8 @@ class RedshiftHook(AwsBaseHook):
         return response["Cluster"] if response["Cluster"] else None
 
     def describe_cluster_snapshots(self, cluster_identifier: str) -> list[str] | None:
-        """List snapshots for a cluster.
+        """
+        List snapshots for a cluster.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.describe_cluster_snapshots`
@@ -136,7 +141,8 @@ class RedshiftHook(AwsBaseHook):
         return snapshots
 
     def restore_from_cluster_snapshot(self, cluster_identifier: str, snapshot_identifier: str) -> str:
-        """Restore a cluster from its snapshot.
+        """
+        Restore a cluster from its snapshot.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.restore_from_cluster_snapshot`
@@ -156,7 +162,8 @@ class RedshiftHook(AwsBaseHook):
         retention_period: int = -1,
         tags: list[Any] | None = None,
     ) -> str:
-        """Create a snapshot of a cluster.
+        """
+        Create a snapshot of a cluster.
 
         .. seealso::
             - :external+boto3:py:meth:`Redshift.Client.create_cluster_snapshot`
@@ -178,7 +185,8 @@ class RedshiftHook(AwsBaseHook):
         return response["Snapshot"] if response["Snapshot"] else None
 
     def get_cluster_snapshot_status(self, snapshot_identifier: str):
-        """Get Redshift cluster snapshot status.
+        """
+        Get Redshift cluster snapshot status.
 
         If cluster snapshot not found, *None* is returned.
 
@@ -210,7 +218,8 @@ class RedshiftAsyncHook(AwsBaseAsyncHook):
         super().__init__(*args, **kwargs)
 
     async def cluster_status(self, cluster_identifier: str, delete_operation: bool = False) -> dict[str, Any]:
-        """Get the cluster status.
+        """
+        Get the cluster status.
 
         :param cluster_identifier: unique identifier of a cluster
         :param delete_operation: whether the method has been called as part of delete cluster operation
@@ -228,7 +237,8 @@ class RedshiftAsyncHook(AwsBaseAsyncHook):
                 return {"status": "error", "message": str(error)}
 
     async def pause_cluster(self, cluster_identifier: str, poll_interval: float = 5.0) -> dict[str, Any]:
-        """Pause the cluster.
+        """
+        Pause the cluster.
 
         :param cluster_identifier: unique identifier of a cluster
         :param poll_interval: polling period in seconds to check for the status
@@ -255,7 +265,8 @@ class RedshiftAsyncHook(AwsBaseAsyncHook):
         cluster_identifier: str,
         polling_period_seconds: float = 5.0,
     ) -> dict[str, Any]:
-        """Resume the cluster.
+        """
+        Resume the cluster.
 
         :param cluster_identifier: unique identifier of a cluster
         :param polling_period_seconds: polling period in seconds to check for the status
@@ -284,7 +295,8 @@ class RedshiftAsyncHook(AwsBaseAsyncHook):
         flag: asyncio.Event,
         delete_operation: bool = False,
     ) -> dict[str, Any]:
-        """Check for expected Redshift cluster state.
+        """
+        Check for expected Redshift cluster state.
 
         :param cluster_identifier: unique identifier of a cluster
         :param expected_state: expected_state example("available", "pausing", "paused"")

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

@@ -231,7 +231,8 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
         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.
+        """
+        Async function to check whether the query is still running.
 
         :param statement_id: the UUID of the statement
         """
@@ -240,7 +241,8 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
             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.
+        """
+        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.

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

@@ -39,7 +39,8 @@ if TYPE_CHECKING:
 
 
 class RedshiftSQLHook(DbApiHook):
-    """Execute statements against Amazon Redshift.
+    """
+    Execute statements against Amazon Redshift.
 
     This hook requires the redshift_conn_id connection.
 
@@ -103,7 +104,8 @@ class RedshiftSQLHook(DbApiHook):
         return conn_params
 
     def get_iam_token(self, conn: Connection) -> tuple[str, str, int]:
-        """Retrieve a temporary password to connect to Redshift.
+        """
+        Retrieve a temporary password to connect to Redshift.
 
         Port is required. If none is provided, default is used for each service.
         """
@@ -177,7 +179,8 @@ class RedshiftSQLHook(DbApiHook):
         return create_engine(self.get_uri(), **engine_kwargs)
 
     def get_table_primary_key(self, table: str, schema: str | None = "public") -> list[str] | None:
-        """Get the table's primary key.
+        """
+        Get the table's primary key.
 
         :param table: Name of the target table
         :param schema: Name of the target schema, public by default

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

@@ -22,6 +22,7 @@ from __future__ import annotations
 import asyncio
 import fnmatch
 import gzip as gz
+import inspect
 import logging
 import os
 import re
@@ -36,7 +37,7 @@ from inspect import signature
 from io import BytesIO
 from pathlib import Path
 from tempfile import NamedTemporaryFile, gettempdir
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable
 from urllib.parse import urlsplit
 from uuid import uuid4
 
@@ -61,42 +62,24 @@ from airflow.utils.helpers import chunks
 logger = logging.getLogger(__name__)
 
 
+# Explicit value that would remove ACLs from a copy
+# No conflicts with Canned ACLs:
+#   https://docs.aws.amazon.com/AmazonS3/latest/userguide/acl-overview.html#canned-acl
+NO_ACL = "no-acl"
+
+
 def provide_bucket_name(func: Callable) -> Callable:
     """Provide a bucket name taken from the connection if no bucket name has been passed to the function."""
     if hasattr(func, "_unify_bucket_name_and_key_wrapped"):
         logger.warning("`unify_bucket_name_and_key` should wrap `provide_bucket_name`.")
-    function_signature = signature(func)
-
-    @wraps(func)
-    def wrapper(*args, **kwargs) -> Callable:
-        bound_args = function_signature.bind(*args, **kwargs)
-
-        if "bucket_name" not in bound_args.arguments:
-            self = args[0]
-
-            if "bucket_name" in self.service_config:
-                bound_args.arguments["bucket_name"] = self.service_config["bucket_name"]
-            elif self.conn_config and self.conn_config.schema:
-                warnings.warn(
-                    "s3 conn_type, and the associated schema field, is deprecated. "
-                    "Please use aws conn_type instead, and specify `bucket_name` "
-                    "in `service_config.s3` within `extras`.",
-                    AirflowProviderDeprecationWarning,
-                    stacklevel=2,
-                )
-                bound_args.arguments["bucket_name"] = self.conn_config.schema
-
-        return func(*bound_args.args, **bound_args.kwargs)
 
-    return wrapper
-
-
-def provide_bucket_name_async(func: Callable) -> Callable:
-    """Provide a bucket name taken from the connection if no bucket name has been passed to the function."""
     function_signature = signature(func)
+    if "bucket_name" not in function_signature.parameters:
+        raise RuntimeError(
+            "Decorator provide_bucket_name should only wrap a function with param 'bucket_name'."
+        )
 
-    @wraps(func)
-    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+    async def maybe_add_bucket_name(*args, **kwargs):
         bound_args = function_signature.bind(*args, **kwargs)
 
         if "bucket_name" not in bound_args.arguments:
@@ -105,8 +88,46 @@ def provide_bucket_name_async(func: Callable) -> Callable:
                 connection = await sync_to_async(self.get_connection)(self.aws_conn_id)
                 if connection.schema:
                     bound_args.arguments["bucket_name"] = connection.schema
+        return bound_args
+
+    if inspect.iscoroutinefunction(func):
+
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            bound_args = await maybe_add_bucket_name(*args, **kwargs)
+            print(f"invoking async function {func=}")
+            return await func(*bound_args.args, **bound_args.kwargs)
+
+    elif inspect.isasyncgenfunction(func):
+
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            bound_args = await maybe_add_bucket_name(*args, **kwargs)
+            async for thing in func(*bound_args.args, **bound_args.kwargs):
+                yield thing
+
+    else:
+
+        @wraps(func)
+        def wrapper(*args, **kwargs) -> Callable:
+            bound_args = function_signature.bind(*args, **kwargs)
+
+            if "bucket_name" not in bound_args.arguments:
+                self = args[0]
+
+                if "bucket_name" in self.service_config:
+                    bound_args.arguments["bucket_name"] = self.service_config["bucket_name"]
+                elif self.conn_config and self.conn_config.schema:
+                    warnings.warn(
+                        "s3 conn_type, and the associated schema field, is deprecated. "
+                        "Please use aws conn_type instead, and specify `bucket_name` "
+                        "in `service_config.s3` within `extras`.",
+                        AirflowProviderDeprecationWarning,
+                        stacklevel=2,
+                    )
+                    bound_args.arguments["bucket_name"] = self.conn_config.schema
 
-        return await func(*bound_args.args, **bound_args.kwargs)
+            return func(*bound_args.args, **bound_args.kwargs)
 
     return wrapper
 
@@ -400,8 +421,8 @@ class S3Hook(AwsBaseHook):
 
         return prefixes
 
-    @provide_bucket_name_async
     @unify_bucket_name_and_key
+    @provide_bucket_name
     async def get_head_object_async(
         self, client: AioBaseClient, key: str, bucket_name: str | None = None
     ) -> dict[str, Any] | None:
@@ -462,10 +483,10 @@ class S3Hook(AwsBaseHook):
 
         return prefixes
 
-    @provide_bucket_name_async
+    @provide_bucket_name
     async def get_file_metadata_async(
         self, client: AioBaseClient, bucket_name: str, key: str | None = None
-    ) -> list[Any]:
+    ) -> AsyncIterator[Any]:
         """
         Get a list of files that a key matching a wildcard expression exists in a bucket asynchronously.
 
@@ -477,11 +498,10 @@ class S3Hook(AwsBaseHook):
         delimiter = ""
         paginator = client.get_paginator("list_objects_v2")
         response = paginator.paginate(Bucket=bucket_name, Prefix=prefix, Delimiter=delimiter)
-        files = []
         async for page in response:
             if "Contents" in page:
-                files += page["Contents"]
-        return files
+                for row in page["Contents"]:
+                    yield row
 
     async def _check_key_async(
         self,
@@ -506,21 +526,16 @@ class S3Hook(AwsBaseHook):
         """
         bucket_name, key = self.get_s3_bucket_key(bucket_val, key, "bucket_name", "bucket_key")
         if wildcard_match:
-            keys = await self.get_file_metadata_async(client, bucket_name, key)
-            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:
-                return False
-
-        return True
+            async for k in self.get_file_metadata_async(client, bucket_name, key):
+                if fnmatch.fnmatch(k["Key"], key):
+                    return True
+            return False
+        if use_regex:
+            async for k in self.get_file_metadata_async(client, bucket_name):
+                if re.match(pattern=key, string=k["Key"]):
+                    return True
+            return False
+        return bool(await self.get_head_object_async(client, key, bucket_name))
 
     async def check_key_async(
         self,
@@ -1276,6 +1291,8 @@ class S3Hook(AwsBaseHook):
             object to be copied which is private by default.
         """
         acl_policy = acl_policy or "private"
+        if acl_policy != NO_ACL:
+            kwargs["ACL"] = acl_policy
 
         dest_bucket_name, dest_bucket_key = self.get_s3_bucket_key(
             dest_bucket_name, dest_bucket_key, "dest_bucket_name", "dest_bucket_key"
@@ -1287,7 +1304,7 @@ class S3Hook(AwsBaseHook):
 
         copy_source = {"Bucket": source_bucket_name, "Key": source_bucket_key, "VersionId": source_version_id}
         response = self.get_conn().copy_object(
-            Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, ACL=acl_policy, **kwargs
+            Bucket=dest_bucket_name, Key=dest_bucket_key, CopySource=copy_source, **kwargs
         )
         return response