apache-airflow-providers-amazon 9.1.0rc4__py3-none-any.whl → 9.2.0__py3-none-any.whl

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

@@ -18,13 +18,15 @@
 from __future__ import annotations
 
 import ast
+import warnings
+from collections.abc import Sequence
 from datetime import timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 from uuid import uuid4
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.emr import EmrContainerHook, EmrHook, EmrServerlessHook
 from airflow.providers.amazon.aws.links.emr import (
@@ -49,7 +51,11 @@ from airflow.providers.amazon.aws.triggers.emr import (
     EmrTerminateJobFlowTrigger,
 )
 from airflow.providers.amazon.aws.utils import validate_execute_complete_event
-from airflow.providers.amazon.aws.utils.waiter import waiter
+from airflow.providers.amazon.aws.utils.waiter import (
+    WAITER_POLICY_NAME_MAPPING,
+    WaitPolicy,
+    waiter,
+)
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
 from airflow.utils.helpers import exactly_one, prune_dict
 from airflow.utils.types import NOTSET, ArgNotSet
@@ -636,8 +642,14 @@ class EmrCreateJobFlowOperator(BaseOperator):
     :param job_flow_overrides: boto3 style arguments or reference to an arguments file
         (must be '.json') to override specific ``emr_conn_id`` extra parameters. (templated)
     :param region_name: Region named passed to EmrHook
-    :param wait_for_completion: Whether to finish task immediately after creation (False) or wait for jobflow
+    :param wait_for_completion: Deprecated - use `wait_policy` instead.
+        Whether to finish task immediately after creation (False) or wait for jobflow
         completion (True)
+        (default: None)
+    :param wait_policy: Whether to finish the task immediately after creation (None) or:
+        - wait for the jobflow completion (WaitPolicy.WAIT_FOR_COMPLETION)
+        - wait for the jobflow completion and cluster to terminate (WaitPolicy.WAIT_FOR_STEPS_COMPLETION)
+        (default: None)
     :param waiter_max_attempts: Maximum number of tries before failing.
     :param waiter_delay: Number of seconds between polling the state of the notebook.
     :param deferrable: If True, the operator will wait asynchronously for the crawl to complete.
@@ -665,7 +677,8 @@ class EmrCreateJobFlowOperator(BaseOperator):
         emr_conn_id: str | None = "emr_default",
         job_flow_overrides: str | dict[str, Any] | None = None,
         region_name: str | None = None,
-        wait_for_completion: bool = False,
+        wait_for_completion: bool | None = None,
+        wait_policy: WaitPolicy | None = None,
         waiter_max_attempts: int | None = None,
         waiter_delay: int | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
@@ -676,11 +689,20 @@ class EmrCreateJobFlowOperator(BaseOperator):
         self.emr_conn_id = emr_conn_id
         self.job_flow_overrides = job_flow_overrides or {}
         self.region_name = region_name
-        self.wait_for_completion = wait_for_completion
+        self.wait_policy = wait_policy
         self.waiter_max_attempts = waiter_max_attempts or 60
         self.waiter_delay = waiter_delay or 60
         self.deferrable = deferrable
 
+        if wait_for_completion is not None:
+            warnings.warn(
+                "`wait_for_completion` parameter is deprecated, please use `wait_policy` instead.",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+            # preserve previous behaviour
+            self.wait_policy = WaitPolicy.WAIT_FOR_COMPLETION if wait_for_completion else None
+
     @cached_property
     def _emr_hook(self) -> EmrHook:
         """Create and return an EmrHook."""
@@ -733,8 +755,9 @@ class EmrCreateJobFlowOperator(BaseOperator):
                 # 60 seconds is added to allow the trigger to exit gracefully (i.e. yield TriggerEvent)
                 timeout=timedelta(seconds=self.waiter_max_attempts * self.waiter_delay + 60),
             )
-        if self.wait_for_completion:
-            self._emr_hook.get_waiter("job_flow_waiting").wait(
+        if self.wait_policy:
+            waiter_name = WAITER_POLICY_NAME_MAPPING[self.wait_policy]
+            self._emr_hook.get_waiter(waiter_name).wait(
                 ClusterId=self._job_flow_id,
                 WaiterConfig=prune_dict(
                     {

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.eventbridge import EventBridgeHook

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator

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

@@ -19,8 +19,9 @@ from __future__ import annotations
 
 import os
 import urllib.parse
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from botocore.exceptions import ClientError
 

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

@@ -18,8 +18,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator

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

@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from datetime import timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -21,8 +21,9 @@ from __future__ import annotations
 
 import subprocess
 import sys
+from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 import pytz
 from dateutil import parser
@@ -282,6 +283,8 @@ class S3CopyObjectOperator(BaseOperator):
                  CA cert bundle than the one used by botocore.
     :param acl_policy: String specifying the canned ACL policy for the file being
         uploaded to the S3 bucket.
+    :param meta_data_directive: Whether to `COPY` the metadata from the source object or `REPLACE` it with
+        metadata that's provided in the request.
     """
 
     template_fields: Sequence[str] = (
@@ -302,6 +305,7 @@ class S3CopyObjectOperator(BaseOperator):
         aws_conn_id: str | None = "aws_default",
         verify: str | bool | None = None,
         acl_policy: str | None = None,
+        meta_data_directive: str | None = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -314,6 +318,7 @@ class S3CopyObjectOperator(BaseOperator):
         self.aws_conn_id = aws_conn_id
         self.verify = verify
         self.acl_policy = acl_policy
+        self.meta_data_directive = meta_data_directive
 
     def execute(self, context: Context):
         s3_hook = S3Hook(aws_conn_id=self.aws_conn_id, verify=self.verify)
@@ -324,6 +329,7 @@ class S3CopyObjectOperator(BaseOperator):
             self.dest_bucket_name,
             self.source_version_id,
             self.acl_policy,
+            self.meta_data_directive,
         )
 
     def get_openlineage_facets_on_start(self):

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

@@ -19,8 +19,9 @@ from __future__ import annotations
 import datetime
 import json
 import time
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, ClassVar, Sequence
+from typing import TYPE_CHECKING, Any, Callable, ClassVar
 
 from botocore.exceptions import ClientError
 

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.sns import SnsHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator

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

@@ -18,7 +18,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator

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

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields

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

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from collections.abc import Sequence
 
 from airflow.providers.amazon.aws.utils.mixins import (
     AwsBaseHookMixin,

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

@@ -16,9 +16,10 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from datetime import timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Any, Sequence, TypeVar
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -19,7 +19,8 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor
 from airflow.providers.amazon.aws.utils.mixins import aws_template_fields

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

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable, Sequence
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.dms import DmsHook

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Iterable, Sequence
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 

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

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -16,8 +16,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.ecs import (

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

@@ -19,8 +19,9 @@
 from __future__ import annotations
 
 from abc import abstractmethod
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.eks import (

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

@@ -17,9 +17,10 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Iterable, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import (

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

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.glacier import GlacierHook

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

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,8 +17,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.glue_crawler import GlueCrawlerHook

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.lambda_function import LambdaHook

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.quicksight import QuickSightHook

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

@@ -16,8 +16,9 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowNotFoundException
 from airflow.providers.amazon.aws.hooks.rds import RdsHook

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

@@ -16,9 +16,10 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Sequence
 from datetime import timedelta
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

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

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

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.sagemaker import LogState, SageMakerHook

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

@@ -19,8 +19,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Collection, Sequence
 from datetime import timedelta
-from typing import TYPE_CHECKING, Any, Collection, Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -17,7 +17,8 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.step_function import StepFunctionHook

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

@@ -19,7 +19,8 @@ from __future__ import annotations
 
 import os
 import tempfile
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook

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

@@ -19,7 +19,7 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from collections.abc import Sequence
 
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

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

@@ -20,12 +20,13 @@ from __future__ import annotations
 
 import json
 import os
+from collections.abc import 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, Sequence
+from typing import IO, TYPE_CHECKING, Any, Callable
 from uuid import uuid4
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook