apache-airflow-providers-amazon 9.4.0__py3-none-any.whl → 9.5.0rc1__py3-none-any.whl

airflow/providers/amazon/__init__.py

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

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

@@ -20,7 +20,7 @@ from enum import Enum
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from airflow.auth.managers.base_auth_manager import ResourceMethod
+    from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
 
 AVP_PREFIX_ENTITIES = "Airflow::"
 

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

@@ -36,7 +36,7 @@ from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
-    from airflow.auth.managers.base_auth_manager import ResourceMethod
+    from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
     from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
 
 

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

@@ -18,15 +18,14 @@ from __future__ import annotations
 
 import argparse
 from collections import defaultdict
-from collections.abc import Container, Sequence
+from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
 
 from fastapi import FastAPI
-from flask import session
 
-from airflow.auth.managers.base_auth_manager import BaseAuthManager
-from airflow.auth.managers.models.resource_details import (
+from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
+from airflow.api_fastapi.auth.managers.models.resource_details import (
     AccessView,
     ConnectionDetails,
     DagAccessEntity,
@@ -49,16 +48,14 @@ from airflow.providers.amazon.aws.auth_manager.user import AwsAuthManagerUser
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
-    from flask_appbuilder.menu import MenuItem
-
-    from airflow.auth.managers.base_auth_manager import ResourceMethod
-    from airflow.auth.managers.models.batch_apis import (
+    from airflow.api_fastapi.auth.managers.base_auth_manager import ResourceMethod
+    from airflow.api_fastapi.auth.managers.models.batch_apis import (
         IsAuthorizedConnectionRequest,
         IsAuthorizedDagRequest,
         IsAuthorizedPoolRequest,
         IsAuthorizedVariableRequest,
     )
-    from airflow.auth.managers.models.resource_details import AssetDetails, ConfigurationDetails
+    from airflow.api_fastapi.auth.managers.models.resource_details import AssetDetails, ConfigurationDetails
 
 
 class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
@@ -83,14 +80,8 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         return AwsAuthManagerAmazonVerifiedPermissionsFacade()
 
     @cached_property
-    def fastapi_endpoint(self) -> str:
-        return conf.get("fastapi", "base_url")
-
-    def get_user(self) -> AwsAuthManagerUser | None:
-        return session["aws_user"] if self.is_logged_in() else None
-
-    def is_logged_in(self) -> bool:
-        return "aws_user" in session
+    def apiserver_endpoint(self) -> str:
+        return conf.get("api", "base_url")
 
     def deserialize_user(self, token: dict[str, Any]) -> AwsAuthManagerUser:
         return AwsAuthManagerUser(**token)
@@ -162,9 +153,9 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
     def is_authorized_asset(
         self, *, method: ResourceMethod, user: AwsAuthManagerUser, details: AssetDetails | None = None
     ) -> bool:
-        asset_uri = details.uri if details else None
+        asset_id = details.id if details else None
         return self.avp_facade.is_authorized(
-            method=method, entity_type=AvpEntities.ASSET, user=user, entity_id=asset_uri
+            method=method, entity_type=AvpEntities.ASSET, user=user, entity_id=asset_id
         )
 
     def is_authorized_pool(
@@ -204,7 +195,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
 
     def is_authorized_custom_view(
         self, *, method: ResourceMethod | str, resource_name: str, user: AwsAuthManagerUser
-    ):
+    ) -> bool:
         return self.avp_facade.is_authorized(
             method=method,
             entity_type=AvpEntities.CUSTOM,
@@ -292,23 +283,18 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         *,
         dag_ids: set[str],
         user: AwsAuthManagerUser,
-        methods: Container[ResourceMethod] | None = None,
+        method: ResourceMethod = "GET",
     ):
-        if not methods:
-            methods = ["PUT", "GET"]
-
         requests: dict[str, dict[ResourceMethod, IsAuthorizedRequest]] = defaultdict(dict)
         requests_list: list[IsAuthorizedRequest] = []
         for dag_id in dag_ids:
-            for method in ["GET", "PUT"]:
-                if method in methods:
-                    request: IsAuthorizedRequest = {
-                        "method": cast("ResourceMethod", method),
-                        "entity_type": AvpEntities.DAG,
-                        "entity_id": dag_id,
-                    }
-                    requests[dag_id][cast("ResourceMethod", method)] = request
-                    requests_list.append(request)
+            request: IsAuthorizedRequest = {
+                "method": method,
+                "entity_type": AvpEntities.DAG,
+                "entity_id": dag_id,
+            }
+            requests[dag_id][method] = request
+            requests_list.append(request)
 
         batch_is_authorized_results = self.avp_facade.get_batch_is_authorized_results(
             requests=requests_list, user=user
@@ -320,67 +306,10 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             )
             return result["decision"] == "ALLOW"
 
-        return {
-            dag_id
-            for dag_id in dag_ids
-            if (
-                "GET" in methods
-                and _has_access_to_dag(requests[dag_id]["GET"])
-                or "PUT" in methods
-                and _has_access_to_dag(requests[dag_id]["PUT"])
-            )
-        }
-
-    def filter_permitted_menu_items(self, menu_items: list[MenuItem]) -> list[MenuItem]:
-        """
-        Filter menu items based on user permissions.
-
-        :param menu_items: list of all menu items
-        """
-        user = self.get_user()
-        if not user:
-            return []
-
-        requests: dict[str, IsAuthorizedRequest] = {}
-        for menu_item in menu_items:
-            if menu_item.childs:
-                for child in menu_item.childs:
-                    requests[child.name] = self._get_menu_item_request(child.name)
-            else:
-                requests[menu_item.name] = self._get_menu_item_request(menu_item.name)
-
-        batch_is_authorized_results = self.avp_facade.get_batch_is_authorized_results(
-            requests=list(requests.values()), user=user
-        )
-
-        def _has_access_to_menu_item(request: IsAuthorizedRequest):
-            result = self.avp_facade.get_batch_is_authorized_single_result(
-                batch_is_authorized_results=batch_is_authorized_results, request=request, user=user
-            )
-            return result["decision"] == "ALLOW"
-
-        accessible_items = []
-        for menu_item in menu_items:
-            if menu_item.childs:
-                accessible_children = []
-                for child in menu_item.childs:
-                    if _has_access_to_menu_item(requests[child.name]):
-                        accessible_children.append(child)
-                menu_item.childs = accessible_children
-
-                # Display the menu if the user has access to at least one sub item
-                if len(accessible_children) > 0:
-                    accessible_items.append(menu_item)
-            elif _has_access_to_menu_item(requests[menu_item.name]):
-                accessible_items.append(menu_item)
-
-        return accessible_items
+        return {dag_id for dag_id in dag_ids if _has_access_to_dag(requests[dag_id][method])}
 
     def get_url_login(self, **kwargs) -> str:
-        return f"{self.fastapi_endpoint}/auth/login"
-
-    def get_url_logout(self) -> str:
-        raise NotImplementedError()
+        return f"{self.apiserver_endpoint}/auth/login"
 
     @staticmethod
     def get_cli_commands() -> list[CLICommand]:
@@ -408,14 +337,6 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
 
         return app
 
-    @staticmethod
-    def _get_menu_item_request(resource_name: str) -> IsAuthorizedRequest:
-        return {
-            "method": "MENU",
-            "entity_type": AvpEntities.MENU,
-            "entity_id": resource_name,
-        }
-
     def _check_avp_schema_version(self):
         if not self.avp_facade.is_policy_store_schema_up_to_date():
             self.log.warning(

airflow/providers/amazon/aws/auth_manager/router/login.py

@@ -79,12 +79,13 @@ def login_callback(request: Request):
         username=saml_auth.get_nameid(),
         email=attributes["email"][0] if "email" in attributes else None,
     )
-    return RedirectResponse(url=f"/webapp?token={get_auth_manager().get_jwt_token(user)}", status_code=303)
+    url = f"{conf.get('api', 'base_url')}/?token={get_auth_manager().get_jwt_token(user)}"
+    return RedirectResponse(url=url, status_code=303)
 
 
 def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
     request_data = _prepare_request(request)
-    base_url = conf.get(section="fastapi", key="base_url")
+    base_url = conf.get(section="api", key="base_url")
     settings = {
         # We want to keep this flag on in case of errors.
         # It provides an error reasons, if turned off, it does not

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

@@ -19,11 +19,14 @@ from __future__ import annotations
 from airflow.exceptions import AirflowOptionalProviderFeatureException
 
 try:
-    from airflow.auth.managers.models.base_user import BaseUser
+    from airflow.api_fastapi.auth.managers.models.base_user import BaseUser
 except ImportError:
-    raise AirflowOptionalProviderFeatureException(
-        "Failed to import BaseUser. This feature is only available in Airflow versions >= 2.8.0"
-    )
+    try:
+        from airflow.auth.managers.models.base_user import BaseUser  # type: ignore[no-redef]
+    except ImportError:
+        raise AirflowOptionalProviderFeatureException(
+            "Failed to import BaseUser. This feature is only available in Airflow versions >= 2.8.0"
+        ) from None
 
 
 class AwsAuthManagerUser(BaseUser):

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

@@ -30,6 +30,7 @@ import inspect
 import json
 import logging
 import os
+import warnings
 from copy import deepcopy
 from functools import cached_property, wraps
 from pathlib import Path
@@ -41,6 +42,7 @@ import botocore.session
 import jinja2
 import requests
 import tenacity
+from asgiref.sync import sync_to_async
 from botocore.config import Config
 from botocore.waiter import Waiter, WaiterModel
 from dateutil.tz import tzlocal
@@ -50,6 +52,7 @@ from airflow.configuration import conf
 from airflow.exceptions import (
     AirflowException,
     AirflowNotFoundException,
+    AirflowProviderDeprecationWarning,
 )
 from airflow.hooks.base import BaseHook
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
@@ -747,7 +750,29 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
     @property
     def async_conn(self):
+        """
+        [DEPRECATED] Get an aiobotocore client to use for async operations.
+
+        This property is deprecated. Accessing it in an async context will cause the event loop to block.
+        Use the async method `get_async_conn` instead.
+        """
+        warnings.warn(
+            "The property `async_conn` is deprecated. Accessing it in an async context will cause the event loop to block. "
+            "Use the async method `get_async_conn` instead.",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+
+        return self._get_async_conn()
+
+    async def get_async_conn(self):
         """Get an aiobotocore client to use for async operations."""
+        # We have to wrap the call `self.get_client_type` in another call `_get_async_conn`,
+        # because one of it's arguments `self.region_name` is a `@property` decorated function
+        # calling the cached property `self.conn_config` at the end.
+        return await sync_to_async(self._get_async_conn)()
+
+    def _get_async_conn(self):
         if not self.client_type:
             raise ValueError("client_type must be specified.")
 

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

@@ -173,7 +173,7 @@ class EC2Hook(AwsBaseHook):
         return [instance["InstanceId"] for instance in self.get_instances(filters=filters)]
 
     async def get_instance_state_async(self, instance_id: str) -> str:
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             response = await client.describe_instances(InstanceIds=[instance_id])
             return response["Reservations"][0]["Instances"][0]["State"]["Name"]
 

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

@@ -211,7 +211,7 @@ class GlueJobHook(AwsBaseHook):
 
         The async version of get_job_state.
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             job_run = await client.get_job_run(JobName=job_name, RunId=run_id)
         return job_run["JobRun"]["JobRunState"]
 
@@ -236,6 +236,9 @@ class GlueJobHook(AwsBaseHook):
         """
         log_client = self.logs_hook.get_conn()
         paginator = log_client.get_paginator("filter_log_events")
+        job_run = self.conn.get_job_run(JobName=job_name, RunId=run_id)["JobRun"]
+        # StartTime needs to be an int and is Epoch time in milliseconds
+        start_time = int(job_run["StartedOn"].timestamp() * 1000)
 
         def display_logs_from(log_group: str, continuation_token: str | None) -> str | None:
             """Mutualize iteration over the 2 different log streams glue jobs write to."""
@@ -245,6 +248,7 @@ class GlueJobHook(AwsBaseHook):
                 for response in paginator.paginate(
                     logGroupName=log_group,
                     logStreamNames=[run_id],
+                    startTime=start_time,
                     PaginationConfig={"StartingToken": continuation_token},
                 ):
                     fetched_logs.extend([event["message"] for event in response["events"]])
@@ -270,7 +274,7 @@ class GlueJobHook(AwsBaseHook):
                 self.log.info("No new log from the Glue Job in %s", log_group)
             return next_token
 
-        log_group_prefix = self.conn.get_job_run(JobName=job_name, RunId=run_id)["JobRun"]["LogGroupName"]
+        log_group_prefix = job_run["LogGroupName"]
         log_group_default = f"{log_group_prefix}/{DEFAULT_LOG_SUFFIX}"
         log_group_error = f"{log_group_prefix}/{ERROR_LOG_SUFFIX}"
         # one would think that the error log group would contain only errors, but it actually contains

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

@@ -152,7 +152,7 @@ class AwsLogsHook(AwsBaseHook):
          If the value is LastEventTime , the results are ordered by the event time. The default value is LogStreamName.
         :param count: The maximum number of items returned
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             try:
                 response: dict[str, Any] = await client.describe_log_streams(
                     logGroupName=log_group,
@@ -194,7 +194,7 @@ class AwsLogsHook(AwsBaseHook):
             else:
                 token_arg = {}
 
-            async with self.async_conn as client:
+            async with await self.get_async_conn() as client:
                 response = await client.get_log_events(
                     logGroupName=log_group,
                     logStreamName=log_stream_name,

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

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import requests
 from botocore.exceptions import ClientError
 
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
@@ -29,6 +30,12 @@ class MwaaHook(AwsBaseHook):
 
     Provide thin wrapper around :external+boto3:py:class:`boto3.client("mwaa") <MWAA.Client>`
 
+    If your IAM policy doesn't have `airflow:InvokeRestApi` permission, the hook will use a fallback method
+    that uses the AWS credential to generate a local web login token for the Airflow Web UI and then directly
+    make requests to the Airflow API. This fallback method can be set as the default (and only) method used by
+    setting `generate_local_token` to True.  Learn more here:
+    https://docs.aws.amazon.com/mwaa/latest/userguide/access-mwaa-apache-airflow-rest-api.html#granting-access-MWAA-Enhanced-REST-API
+
     Additional arguments (such as ``aws_conn_id``) may be specified and
     are passed down to the underlying AwsBaseHook.
 
@@ -47,6 +54,7 @@ class MwaaHook(AwsBaseHook):
         method: str,
         body: dict | None = None,
         query_params: dict | None = None,
+        generate_local_token: bool = False,
     ) -> dict:
         """
         Invoke the REST API on the Airflow webserver with the specified inputs.
@@ -56,30 +64,86 @@ class MwaaHook(AwsBaseHook):
 
         :param env_name: name of the MWAA environment
         :param path: Apache Airflow REST API endpoint path to be called
-        :param method: HTTP method used for making Airflow REST API calls
+        :param method: HTTP method used for making Airflow REST API calls: 'GET'|'PUT'|'POST'|'PATCH'|'DELETE'
         :param body: Request body for the Apache Airflow REST API call
         :param query_params: Query parameters to be included in the Apache Airflow REST API call
+        :param generate_local_token: If True, only the local web token method is used without trying boto's
+            `invoke_rest_api` first. If False, the local web token method is used as a fallback after trying
+            boto's `invoke_rest_api`
         """
-        body = body or {}
+        # Filter out keys with None values because Airflow REST API doesn't accept requests otherwise
+        body = {k: v for k, v in body.items() if v is not None} if body else {}
+        query_params = query_params or {}
         api_kwargs = {
             "Name": env_name,
             "Path": path,
             "Method": method,
-            # Filter out keys with None values because Airflow REST API doesn't accept requests otherwise
-            "Body": {k: v for k, v in body.items() if v is not None},
-            "QueryParameters": query_params if query_params else {},
+            "Body": body,
+            "QueryParameters": query_params,
         }
+
+        if generate_local_token:
+            return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+
         try:
-            result = self.conn.invoke_rest_api(**api_kwargs)
+            response = self.conn.invoke_rest_api(**api_kwargs)
             # ResponseMetadata is removed because it contains data that is either very unlikely to be useful
             # in XComs and logs, or redundant given the data already included in the response
-            result.pop("ResponseMetadata", None)
-            return result
+            response.pop("ResponseMetadata", None)
+            return response
+
         except ClientError as e:
-            to_log = e.response
-            # ResponseMetadata and Error are removed because they contain data that is either very unlikely to
-            # be useful in XComs and logs, or redundant given the data already included in the response
-            to_log.pop("ResponseMetadata", None)
-            to_log.pop("Error", None)
-            self.log.error(to_log)
-            raise e
+            if (
+                e.response["Error"]["Code"] == "AccessDeniedException"
+                and "Airflow role" in e.response["Error"]["Message"]
+            ):
+                self.log.info(
+                    "Access Denied due to missing airflow:InvokeRestApi in IAM policy. Trying again by generating local token..."
+                )
+                return self._invoke_rest_api_using_local_session_token(**api_kwargs)
+            else:
+                to_log = e.response
+                # ResponseMetadata is removed because it contains data that is either very unlikely to be
+                # useful in XComs and logs, or redundant given the data already included in the response
+                to_log.pop("ResponseMetadata", None)
+                self.log.error(to_log)
+                raise
+
+    def _invoke_rest_api_using_local_session_token(
+        self,
+        **api_kwargs,
+    ) -> dict:
+        try:
+            session, hostname = self._get_session_conn(api_kwargs["Name"])
+
+            response = session.request(
+                method=api_kwargs["Method"],
+                url=f"https://{hostname}/api/v1{api_kwargs['Path']}",
+                params=api_kwargs["QueryParameters"],
+                json=api_kwargs["Body"],
+                timeout=10,
+            )
+            response.raise_for_status()
+
+        except requests.HTTPError as e:
+            self.log.error(e.response.json())
+            raise
+
+        return {
+            "RestApiStatusCode": response.status_code,
+            "RestApiResponse": response.json(),
+        }
+
+    # Based on: https://docs.aws.amazon.com/mwaa/latest/userguide/access-mwaa-apache-airflow-rest-api.html#create-web-server-session-token
+    def _get_session_conn(self, env_name: str) -> tuple:
+        create_token_response = self.conn.create_web_login_token(Name=env_name)
+        web_server_hostname = create_token_response["WebServerHostname"]
+        web_token = create_token_response["WebToken"]
+
+        login_url = f"https://{web_server_hostname}/aws_mwaa/login"
+        login_payload = {"token": web_token}
+        session = requests.Session()
+        login_response = session.post(login_url, data=login_payload, timeout=10)
+        login_response.raise_for_status()
+
+        return session, web_server_hostname

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

@@ -93,7 +93,7 @@ class RedshiftHook(AwsBaseHook):
             return "cluster_not_found"
 
     async def cluster_status_async(self, cluster_identifier: str) -> str:
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             response = await client.describe_clusters(ClusterIdentifier=cluster_identifier)
             return response["Clusters"][0]["ClusterStatus"] if response else None
 

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

@@ -275,7 +275,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         :param statement_id: the UUID of the statement
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             desc = await client.describe_statement(Id=statement_id)
             return desc["Status"] in RUNNING_STATES
 
@@ -288,6 +288,6 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
 
         :param statement_id: the UUID of the statement
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             resp = await client.describe_statement(Id=statement_id)
             return self.parse_statement_response(resp)

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

@@ -1318,7 +1318,7 @@ class SageMakerHook(AwsBaseHook):
 
         :param job_name: the name of the training job
         """
-        async with self.async_conn as client:
+        async with await self.get_async_conn() as client:
             response: dict[str, Any] = await client.describe_training_job(TrainingJobName=job_name)
             return response