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

airflow/providers/amazon/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "8.25.0"
+__version__ = "8.26.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.7.0"

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

@@ -50,7 +50,8 @@ def query_params_to_string(params: dict[str, str | Collection[str]]) -> str:
 
 
 class AthenaHook(AwsBaseHook):
-    """Interact with Amazon Athena.
+    """
+    Interact with Amazon Athena.
 
     Provide thick wrapper around
     :external+boto3:py:class:`boto3.client("athena") <Athena.Client>`.
@@ -106,7 +107,8 @@ class AthenaHook(AwsBaseHook):
         client_request_token: str | None = None,
         workgroup: str = "primary",
     ) -> str:
-        """Run a Trino/Presto query on Athena with provided config.
+        """
+        Run a Trino/Presto query on Athena with provided config.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.start_query_execution`
@@ -136,7 +138,8 @@ class AthenaHook(AwsBaseHook):
         return query_execution_id
 
     def get_query_info(self, query_execution_id: str, use_cache: bool = False) -> dict:
-        """Get information about a single execution of a query.
+        """
+        Get information about a single execution of a query.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.get_query_execution`
@@ -152,7 +155,8 @@ class AthenaHook(AwsBaseHook):
         return response
 
     def check_query_status(self, query_execution_id: str, use_cache: bool = False) -> str | None:
-        """Fetch the state of a submitted query.
+        """
+        Fetch the state of a submitted query.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.get_query_execution`
@@ -200,7 +204,8 @@ class AthenaHook(AwsBaseHook):
     def get_query_results(
         self, query_execution_id: str, next_token_id: str | None = None, max_results: int = 1000
     ) -> dict | None:
-        """Fetch submitted query results.
+        """
+        Fetch submitted query results.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.get_query_results`
@@ -234,7 +239,8 @@ class AthenaHook(AwsBaseHook):
         page_size: int | None = None,
         starting_token: str | None = None,
     ) -> PageIterator | None:
-        """Fetch submitted Athena query results.
+        """
+        Fetch submitted Athena query results.
 
         .. seealso::
             - :external+boto3:py:class:`Athena.Paginator.GetQueryResults`
@@ -274,7 +280,8 @@ class AthenaHook(AwsBaseHook):
     def poll_query_status(
         self, query_execution_id: str, max_polling_attempts: int | None = None, sleep_time: int | None = None
     ) -> str | None:
-        """Poll the state of a submitted query until it reaches final state.
+        """
+        Poll the state of a submitted query until it reaches final state.
 
         :param query_execution_id: ID of submitted athena query
         :param max_polling_attempts: Number of times to poll for query state before function exits
@@ -299,7 +306,8 @@ class AthenaHook(AwsBaseHook):
             return self.check_query_status(query_execution_id)
 
     def get_output_location(self, query_execution_id: str) -> str:
