apache-airflow-providers-amazon 8.19.0rc1__py3-none-any.whl → 8.20.0rc1__py3-none-any.whl

airflow/providers/amazon/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.19.0"
+__version__ = "8.20.0"
 
 try:
     from airflow import __version__ as airflow_version

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

@@ -29,14 +29,16 @@ class AvpEntities(Enum):
     """Enum of Amazon Verified Permissions entities."""
 
     ACTION = "Action"
-    ROLE = "Role"
+    GROUP = "Group"
     USER = "User"
 
     # Resource types
     CONFIGURATION = "Configuration"
     CONNECTION = "Connection"
+    CUSTOM = "Custom"
     DAG = "Dag"
     DATASET = "Dataset"
+    MENU = "Menu"
     POOL = "Pool"
     VARIABLE = "Variable"
     VIEW = "View"
@@ -48,7 +50,7 @@ def get_entity_type(resource_type: AvpEntities) -> str:
 
     :param resource_type: Resource type.
 
-    Example: Airflow::Action, Airflow::Role, Airflow::Variable, Airflow::User.
+    Example: Airflow::Action, Airflow::Group, Airflow::Variable, Airflow::User.
     """
     return AVP_PREFIX_ENTITIES + resource_type.value
 

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

@@ -16,7 +16,9 @@
 # under the License.
 from __future__ import annotations
 
+import json
 from functools import cached_property
+from pathlib import Path
 from typing import TYPE_CHECKING, Sequence, TypedDict
 
 from airflow.configuration import conf
@@ -94,7 +96,7 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
         if user is None:
             return False
 
-        entity_list = self._get_user_role_entities(user)
+        entity_list = self._get_user_group_entities(user)
 
         self.log.debug(
             "Making authorization request for user=%s, method=%s, entity_type=%s, entity_id=%s",
@@ -144,7 +146,7 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
         :param requests: the list of requests containing the method, the entity_type and the entity ID
         :param user: the user
         """
-        entity_list = self._get_user_role_entities(user)
+        entity_list = self._get_user_group_entities(user)
 
         self.log.debug("Making batch authorization request for user=%s, requests=%s", user.get_id(), requests)
 
@@ -222,20 +224,33 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
         )
         raise AirflowException("Could not find the authorization result.")
 
+    def is_policy_store_schema_up_to_date(self) -> bool:
+        """Return whether the policy store schema equals the latest version of the schema."""
+        resp = self.avp_client.get_schema(
+            policyStoreId=self.avp_policy_store_id,
+        )
+        policy_store_schema = json.loads(resp["schema"])
+
+        schema_path = Path(__file__).parents[0] / "schema.json"
+        with open(schema_path) as schema_file:
+            latest_schema = json.load(schema_file)
+
+        return policy_store_schema == latest_schema
+
     @staticmethod
-    def _get_user_role_entities(user: AwsAuthManagerUser) -> list[dict]:
+    def _get_user_group_entities(user: AwsAuthManagerUser) -> list[dict]:
         user_entity = {
             "identifier": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()},
             "parents": [
-                {"entityType": get_entity_type(AvpEntities.ROLE), "entityId": group}
+                {"entityType": get_entity_type(AvpEntities.GROUP), "entityId": group}
                 for group in user.get_groups()
             ],
         }
-        role_entities = [
-            {"identifier": {"entityType": get_entity_type(AvpEntities.ROLE), "entityId": group}}
+        group_entities = [
+            {"identifier": {"entityType": get_entity_type(AvpEntities.GROUP), "entityId": group}}
             for group in user.get_groups()
         ]
-        return [user_entity, *role_entities]
+        return [user_entity, *group_entities]
 
     @staticmethod
     def _build_context(context: dict | None) -> dict | None:

airflow/providers/amazon/aws/auth_manager/{cli → avp}/schema.json

@@ -25,6 +25,30 @@
                     "resourceTypes": ["Connection"]
                 }
             },
+            "Custom.DELETE": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Custom"]
+                }
+            },
+            "Custom.GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Custom"]
+                }
+            },
+            "Custom.POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Custom"]
+                }
+            },
+            "Custom.PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Custom"]
+                }
+            },
             "Configuration.GET": {
                 "appliesTo": {
                     "principalTypes": ["User"],
@@ -97,6 +121,12 @@
                     "resourceTypes": ["Dataset"]
                 }
             },
+            "Menu.MENU": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Menu"]
+                }
+            },
             "Pool.DELETE": {
                 "appliesTo": {
                     "principalTypes": ["User"],
@@ -155,13 +185,15 @@
         "entityTypes": {
             "Configuration": {},
             "Connection": {},
+            "Custom": {},
             "Dag": {},
             "Dataset": {},
+            "Menu": {},
             "Pool": {},
-            "Role": {},
+            "Group": {},
             "User": {
                 "memberOfTypes": [
-                    "Role"
+                    "Group"
                 ]
             },
             "Variable": {},

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

@@ -17,14 +17,15 @@
 from __future__ import annotations
 
 import argparse
+from collections import defaultdict
 from functools import cached_property
-from typing import TYPE_CHECKING, Sequence, cast
+from typing import TYPE_CHECKING, Container, Sequence, cast
 
 from flask import session, url_for
 
 from airflow.cli.cli_config import CLICommand, DefaultHelpParser, GroupCommand
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowOptionalProviderFeatureException
+from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
 from airflow.providers.amazon.aws.auth_manager.avp.facade import (
     AwsAuthManagerAmazonVerifiedPermissionsFacade,
@@ -40,28 +41,6 @@ from airflow.providers.amazon.aws.auth_manager.constants import (
 from airflow.providers.amazon.aws.auth_manager.security_manager.aws_security_manager_override import (
     AwsSecurityManagerOverride,
 )
-from airflow.security.permissions import (
-    RESOURCE_AUDIT_LOG,
-    RESOURCE_CLUSTER_ACTIVITY,
-    RESOURCE_CONFIG,
-    RESOURCE_CONNECTION,
-    RESOURCE_DAG,
-    RESOURCE_DAG_CODE,
-    RESOURCE_DAG_DEPENDENCIES,
-    RESOURCE_DAG_RUN,
-    RESOURCE_DATASET,
-    RESOURCE_DOCS,
-    RESOURCE_JOB,
-    RESOURCE_PLUGIN,
-    RESOURCE_POOL,
-    RESOURCE_PROVIDER,
-    RESOURCE_SLA_MISS,
-    RESOURCE_TASK_INSTANCE,
-    RESOURCE_TASK_RESCHEDULE,
-    RESOURCE_TRIGGER,
-    RESOURCE_VARIABLE,
-    RESOURCE_XCOM,
-)
 
 try:
     from airflow.auth.managers.base_auth_manager import BaseAuthManager, ResourceMethod
@@ -96,136 +75,6 @@ if TYPE_CHECKING:
     from airflow.www.extensions.init_appbuilder import AirflowAppBuilder
 
 
-_MENU_ITEM_REQUESTS: dict[str, IsAuthorizedRequest] = {
-    RESOURCE_AUDIT_LOG: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.AUDIT_LOG.value,
-            },
-        },
-    },
-    RESOURCE_CLUSTER_ACTIVITY: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.CLUSTER_ACTIVITY.value,
-    },
-    RESOURCE_CONFIG: {
-        "method": "GET",
-        "entity_type": AvpEntities.CONFIGURATION,
-    },
-    RESOURCE_CONNECTION: {
-        "method": "GET",
-        "entity_type": AvpEntities.CONNECTION,
-    },
-    RESOURCE_DAG: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-    },
-    RESOURCE_DAG_CODE: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.CODE.value,
-            },
-        },
-    },
-    RESOURCE_DAG_DEPENDENCIES: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.DEPENDENCIES.value,
-            },
-        },
-    },
-    RESOURCE_DAG_RUN: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.RUN.value,
-            },
-        },
-    },
-    RESOURCE_DATASET: {
-        "method": "GET",
-        "entity_type": AvpEntities.DATASET,
-    },
-    RESOURCE_DOCS: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.DOCS.value,
-    },
-    RESOURCE_PLUGIN: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.PLUGINS.value,
-    },
-    RESOURCE_JOB: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.JOBS.value,
-    },
-    RESOURCE_POOL: {
-        "method": "GET",
-        "entity_type": AvpEntities.POOL,
-    },
-    RESOURCE_PROVIDER: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.PROVIDERS.value,
-    },
-    RESOURCE_SLA_MISS: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.SLA_MISS.value,
-            },
-        },
-    },
-    RESOURCE_TASK_INSTANCE: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.TASK_INSTANCE.value,
-            },
-        },
-    },
-    RESOURCE_TASK_RESCHEDULE: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.TASK_RESCHEDULE.value,
-            },
-        },
-    },
-    RESOURCE_TRIGGER: {
-        "method": "GET",
-        "entity_type": AvpEntities.VIEW,
-        "entity_id": AccessView.TRIGGERS.value,
-    },
-    RESOURCE_VARIABLE: {
-        "method": "GET",
-        "entity_type": AvpEntities.VARIABLE,
-    },
-    RESOURCE_XCOM: {
-        "method": "GET",
-        "entity_type": AvpEntities.DAG,
-        "context": {
-            "dag_entity": {
-                "string": DagAccessEntity.XCOM.value,
-            },
-        },
-    },
-}
-
-
 class AwsAuthManager(BaseAuthManager):
     """
     AWS auth manager.
@@ -239,6 +88,7 @@ class AwsAuthManager(BaseAuthManager):
     def __init__(self, appbuilder: AirflowAppBuilder) -> None:
         super().__init__(appbuilder)
         enable = conf.getboolean(CONF_SECTION_NAME, CONF_ENABLE_KEY)
+        self._check_avp_schema_version()
         if not enable:
             raise NotImplementedError(
                 "The AWS auth manager is currently being built. It is not finalized. It is not intended to be used yet."
@@ -356,6 +206,16 @@ class AwsAuthManager(BaseAuthManager):
             entity_id=access_view.value,
         )
 
+    def is_authorized_custom_view(
+        self, *, method: ResourceMethod, resource_name: str, user: BaseUser | None = None
+    ):
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.CUSTOM,
+            user=user or self.get_user(),
+            entity_id=resource_name,
+        )
+
     def batch_is_authorized_connection(
         self,
         requests: Sequence[IsAuthorizedConnectionRequest],
@@ -443,6 +303,60 @@ class AwsAuthManager(BaseAuthManager):
         ]
         return self.avp_facade.batch_is_authorized(requests=facade_requests, user=self.get_user())
 
+    def filter_permitted_dag_ids(
+        self,
+        *,
+        dag_ids: set[str],
+        methods: Container[ResourceMethod] | None = None,
+        user=None,
+    ):
+        """
+        Filter readable or writable DAGs for user.
+
+        :param dag_ids: the list of DAG ids
+        :param methods: whether filter readable or writable
+        :param user: the current user
+        """
+        if not methods:
+            methods = ["PUT", "GET"]
+
+        if not user:
+            user = self.get_user()
+
+        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)
+
+        batch_is_authorized_results = self.avp_facade.get_batch_is_authorized_results(
+            requests=requests_list, user=user
+        )
+
+        def _has_access_to_dag(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 {
+            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.
@@ -465,19 +379,25 @@ class AwsAuthManager(BaseAuthManager):
             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 self._has_access_to_menu_item(batch_is_authorized_results, requests[child.name], user):
+                    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 self._has_access_to_menu_item(batch_is_authorized_results, requests[menu_item.name], user):
+            elif _has_access_to_menu_item(requests[menu_item.name]):
                 accessible_items.append(menu_item)
 
         return accessible_items
@@ -504,20 +424,21 @@ class AwsAuthManager(BaseAuthManager):
         ]
 
     @staticmethod
-    def _get_menu_item_request(fab_resource_name: str) -> IsAuthorizedRequest:
-        menu_item_request = _MENU_ITEM_REQUESTS.get(fab_resource_name)
-        if menu_item_request:
-            return menu_item_request
-        else:
-            raise AirflowException(f"Unknown resource name {fab_resource_name}")
-
-    def _has_access_to_menu_item(
-        self, batch_is_authorized_results: list[dict], request: IsAuthorizedRequest, user: AwsAuthManagerUser
-    ):
-        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"
+    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(
+                "The Amazon Verified Permissions policy store schema is different from the latest version "
+                "(https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/auth_manager/avp/schema.json). "
+                "Please update it to its latest version. "
+                "See doc: https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/auth-manager/setup/amazon-verified-permissions.html#update-the-policy-store-schema."
+            )
 
 
 def get_parser() -> argparse.ArgumentParser:

airflow/providers/amazon/aws/auth_manager/cli/avp_commands.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """User sub-commands."""
+
 from __future__ import annotations
 
 import json
@@ -63,8 +64,9 @@ def init_avp(args):
         _set_schema(client, policy_store_id, args)
 
     if not args.dry_run:
-        print("Amazon Verified Permissions resources created successfully.")
-        print("Please set them in Airflow configuration under AIRFLOW__AWS_AUTH_MANAGER__<config name>.")
+        print(
+            "Please set configs below in Airflow configuration under AIRFLOW__AWS_AUTH_MANAGER__<config name>."
+        )
         print(json.dumps({"avp_policy_store_id": policy_store_id}, indent=4))
 
 
@@ -75,9 +77,6 @@ def update_schema(args):
     client = _get_client()
     _set_schema(client, args.policy_store_id, args)
 
-    if not args.dry_run:
-        print("Amazon Verified Permissions policy store schema updated successfully.")
-
 
 def _get_client():
     """Return Amazon Verified Permissions client."""
@@ -121,7 +120,7 @@ def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
 
         response = client.create_policy_store(
             validationSettings={
-                "mode": "OFF",
+                "mode": "STRICT",
             },
             description=args.policy_store_description,
         )
@@ -139,20 +138,7 @@ def _set_schema(client: BaseClient, policy_store_id: str, args) -> None:
         print(f"Dry run, not updating the schema of the policy store with ID '{policy_store_id}'.")
         return
 
-    if args.verbose:
-        log.debug("Disabling schema validation before updating schema")
-
-    response = client.update_policy_store(
-        policyStoreId=policy_store_id,
-        validationSettings={
-            "mode": "OFF",
-        },
-    )
-
-    if args.verbose:
-        log.debug("Response from update_policy_store: %s", response)
-
-    schema_path = Path(__file__).parents[0].joinpath("schema.json").resolve()
+    schema_path = Path(__file__).parents[1] / "avp" / "schema.json"
     with open(schema_path) as schema_file:
         response = client.put_schema(
             policyStoreId=policy_store_id,
@@ -164,15 +150,4 @@ def _set_schema(client: BaseClient, policy_store_id: str, args) -> None:
         if args.verbose:
             log.debug("Response from put_schema: %s", response)
 
-    if args.verbose:
-        log.debug("Enabling schema validation after updating schema")
-
-    response = client.update_policy_store(
-        policyStoreId=policy_store_id,
-        validationSettings={
-            "mode": "STRICT",
-        },
-    )
-
-    if args.verbose:
-        log.debug("Response from update_policy_store: %s", response)
+    print("Policy store schema updated.")

airflow/providers/amazon/aws/auth_manager/cli/definition.py

@@ -69,7 +69,7 @@ AWS_AUTH_MANAGER_COMMANDS = (
     ),
     ActionCommand(
         name="update-avp-schema",
-        help="Update Amazon Verified permissions policy store schema to the latest version in 'airflow/providers/amazon/aws/auth_manager/cli/schema.json'",
+        help="Update Amazon Verified permissions policy store schema to the latest version in 'airflow/providers/amazon/aws/auth_manager/avp/schema.json'",
         func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.avp_commands.update_schema"),
         args=(ARG_POLICY_STORE_ID, ARG_DRY_RUN, ARG_VERBOSE),
     ),

airflow/providers/amazon/aws/auth_manager/cli/idc_commands.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """User sub-commands."""
+
 from __future__ import annotations
 
 import logging

airflow/providers/amazon/aws/auth_manager/views/auth.py

@@ -93,7 +93,7 @@ class AwsAuthManagerAuthenticationViews(AirflowBaseView):
             user_id=attributes["id"][0],
             groups=attributes["groups"],
             username=saml_auth.get_nameid(),
-            email=attributes["email"][0],
+            email=attributes["email"][0] if "email" in attributes else None,
         )
         session["aws_user"] = user
 

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

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.