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

airflow/providers/amazon/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "9.1.0"
+__version__ = "9.2.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.8.0"
+    "2.9.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.8.0+"
+        f"The package `apache-airflow-providers-amazon:{__version__}` needs Apache Airflow 2.9.0+"
     )

airflow/providers/amazon/aws/auth_manager/avp/facade.py

@@ -17,9 +17,10 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Sequence, TypedDict
+from typing import TYPE_CHECKING, TypedDict
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -19,8 +19,9 @@ from __future__ import annotations
 import argparse
 import warnings
 from collections import defaultdict
+from collections.abc import Container, Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Container, Sequence, cast
+from typing import TYPE_CHECKING, cast
 
 from flask import session, url_for
 
@@ -80,16 +81,6 @@ class AwsAuthManager(BaseAuthManager):
     """
 
     def __init__(self, appbuilder: AirflowAppBuilder) -> None:
-        from packaging.version import Version
-
-        from airflow.version import version
-
-        # TODO: remove this if block when min_airflow_version is set to higher than 2.9.0
-        if Version(version) < Version("2.9"):
-            raise AirflowOptionalProviderFeatureException(
-                "``AwsAuthManager`` is compatible with Airflow versions >= 2.9."
-            )
-
         super().__init__(appbuilder)
         self._check_avp_schema_version()
 
@@ -430,7 +421,8 @@ class AwsAuthManager(BaseAuthManager):
         ]
 
     def register_views(self) -> None:
-        self.appbuilder.add_view_no_menu(AwsAuthManagerAuthenticationViews())
+        if self.appbuilder:
+            self.appbuilder.add_view_no_menu(AwsAuthManagerAuthenticationViews())
 
     @staticmethod
     def _get_menu_item_request(resource_name: str) -> IsAuthorizedRequest:

airflow/providers/amazon/aws/executors/batch/batch_executor.py

@@ -21,9 +21,10 @@ from __future__ import annotations
 
 import time
 from collections import deque
+from collections.abc import Sequence
 from contextlib import suppress
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict, List, Sequence
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError, NoCredentialsError
 
@@ -55,8 +56,8 @@ from airflow.providers.amazon.aws.executors.batch.utils import (
 )
 from airflow.utils.state import State
 
-CommandType = List[str]
-ExecutorConfigType = Dict[str, Any]
+CommandType = list[str]
+ExecutorConfigType = dict[str, Any]
 
 INVALID_CREDENTIALS_EXCEPTIONS = [
     "ExpiredTokenException",

airflow/providers/amazon/aws/executors/batch/utils.py

@@ -19,7 +19,7 @@ from __future__ import annotations
 import datetime
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import TYPE_CHECKING, Any
 
 from airflow.providers.amazon.aws.executors.utils.base_config_keys import BaseConfigKeys
 from airflow.utils.state import State
@@ -27,8 +27,8 @@ from airflow.utils.state import State
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstanceKey
 
-CommandType = List[str]
-ExecutorConfigType = Dict[str, Any]
+CommandType = list[str]
+ExecutorConfigType = dict[str, Any]
 
 CONFIG_GROUP_NAME = "aws_batch_executor"
 

airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -25,9 +25,10 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict, deque
+from collections.abc import Sequence
 from contextlib import suppress
 from copy import deepcopy
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING
 
 from botocore.exceptions import ClientError, NoCredentialsError
 

airflow/providers/amazon/aws/executors/ecs/utils.py

@@ -26,7 +26,7 @@ from __future__ import annotations
 import datetime
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Callable, Dict, List
+from typing import TYPE_CHECKING, Any, Callable
 
 from inflection import camelize
 
@@ -36,9 +36,9 @@ from airflow.utils.state import State
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstanceKey
 
-CommandType = List[str]
+CommandType = list[str]
 ExecutorConfigFunctionType = Callable[[CommandType], dict]
-ExecutorConfigType = Dict[str, Any]
+ExecutorConfigType = dict[str, Any]
 
 ECS_LAUNCH_TYPE_EC2 = "EC2"
 ECS_LAUNCH_TYPE_FARGATE = "FARGATE"

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

@@ -19,7 +19,7 @@ from __future__ import annotations
 import asyncio
 import logging
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, Dict
+from typing import TYPE_CHECKING, Any, Callable
 
 import requests
 from botocore import UNSIGNED
@@ -32,7 +32,7 @@ if TYPE_CHECKING:
     from fsspec import AbstractFileSystem
 
 
-Properties = Dict[str, str]
+Properties = dict[str, str]
 
 S3_PROXY_URI = "proxy-uri"
 

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

@@ -16,7 +16,15 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, cast
+
+from mypy_boto3_appflow.type_defs import (
+    DestinationFlowConfigTypeDef,
+    SourceFlowConfigTypeDef,
+    TaskTypeDef,
+    TriggerConfigTypeDef,
+)
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
 from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
@@ -117,9 +125,11 @@ class AppflowHook(AwsGenericHook["AppflowClient"]):
 
         self.conn.update_flow(
             flowName=response["flowName"],
-            destinationFlowConfigList=response["destinationFlowConfigList"],
-            sourceFlowConfig=response["sourceFlowConfig"],
-            triggerConfig=response["triggerConfig"],
+            destinationFlowConfigList=cast(
+                Sequence[DestinationFlowConfigTypeDef], response["destinationFlowConfigList"]
+            ),
+            sourceFlowConfig=cast(SourceFlowConfigTypeDef, response["sourceFlowConfig"]),
+            triggerConfig=cast(TriggerConfigTypeDef, response["triggerConfig"]),
             description=response.get("description", "Flow description."),
-            tasks=tasks,
+            tasks=cast(Sequence[TaskTypeDef], tasks),
         )

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

@@ -25,7 +25,8 @@ This module contains AWS Athena hook.
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Collection
+from collections.abc import Collection
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook

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

@@ -18,7 +18,12 @@
 from __future__ import annotations
 
 import json
+from datetime import datetime
 from enum import Enum
+from typing import Any
+
+from botocore.exceptions import ClientError
+from dateutil import parser
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
@@ -219,3 +224,159 @@ class DmsHook(AwsBaseHook):
             ],
             WithoutSettings=True,
         )
+
+    def describe_replication_configs(self, filters: list[dict] | None = None, **kwargs) -> list[dict]:
+        """
+        Return list of serverless replication configs.
+
+        .. seealso::
+            - :external+boto3:py:meth:`DatabaseMigrationService.Client.describe_replication_configs`
+
+        :param filters: List of filter objects
+        :return: List of replication tasks
+        """
+        filters = filters if filters is not None else []
+
+        try:
+            resp = self.conn.describe_replication_configs(Filters=filters, **kwargs)
+            return resp.get("ReplicationConfigs", [])
+        except Exception as ex:
+            self.log.error("Error while describing replication configs: %s", str(ex))
+            raise ex
+
+    def create_replication_config(
+        self,
+        replication_config_id: str,
+        source_endpoint_arn: str,
+        target_endpoint_arn: str,
+        compute_config: dict[str, Any],
+        replication_type: str,
+        table_mappings: str,
+        additional_config_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        """
+        Create an AWS DMS Serverless configuration that can be used to start an DMS Serverless replication.
+
+        .. seealso::
+            - :external+boto3:py:meth:`DatabaseMigrationService.Client.create_replication_config`
+
+        :param replicationConfigId: Unique identifier used to create a ReplicationConfigArn.
+        :param sourceEndpointArn: ARN of the source endpoint
+        :param targetEndpointArn: ARN of the target endpoint
+        :param computeConfig: Parameters for provisioning an DMS Serverless replication.
+        :param replicationType: type of DMS Serverless replication
+        :param tableMappings: JSON table mappings
+        :param tags: Key-value tag pairs
+        :param resourceId: Unique value or name that you set for a given resource that can be used to construct an Amazon Resource Name (ARN) for that resource.
+        :param supplementalSettings: JSON settings for specifying supplemental data
+        :param replicationSettings: JSON settings for DMS Serverless replications
+
+        :return: ReplicationConfigArn
+
+        """
+        if additional_config_kwargs is None:
+            additional_config_kwargs = {}
+        try:
+            resp = self.conn.create_replication_config(
+                ReplicationConfigIdentifier=replication_config_id,
+                SourceEndpointArn=source_endpoint_arn,
+                TargetEndpointArn=target_endpoint_arn,
+                ComputeConfig=compute_config,
+                ReplicationType=replication_type,
+                TableMappings=table_mappings,
+                **additional_config_kwargs,
+            )
+            arn = resp.get("ReplicationConfig", {}).get("ReplicationConfigArn")
+            self.log.info("Successfully created replication config: %s", arn)
+            return arn
+
+        except ClientError as err:
+            err_str = f"Error: {err.get('Error','').get('Code','')}: {err.get('Error','').get('Message','')}"
+            self.log.error("Error while creating replication config: %s", err_str)
+            raise err
+
+    def describe_replications(self, filters: list[dict[str, Any]] | None = None, **kwargs) -> list[dict]:
+        """
+        Return list of serverless replications.
+
+        .. seealso::
+            - :external+boto3:py:meth:`DatabaseMigrationService.Client.describe_replications`
+
+        :param filters: List of filter objects
+        :return: List of replications
+        """
+        filters = filters if filters is not None else []
+        try:
+            resp = self.conn.describe_replications(Filters=filters, **kwargs)
+            return resp.get("Replications", [])
+        except Exception as ex:
+            self.log.error("Error while describing replications: %s", str(ex))
+            raise ex
+
+    def delete_replication_config(
+        self, replication_config_arn: str, delay: int = 60, max_attempts: int = 120
+    ):
+        """
+        Delete an AWS DMS Serverless configuration.
+
+        .. seealso::
+            - :external+boto3:py:meth:`DatabaseMigrationService.Client.delete_replication_config`
+
+        :param replication_config_arn: ReplicationConfigArn
+        """
+        try:
+            self.log.info("Deleting replication config: %s", replication_config_arn)
+
+            self.conn.delete_replication_config(ReplicationConfigArn=replication_config_arn)
+
+        except ClientError as err:
+            err_str = (
+                f"Error: {err.get('Error', '').get('Code', '')}: {err.get('Error', '').get('Message', '')}"
+            )
+            self.log.error("Error while deleting replication config: %s", err_str)
+            raise err
+
+    def start_replication(
+        self,
+        replication_config_arn: str,
+        start_replication_type: str,
+        cdc_start_time: datetime | str | None = None,
+        cdc_start_pos: str | None = None,
+        cdc_stop_pos: str | None = None,
+    ):
+        additional_args: dict[str, Any] = {}
+
+        if cdc_start_time:
+            additional_args["CdcStartTime"] = (
+                cdc_start_time if isinstance(cdc_start_time, datetime) else parser.parse(cdc_start_time)
+            )
+        if cdc_start_pos:
+            additional_args["CdcStartPosition"] = cdc_start_pos
+        if cdc_stop_pos:
+            additional_args["CdcStopPosition"] = cdc_stop_pos
+
+        try:
+            resp = self.conn.start_replication(
+                ReplicationConfigArn=replication_config_arn,
+                StartReplicationType=start_replication_type,
+                **additional_args,
+            )
+
+            return resp
+        except Exception as ex:
+            self.log.error("Error while starting replication: %s", str(ex))
+            raise ex
+
+    def stop_replication(self, replication_config_arn: str):
+        resp = self.conn.stop_replication(ReplicationConfigArn=replication_config_arn)
+        return resp
+
+    def get_provision_status(self, replication_config_arn: str) -> str:
+        """Get the provisioning status for a serverless replication."""
+        result = self.describe_replications(
+            filters=[{"Name": "replication-config-arn", "Values": [replication_config_arn]}]
+        )
+
+        provision_status = result[0].get("ProvisionData", {}).get("ProvisionState", "")
+        return provision_status

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

@@ -19,8 +19,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Iterable
 from functools import cached_property
-from typing import TYPE_CHECKING, Iterable
+from typing import TYPE_CHECKING
 
 from botocore.exceptions import ClientError
 

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

@@ -22,10 +22,11 @@ import base64
 import json
 import sys
 import tempfile
+from collections.abc import Generator
 from contextlib import contextmanager
 from enum import Enum
 from functools import partial
-from typing import Callable, Generator
+from typing import Callable
 
 from botocore.exceptions import ClientError
 from botocore.signers import RequestSigner

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

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

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

@@ -18,7 +18,8 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, AsyncGenerator, Generator
+from collections.abc import AsyncGenerator, Generator
+from typing import Any
 
 from botocore.exceptions import ClientError
 

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

@@ -16,7 +16,8 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 
@@ -93,8 +94,8 @@ class RedshiftHook(AwsBaseHook):
 
     async def cluster_status_async(self, cluster_identifier: str) -> str:
         async with self.async_conn as client:
-            response = await client.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
-            return response[0]["ClusterStatus"] if response else None
+            response = await client.describe_clusters(ClusterIdentifier=cluster_identifier)
+            return response["Clusters"][0]["ClusterStatus"] if response else None
 
     def delete_cluster(
         self,

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

@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 import time
+from collections.abc import Iterable
 from dataclasses import dataclass
 from pprint import pformat
-from typing import TYPE_CHECKING, Any, Iterable
+from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
 from pendulum import duration

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

@@ -20,19 +20,15 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 
 import redshift_connector
-from packaging.version import Version
 from redshift_connector import Connection as RedshiftConnection
 from sqlalchemy import create_engine
 from sqlalchemy.engine.url import URL
 
-from airflow import __version__ as AIRFLOW_VERSION
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.version_compat import AIRFLOW_V_2_10_PLUS
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 
-_IS_AIRFLOW_2_10_OR_HIGHER = Version(Version(AIRFLOW_VERSION).base_version) >= Version("2.10.0")
-
-
 if TYPE_CHECKING:
     from airflow.models.connection import Connection
     from airflow.providers.openlineage.sqlparser import DatabaseInfo
@@ -265,6 +261,6 @@ class RedshiftSQLHook(DbApiHook):
 
     def get_openlineage_default_schema(self) -> str | None:
         """Return current schema. This is usually changed with ``SEARCH_PATH`` parameter."""
-        if _IS_AIRFLOW_2_10_OR_HIGHER:
+        if AIRFLOW_V_2_10_PLUS:
             return self.get_first("SELECT CURRENT_SCHEMA();")[0]
         return super().get_openlineage_default_schema()

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

@@ -28,6 +28,7 @@ import os
 import re
 import shutil
 import time
+from collections.abc import AsyncIterator
 from contextlib import suppress
 from copy import deepcopy
 from datetime import datetime
@@ -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, AsyncIterator, Callable
+from typing import TYPE_CHECKING, Any, Callable
 from urllib.parse import urlsplit
 from uuid import uuid4
 
@@ -1297,6 +1298,7 @@ class S3Hook(AwsBaseHook):
         dest_bucket_name: str | None = None,
         source_version_id: str | None = None,
         acl_policy: str | None = None,
+        meta_data_directive: str | None = None,
         **kwargs,
     ) -> None:
         """