-        """Get the output location of the query results in S3 URI format.
+        """
+        Get the output location of the query results in S3 URI format.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.get_query_execution`
@@ -319,7 +327,8 @@ class AthenaHook(AwsBaseHook):
             raise
 
     def stop_query(self, query_execution_id: str) -> dict:
-        """Cancel the submitted query.
+        """
+        Cancel the submitted query.
 
         .. seealso::
             - :external+boto3:py:meth:`Athena.Client.stop_query_execution`

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

@@ -33,7 +33,8 @@ if TYPE_CHECKING:
 
 
 class AthenaSQLHook(AwsBaseHook, DbApiHook):
-    """Interact with Amazon Athena.
+    """
+    Interact with Amazon Athena.
 
     Provide wrapper around PyAthena library.
 

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

@@ -70,9 +70,23 @@ if TYPE_CHECKING:
 
     from airflow.models.connection import Connection  # Avoid circular imports.
 
+_loader = botocore.loaders.Loader()
+"""
+botocore data loader to be used with async sessions
+
+By default, a botocore session creates and caches an instance of JSONDecoder which
+consumes a lot of memory.  This issue was reported here https://github.com/boto/botocore/issues/3078.
+In the context of triggers which use boto sessions, this can result in excessive
+memory usage and as a result reduced capacity on the triggerer.  We can reduce
+memory footprint by sharing the loader instance across the sessions.
+
+:meta private:
+"""
+
 
 class BaseSessionFactory(LoggingMixin):
-    """Base AWS Session Factory class.
+    """
+    Base AWS Session Factory class.
 
     This handles synchronous and async boto session creation. It can handle most
     of the AWS supported authentication methods.
@@ -154,7 +168,9 @@ class BaseSessionFactory(LoggingMixin):
     def get_async_session(self):
         from aiobotocore.session import get_session as async_get_session
 
-        return async_get_session()
+        session = async_get_session()
+        session.register_component("data_loader", _loader)
+        return session
 
     def create_session(
         self, deferrable: bool = False
@@ -435,7 +451,8 @@ class BaseSessionFactory(LoggingMixin):
 
 
 class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
-    """Generic class for interact with AWS.
+    """
+    Generic class for interact with AWS.
 
     This class provide a thin wrapper around the boto3 Python library.
 
@@ -497,7 +514,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
     @staticmethod
     def _find_operator_class_name(target_function_name: str) -> str | None:
-        """Given a frame off the stack, return the name of the class that made the call.
+        """
+        Given a frame off the stack, return the name of the class that made the call.
 
         This method may raise a ValueError or an IndexError. The caller is
         responsible with catching and handling those.
@@ -547,7 +565,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
     @staticmethod
     @return_on_error("00000000-0000-0000-0000-000000000000")
     def _generate_dag_key() -> str:
-        """Generate a DAG key.
+        """
+        Generate a DAG key.
 
         The Object Identifier (OID) namespace is used to salt the dag_id value.
         That salted value is used to generate a SHA-1 hash which, by definition,
@@ -783,7 +802,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         return creds
 
     def expand_role(self, role: str, region_name: str | None = None) -> str:
-        """Get the Amazon Resource Name (ARN) for the role.
+        """
+        Get the Amazon Resource Name (ARN) for the role.
 
         If IAM role is already an IAM role ARN, the value is returned unchanged.
 
@@ -861,7 +881,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         }
 
     def test_connection(self):
-        """Test the AWS connection by call AWS STS (Security Token Service) GetCallerIdentity API.
+        """
+        Test the AWS connection by call AWS STS (Security Token Service) GetCallerIdentity API.
 
         .. seealso::
             https://docs.aws.amazon.com/STS/latest/APIReference/API_GetCallerIdentity.html
@@ -898,7 +919,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         deferrable: bool = False,
         client=None,
     ) -> Waiter:
-        """Get a waiter by name.
+        """
+        Get a waiter by name.
 
         First checks if there is a custom waiter with the provided waiter_name and
         uses that if it exists, otherwise it will check the service client for a
@@ -969,7 +991,8 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
 
 class AwsBaseHook(AwsGenericHook[Union[boto3.client, boto3.resource]]):
-    """Base class for interact with AWS.
+    """
+    Base class for interact with AWS.
 
     This class provide a thin wrapper around the boto3 Python library.
 
@@ -1128,7 +1151,8 @@ class BaseAsyncSessionFactory(BaseSessionFactory):
     category=AirflowProviderDeprecationWarning,
 )
 class AwsBaseAsyncHook(AwsBaseHook):
-    """Interacts with AWS using aiobotocore asynchronously.
+    """
+    Interacts with AWS using aiobotocore asynchronously.
 
     :param aws_conn_id: The Airflow connection used for AWS credentials.
         If this is None or empty then the default botocore behaviour is used. If

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

@@ -30,7 +30,8 @@ from airflow.providers.http.hooks.http import HttpHook
 
 
 class ChimeWebhookHook(HttpHook):
-    """Interact with Amazon Chime Webhooks to create notifications.
+    """
+    Interact with Amazon Chime Webhooks to create notifications.
 
     .. warning:: This hook is only designed to work with web hooks and not chat bots.
 

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

@@ -136,7 +136,8 @@ class DataSyncHook(AwsBaseHook):
     def create_task(
         self, source_location_arn: str, destination_location_arn: str, **create_task_kwargs
     ) -> str:
-        """Create a Task between the specified source and destination LocationArns.
+        """
+        Create a Task between the specified source and destination LocationArns.
 
         .. seealso::
             - :external+boto3:py:meth:`DataSync.Client.create_task`
@@ -155,7 +156,8 @@ class DataSyncHook(AwsBaseHook):
         return task["TaskArn"]
 
     def update_task(self, task_arn: str, **update_task_kwargs) -> None:
-        """Update a Task.
+        """
+        Update a Task.
 
         .. seealso::
             - :external+boto3:py:meth:`DataSync.Client.update_task`
@@ -166,7 +168,8 @@ class DataSyncHook(AwsBaseHook):
         self.get_conn().update_task(TaskArn=task_arn, **update_task_kwargs)
 
     def delete_task(self, task_arn: str) -> None:
-        """Delete a Task.
+        """
+        Delete a Task.
 
         .. seealso::
             - :external+boto3:py:meth:`DataSync.Client.delete_task`

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

@@ -70,7 +70,8 @@ class EcrHook(AwsBaseHook):
         super().__init__(**kwargs)
 
     def get_temporary_credentials(self, registry_ids: list[str] | str | None = None) -> list[EcrCredentials]:
-        """Get temporary credentials for Amazon ECR.
+        """
+        Get temporary credentials for Amazon ECR.
 
         .. seealso::
             - :external+boto3:py:meth:`ECR.Client.get_authorization_token`

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

@@ -152,42 +152,48 @@ class EcsProtocol(Protocol):
     """
 
     def run_task(self, **kwargs) -> dict:
-        """Run a task.
+        """
+        Run a task.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.run_task
         """
         ...
 
     def get_waiter(self, x: str) -> Waiter:
-        """Get a waiter.
+        """
+        Get a waiter.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.get_waiter
         """
         ...
 
     def describe_tasks(self, cluster: str, tasks) -> dict:
-        """Describe tasks.
+        """
+        Describe tasks.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.describe_tasks
         """
         ...
 
     def stop_task(self, cluster, task, reason: str) -> dict:
-        """Stop a task.
+        """
+        Stop a task.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.stop_task
         """
         ...
 
     def describe_task_definition(self, taskDefinition: str) -> dict:
-        """Describe a task definition.
+        """
+        Describe a task definition.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.describe_task_definition
         """
         ...
 
     def list_tasks(self, cluster: str, launchType: str, desiredStatus: str, family: str) -> dict:
-        """List tasks.
+        """
+        List tasks.
 
         https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ecs.html#ECS.Client.list_tasks
         """

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

@@ -23,7 +23,8 @@ from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
 
 class GlacierHook(AwsBaseHook):
-    """Interact with Amazon Glacier.
+    """
+    Interact with Amazon Glacier.
 
     This is a thin wrapper around
     :external+boto3:py:class:`boto3.client("glacier") <Glacier.Client>`.
@@ -40,7 +41,8 @@ class GlacierHook(AwsBaseHook):
         super().__init__(*args, **kwargs)
 
     def retrieve_inventory(self, vault_name: str) -> dict[str, Any]:
-        """Initiate an Amazon Glacier inventory-retrieval job.
+        """
+        Initiate an Amazon Glacier inventory-retrieval job.
 
         .. seealso::
             - :external+boto3:py:meth:`Glacier.Client.initiate_job`
@@ -55,7 +57,8 @@ class GlacierHook(AwsBaseHook):
         return response
 
     def retrieve_inventory_results(self, vault_name: str, job_id: str) -> dict[str, Any]:
-        """Retrieve the results of an Amazon Glacier inventory-retrieval job.
+        """
+        Retrieve the results of an Amazon Glacier inventory-retrieval job.
 
         .. seealso::
             - :external+boto3:py:meth:`Glacier.Client.get_job_output`
@@ -68,7 +71,8 @@ class GlacierHook(AwsBaseHook):
         return response
 
     def describe_job(self, vault_name: str, job_id: str) -> dict[str, Any]:
-        """Retrieve the status of an Amazon S3 Glacier job.
+        """
+        Retrieve the status of an Amazon S3 Glacier job.
 
         .. seealso::
             - :external+boto3:py:meth:`Glacier.Client.describe_job`

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

@@ -45,7 +45,8 @@ class FirehoseHook(AwsBaseHook):
         super().__init__(*args, **kwargs)
 
     def put_records(self, records: Iterable):
-        """Write batch records to Kinesis Firehose.
+        """
+        Write batch records to Kinesis Firehose.
 
         .. seealso::
             - :external+boto3:py:meth:`Firehose.Client.put_record_batch`

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

@@ -158,7 +158,8 @@ class AwsLogsHook(AwsBaseHook):
     async def describe_log_streams_async(
         self, log_group: str, stream_prefix: str, order_by: str, count: int
     ) -> dict[str, Any] | None:
-        """Async function to get the list of log streams for the specified log group.
+        """
+        Async function to get the list of log streams for the specified log group.
 
         You can list all the log streams or filter the results by prefix. You can also control
         how the results are ordered.
@@ -193,7 +194,8 @@ class AwsLogsHook(AwsBaseHook):
         skip: int = 0,
         start_from_head: bool = True,
     ) -> AsyncGenerator[Any, dict[str, Any]]:
-        """Yield all the available items in a single log stream.
+        """
+        Yield all the available items in a single log stream.
 
         :param log_group: The name of the log group.
         :param log_stream_name: The name of the specific stream.

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