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

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

@@ -34,6 +34,8 @@ class AvpEntities(Enum):
 
     # Resource types
     ASSET = "Asset"
+    ASSET_ALIAS = "AssetAlias"
+    BACKFILL = "Backfills"
     CONFIGURATION = "Configuration"
     CONNECTION = "Connection"
     CUSTOM = "Custom"

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

@@ -21,12 +21,15 @@ from collections import defaultdict
 from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, cast
+from urllib.parse import urljoin
 
 from fastapi import FastAPI
 
+from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
 from airflow.api_fastapi.auth.managers.models.resource_details import (
     AccessView,
+    BackfillDetails,
     ConnectionDetails,
     DagAccessEntity,
     DagDetails,
@@ -55,7 +58,12 @@ if TYPE_CHECKING:
         IsAuthorizedPoolRequest,
         IsAuthorizedVariableRequest,
     )
-    from airflow.api_fastapi.auth.managers.models.resource_details import AssetDetails, ConfigurationDetails
+    from airflow.api_fastapi.auth.managers.models.resource_details import (
+        AssetAliasDetails,
+        AssetDetails,
+        ConfigurationDetails,
+    )
+    from airflow.api_fastapi.common.types import MenuItem
 
 
 class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
@@ -84,11 +92,11 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         return conf.get("api", "base_url")
 
     def deserialize_user(self, token: dict[str, Any]) -> AwsAuthManagerUser:
-        return AwsAuthManagerUser(**token)
+        return AwsAuthManagerUser(user_id=token.pop("sub"), **token)
 
     def serialize_user(self, user: AwsAuthManagerUser) -> dict[str, Any]:
         return {
-            "user_id": user.get_id(),
+            "sub": user.get_id(),
             "groups": user.get_groups(),
             "username": user.username,
             "email": user.email,
@@ -150,6 +158,14 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             context=context,
         )
 
+    def is_authorized_backfill(
+        self, *, method: ResourceMethod, user: AwsAuthManagerUser, details: BackfillDetails | None = None
+    ) -> bool:
+        backfill_id = details.id if details else None
+        return self.avp_facade.is_authorized(
+            method=method, entity_type=AvpEntities.BACKFILL, user=user, entity_id=backfill_id
+        )
+
     def is_authorized_asset(
         self, *, method: ResourceMethod, user: AwsAuthManagerUser, details: AssetDetails | None = None
     ) -> bool:
@@ -158,6 +174,14 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             method=method, entity_type=AvpEntities.ASSET, user=user, entity_id=asset_id
         )
 
+    def is_authorized_asset_alias(
+        self, *, method: ResourceMethod, user: AwsAuthManagerUser, details: AssetAliasDetails | None = None
+    ) -> bool:
+        asset_alias_id = details.id if details else None
+        return self.avp_facade.is_authorized(
+            method=method, entity_type=AvpEntities.ASSET_ALIAS, user=user, entity_id=asset_alias_id
+        )
+
     def is_authorized_pool(
         self, *, method: ResourceMethod, user: AwsAuthManagerUser, details: PoolDetails | None = None
     ) -> bool:
@@ -203,6 +227,25 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             entity_id=resource_name,
         )
 
+    def filter_authorized_menu_items(
+        self, menu_items: list[MenuItem], *, user: AwsAuthManagerUser
+    ) -> list[MenuItem]:
+        requests: dict[str, IsAuthorizedRequest] = {}
+        for menu_item in menu_items:
+            requests[menu_item.value] = self._get_menu_item_request(menu_item.value)
+
+        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"
+
+        return [menu_item for menu_item in menu_items if _has_access_to_menu_item(requests[menu_item.value])]
+
     def batch_is_authorized_connection(
         self,
         requests: Sequence[IsAuthorizedConnectionRequest],
@@ -278,7 +321,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         ]
         return self.avp_facade.batch_is_authorized(requests=facade_requests, user=user)
 
-    def filter_permitted_dag_ids(
+    def filter_authorized_dag_ids(
         self,
         *,
         dag_ids: set[str],
@@ -309,7 +352,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
         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.apiserver_endpoint}/auth/login"
+        return urljoin(self.apiserver_endpoint, f"{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login")
 
     @staticmethod
     def get_cli_commands() -> list[CLICommand]:
@@ -337,6 +380,14 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
 
         return app
 
+    @staticmethod
+    def _get_menu_item_request(menu_item_text: str) -> IsAuthorizedRequest:
+        return {
+            "method": "MENU",
+            "entity_type": AvpEntities.MENU,
+            "entity_id": menu_item_text,
+        }
+
     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

@@ -25,7 +25,8 @@ from fastapi import HTTPException, Request
 from starlette import status
 from starlette.responses import RedirectResponse
 
-from airflow.api_fastapi.app import get_auth_manager
+from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX, get_auth_manager
+from airflow.api_fastapi.auth.managers.base_auth_manager import COOKIE_NAME_JWT_TOKEN
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.auth_manager.constants import CONF_SAML_METADATA_URL_KEY, CONF_SECTION_NAME
@@ -79,8 +80,11 @@ def login_callback(request: Request):
         username=saml_auth.get_nameid(),
         email=attributes["email"][0] if "email" in attributes else None,
     )
-    url = f"{conf.get('api', 'base_url')}/?token={get_auth_manager().get_jwt_token(user)}"
-    return RedirectResponse(url=url, status_code=303)
+    url = conf.get("api", "base_url")
+    token = get_auth_manager().generate_jwt(user)
+    response = RedirectResponse(url=url, status_code=303)
+    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True)
+    return response
 
 
 def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
@@ -93,7 +97,7 @@ def _init_saml_auth(request: Request) -> OneLogin_Saml2_Auth:
         "sp": {
             "entityId": "aws-auth-manager-saml-client",
             "assertionConsumerService": {
-                "url": f"{base_url}/auth/login_callback",
+                "url": f"{base_url}{AUTH_MANAGER_FASTAPI_APP_PREFIX}/login_callback",
                 "binding": "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST",
             },
         },

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

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

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

@@ -943,6 +943,7 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         self,
         waiter_name: str,
         parameters: dict[str, str] | None = None,
+        config_overrides: dict[str, Any] | None = None,
         deferrable: bool = False,
         client=None,
     ) -> Waiter:
@@ -962,6 +963,9 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
         :param parameters: will scan the waiter config for the keys of that dict,
             and replace them with the corresponding value. If a custom waiter has
             such keys to be expanded, they need to be provided here.
+            Note: cannot be used if parameters are included in config_overrides
+        :param config_overrides: will update values of provided keys in the waiter's
+            config. Only specified keys will be updated.
         :param deferrable: If True, the waiter is going to be an async custom waiter.
             An async client must be provided in that case.
         :param client: The client to use for the waiter's operations
@@ -970,14 +974,18 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
 
         if deferrable and not client:
             raise ValueError("client must be provided for a deferrable waiter.")
+        if parameters is not None and config_overrides is not None and "acceptors" in config_overrides:
+            raise ValueError('parameters must be None when "acceptors" is included in config_overrides')
         # Currently, the custom waiter doesn't work with resource_type, only client_type is supported.
         client = client or self._client
         if self.waiter_path and (waiter_name in self._list_custom_waiters()):
             # Technically if waiter_name is in custom_waiters then self.waiter_path must
             # exist but MyPy doesn't like the fact that self.waiter_path could be None.
             with open(self.waiter_path) as config_file:
-                config = json.loads(config_file.read())
+                config: dict = json.loads(config_file.read())
 
+            if config_overrides is not None:
+                config["waiters"][waiter_name].update(config_overrides)
             config = self._apply_parameters_value(config, waiter_name, parameters)
             return BaseBotoWaiter(client=client, model_config=config, deferrable=deferrable).waiter(
                 waiter_name

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

@@ -35,7 +35,6 @@ from botocore.signers import RequestSigner
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.sts import StsHook
 from airflow.utils import yaml
-from airflow.utils.json import AirflowJsonEncoder
 
 DEFAULT_PAGINATION_TOKEN = ""
 STS_TOKEN_EXPIRES_IN = 60
@@ -315,7 +314,7 @@ class EksHook(AwsBaseHook):
         )
         if verbose:
             cluster_data = response.get("cluster")
-            self.log.info("Amazon EKS cluster details: %s", json.dumps(cluster_data, cls=AirflowJsonEncoder))
+            self.log.info("Amazon EKS cluster details: %s", json.dumps(cluster_data, default=repr))
         return response
 
     def describe_nodegroup(self, clusterName: str, nodegroupName: str, verbose: bool = False) -> dict:
@@ -343,7 +342,7 @@ class EksHook(AwsBaseHook):
             nodegroup_data = response.get("nodegroup")
             self.log.info(
                 "Amazon EKS managed node group details: %s",
-                json.dumps(nodegroup_data, cls=AirflowJsonEncoder),
+                json.dumps(nodegroup_data, default=repr),
             )
         return response
 
@@ -374,9 +373,7 @@ class EksHook(AwsBaseHook):
         )
         if verbose:
             fargate_profile_data = response.get("fargateProfile")
-            self.log.info(
-                "AWS Fargate profile details: %s", json.dumps(fargate_profile_data, cls=AirflowJsonEncoder)
-            )
+            self.log.info("AWS Fargate profile details: %s", json.dumps(fargate_profile_data, default=repr))
         return response
 
     def get_cluster_state(self, clusterName: str) -> ClusterStates:

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

@@ -1494,7 +1494,9 @@ class S3Hook(AwsBaseHook):
             get_hook_lineage_collector().add_output_asset(
                 context=self,
                 scheme="file",
-                asset_kwargs={"path": file_path if file_path.is_absolute() else file_path.absolute()},
+                asset_kwargs={
+                    "path": str(file_path) if file_path.is_absolute() else str(file_path.absolute())
+                },
             )
             file = open(file_path, "wb")
         else:

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

@@ -19,7 +19,6 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.models import XCom
 from airflow.providers.amazon.aws.utils.suppress import return_on_error
 from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
 
@@ -30,7 +29,9 @@ if TYPE_CHECKING:
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk import BaseOperatorLink
+    from airflow.sdk.execution_time.xcom import XCom
 else:
+    from airflow.models import XCom  # type: ignore[no-redef]
     from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
 
 

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

@@ -21,12 +21,11 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 
 from airflow.providers.amazon.aws.hooks.chime import ChimeWebhookHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
 
-from airflow.notifications.basenotifier import BaseNotifier
-
 
 class ChimeNotifier(BaseNotifier):
     """

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

@@ -20,8 +20,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from functools import cached_property
 
-from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sns import SnsHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 
 class SnsNotifier(BaseNotifier):

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

@@ -20,8 +20,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from functools import cached_property
 
-from airflow.notifications.basenotifier import BaseNotifier
 from airflow.providers.amazon.aws.hooks.sqs import SqsHook
+from airflow.providers.common.compat.notifier import BaseNotifier
 
 
 class SqsNotifier(BaseNotifier):