@@ -1326,10 +1328,14 @@ class S3Hook(AwsBaseHook):
         :param source_version_id: Version ID of the source object (OPTIONAL)
         :param acl_policy: The string to specify the canned ACL policy for the
             object to be copied which is private by default.
+        :param meta_data_directive: Whether to `COPY` the metadata from the source object or `REPLACE` it
+            with metadata that's provided in the request.
         """
         acl_policy = acl_policy or "private"
         if acl_policy != NO_ACL:
             kwargs["ACL"] = acl_policy
+        if meta_data_directive:
+            kwargs["MetadataDirective"] = meta_data_directive
 
         dest_bucket_name, dest_bucket_key = self.get_s3_bucket_key(
             dest_bucket_name, dest_bucket_key, "dest_bucket_name", "dest_bucket_key"

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

@@ -23,9 +23,10 @@ import tarfile
 import tempfile
 import time
 from collections import Counter, namedtuple
+from collections.abc import AsyncGenerator, Generator
 from datetime import datetime
 from functools import partial
-from typing import Any, AsyncGenerator, Callable, Generator, cast
+from typing import Any, Callable, cast
 
 from asgiref.sync import sync_to_async
 from botocore.exceptions import ClientError

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

@@ -18,7 +18,8 @@
 
 from __future__ import annotations
 
-from typing import Any, Iterable
+from collections.abc import Iterable
+from typing import Any
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.utils.email import build_mime_message

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

@@ -17,8 +17,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import Sequence
 
 from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sns import SnsHook

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

@@ -17,8 +17,8 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from functools import cached_property
-from typing import Sequence
 
 from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook

airflow/providers/amazon/aws/operators/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 urllib.parse import urlparse
 
 from airflow.configuration import conf

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

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

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

@@ -26,9 +26,10 @@ AWS Batch services.
 
 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/operators/bedrock.py

@@ -17,8 +17,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Sequence
 from time import sleep
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any
 
 from botocore.exceptions import ClientError
 

airflow/providers/amazon/aws/operators/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.hooks.cloud_formation import CloudFormationHook
 from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator

airflow/providers/amazon/aws/operators/comprehend.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, Any, ClassVar, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException

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

@@ -20,7 +20,8 @@ from __future__ import annotations
 
 import logging
 import random
-from typing import TYPE_CHECKING, Any, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException, AirflowTaskTimeout
 from airflow.providers.amazon.aws.hooks.datasync import DataSyncHook