apache-airflow-providers-amazon 8.16.0__py3-none-any.whl → 8.17.0__py3-none-any.whl

airflow/providers/amazon/__init__.py

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

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

@@ -35,6 +35,7 @@ class AvpEntities(Enum):
     # Resource types
     CONFIGURATION = "Configuration"
     CONNECTION = "Connection"
+    DAG = "Dag"
     DATASET = "Dataset"
     POOL = "Pool"
     VARIABLE = "Variable"

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

@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 from functools import cached_property
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
@@ -25,9 +25,11 @@ from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities,
 from airflow.providers.amazon.aws.auth_manager.constants import (
     CONF_AVP_POLICY_STORE_ID_KEY,
     CONF_CONN_ID_KEY,
+    CONF_REGION_NAME_KEY,
     CONF_SECTION_NAME,
 )
 from airflow.providers.amazon.aws.hooks.verified_permissions import VerifiedPermissionsHook
+from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
@@ -46,7 +48,8 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
     def avp_client(self):
         """Build Amazon Verified Permissions client."""
         aws_conn_id = conf.get(CONF_SECTION_NAME, CONF_CONN_ID_KEY)
-        return VerifiedPermissionsHook(aws_conn_id=aws_conn_id).conn
+        region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
+        return VerifiedPermissionsHook(aws_conn_id=aws_conn_id, region_name=region_name).conn
 
     @cached_property
     def avp_policy_store_id(self):
@@ -58,9 +61,9 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
         *,
         method: ResourceMethod,
         entity_type: AvpEntities,
-        user: AwsAuthManagerUser,
+        user: AwsAuthManagerUser | None,
         entity_id: str | None = None,
-        entity_fetcher: Callable | None = None,
+        context: dict | None = None,
     ) -> bool:
         """
         Make an authorization decision against Amazon Verified Permissions.
@@ -72,14 +75,12 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
         :param user: the user
         :param entity_id: the entity ID the user accesses. If not provided, all entities of the type will be
             considered.
-        :param entity_fetcher: function that returns list of entities to be passed to Amazon Verified
-            Permissions. Only needed if some resource properties are used in the policies (e.g. DAG folder).
+        :param context: optional additional context to pass to Amazon Verified Permissions.
         """
+        if user is None:
+            return False
+
         entity_list = self._get_user_role_entities(user)
-        if entity_fetcher and entity_id:
-            # If no entity ID is provided, there is no need to fetch entities.
-            # We just need to know whether the user has permissions to access all resources from this type
-            entity_list += entity_fetcher()
 
         self.log.debug(
             "Making authorization request for user=%s, method=%s, entity_type=%s, entity_id=%s",
@@ -89,17 +90,22 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
             entity_id,
         )
 
-        resp = self.avp_client.is_authorized(
-            policyStoreId=self.avp_policy_store_id,
-            principal={"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()},
-            action={
-                "actionType": get_entity_type(AvpEntities.ACTION),
-                "actionId": get_action_id(entity_type, method),
-            },
-            resource={"entityType": get_entity_type(entity_type), "entityId": entity_id or "*"},
-            entities={"entityList": entity_list},
+        request_params = prune_dict(
+            {
+                "policyStoreId": self.avp_policy_store_id,
+                "principal": {"entityType": get_entity_type(AvpEntities.USER), "entityId": user.get_id()},
+                "action": {
+                    "actionType": get_entity_type(AvpEntities.ACTION),
+                    "actionId": get_action_id(entity_type, method),
+                },
+                "resource": {"entityType": get_entity_type(entity_type), "entityId": entity_id or "*"},
+                "entities": {"entityList": entity_list},
+                "context": self._build_context(context),
+            }
         )
 
+        resp = self.avp_client.is_authorized(**request_params)
+
         self.log.debug("Authorization response: %s", resp)
 
         if len(resp.get("errors", [])) > 0:
@@ -124,3 +130,12 @@ class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
             for group in user.get_groups()
         ]
         return [user_entity, *role_entities]
+
+    @staticmethod
+    def _build_context(context: dict | None) -> dict | None:
+        if context is None or len(context) == 0:
+            return None
+
+        return {
+            "contextMap": context,
+        }

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

@@ -16,15 +16,20 @@
 # under the License.
 from __future__ import annotations
 
+import argparse
 from functools import cached_property
 from typing import TYPE_CHECKING
 
 from flask import session, url_for
 
+from airflow.cli.cli_config import CLICommand, DefaultHelpParser, GroupCommand
 from airflow.configuration import conf
 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
+from airflow.providers.amazon.aws.auth_manager.cli.definition import (
+    AWS_AUTH_MANAGER_COMMANDS,
+)
 from airflow.providers.amazon.aws.auth_manager.constants import (
     CONF_ENABLE_KEY,
     CONF_SECTION_NAME,
@@ -122,7 +127,23 @@ class AwsAuthManager(BaseAuthManager):
         details: DagDetails | None = None,
         user: BaseUser | None = None,
     ) -> bool:
-        return self.is_logged_in()
+        dag_id = details.id if details else None
+        context = (
+            None
+            if access_entity is None
+            else {
+                "dag_entity": {
+                    "string": access_entity.value,
+                },
+            }
+        )
+        return self.avp_facade.is_authorized(
+            method=method,
+            entity_type=AvpEntities.DAG,
+            user=user or self.get_user(),
+            entity_id=dag_id,
+            context=context,
+        )
 
     def is_authorized_dataset(
         self, *, method: ResourceMethod, details: DatasetDetails | None = None, user: BaseUser | None = None
@@ -179,3 +200,25 @@ class AwsAuthManager(BaseAuthManager):
     @cached_property
     def security_manager(self) -> AwsSecurityManagerOverride:
         return AwsSecurityManagerOverride(self.appbuilder)
+
+    @staticmethod
+    def get_cli_commands() -> list[CLICommand]:
+        """Vends CLI commands to be included in Airflow CLI."""
+        return [
+            GroupCommand(
+                name="aws-auth-manager",
+                help="Manage resources used by AWS auth manager",
+                subcommands=AWS_AUTH_MANAGER_COMMANDS,
+            ),
+        ]
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Generate documentation; used by Sphinx argparse."""
+    from airflow.cli.cli_parser import AirflowHelpFormatter, _add_command
+
+    parser = DefaultHelpParser(prog="airflow", formatter_class=AirflowHelpFormatter)
+    subparsers = parser.add_subparsers(dest="subcommand", metavar="GROUP_OR_COMMAND")
+    for group_command in AwsAuthManager.get_cli_commands():
+        _add_command(subparsers, group_command)
+    return parser

airflow/providers/amazon/aws/auth_manager/cli/__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.

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

@@ -0,0 +1,178 @@
+# 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.
+"""User sub-commands."""
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import boto3
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.amazon.aws.auth_manager.constants import CONF_REGION_NAME_KEY, CONF_SECTION_NAME
+from airflow.utils import cli as cli_utils
+
+try:
+    from airflow.utils.providers_configuration_loader import providers_configuration_loaded
+except ImportError:
+    raise AirflowOptionalProviderFeatureException(
+        "Failed to import avp_commands. This feature is only available in Airflow "
+        "version >= 2.8.0 where Auth Managers are introduced."
+    )
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+log = logging.getLogger(__name__)
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def init_avp(args):
+    """Initialize Amazon Verified Permissions resources."""
+    client = _get_client()
+
+    # Create the policy store if needed
+    policy_store_id, is_new_policy_store = _create_policy_store(client, args)
+
+    if not is_new_policy_store:
+        print(
+            f"Since an existing policy store with description '{args.policy_store_description}' has been found in Amazon Verified Permissions, "
+            "the CLI nade no changes to this policy store for security reasons. "
+            "Any modification to this policy store must be done manually.",
+        )
+    else:
+        # Set the schema
+        _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(json.dumps({"avp_policy_store_id": policy_store_id}, indent=4))
+
+
+@cli_utils.action_cli
+@providers_configuration_loaded
+def update_schema(args):
+    """Update Amazon Verified Permissions policy store schema."""
+    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():
+    """Returns Amazon Verified Permissions client."""
+    region_name = conf.get(CONF_SECTION_NAME, CONF_REGION_NAME_KEY)
+    return boto3.client("verifiedpermissions", region_name=region_name)
+
+
+def _create_policy_store(client: BaseClient, args) -> tuple[str | None, bool]:
+    """
+    Create if needed the policy store.
+
+    This function returns two elements:
+    - the policy store ID
+    - whether the policy store ID returned refers to a newly created policy store.
+    """
+    paginator = client.get_paginator("list_policy_stores")
+    pages = paginator.paginate()
+    policy_stores = [application for page in pages for application in page["policyStores"]]
+    existing_policy_stores = [
+        policy_store
+        for policy_store in policy_stores
+        if policy_store.get("description") == args.policy_store_description
+    ]
+
+    if args.verbose:
+        log.debug("Policy stores found: %s", policy_stores)
+        log.debug("Existing policy stores found: %s", existing_policy_stores)
+
+    if len(existing_policy_stores) > 0:
+        print(
+            f"There is already a policy store with description '{args.policy_store_description}' in Amazon Verified Permissions: '{existing_policy_stores[0]['policyStoreId']}'."
+        )
+        return existing_policy_stores[0]["policyStoreId"], False
+    else:
+        print(f"No policy store with description '{args.policy_store_description}' found, creating one.")
+        if args.dry_run:
+            print(
+                "Dry run, not creating the policy store with description '{args.policy_store_description}'."
+            )
+            return None, True
+
+        response = client.create_policy_store(
+            validationSettings={
+                "mode": "OFF",
+            },
+            description=args.policy_store_description,
+        )
+        if args.verbose:
+            log.debug("Response from create_policy_store: %s", response)
+
+        print(f"Policy store created: '{response['policyStoreId']}'")
+
+        return response["policyStoreId"], True
+
+
+def _set_schema(client: BaseClient, policy_store_id: str, args) -> None:
+    """Set the policy store schema."""
+    if args.dry_run:
+        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()
+    with open(schema_path) as schema_file:
+        response = client.put_schema(
+            policyStoreId=policy_store_id,
+            definition={
+                "cedarJson": schema_file.read(),
+            },
+        )
+
+        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)

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

@@ -0,0 +1,62 @@
+# 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.
+
+from __future__ import annotations
+
+from airflow.cli.cli_config import (
+    ActionCommand,
+    Arg,
+    lazy_load_command,
+)
+
+############
+# # ARGS # #
+############
+
+ARG_VERBOSE = Arg(("-v", "--verbose"), help="Make logging output more verbose", action="store_true")
+
+ARG_DRY_RUN = Arg(
+    ("--dry-run",),
+    help="Perform a dry run",
+    action="store_true",
+)
+
+# Amazon Verified Permissions
+ARG_POLICY_STORE_DESCRIPTION = Arg(
+    ("--policy-store-description",), help="Policy store description", default="Airflow"
+)
+ARG_POLICY_STORE_ID = Arg(("--policy-store-id",), help="Policy store ID")
+
+
+################
+# # COMMANDS # #
+################
+
+AWS_AUTH_MANAGER_COMMANDS = (
+    ActionCommand(
+        name="init-avp",
+        help="Initialize Amazon Verified resources to be used by AWS manager",
+        func=lazy_load_command("airflow.providers.amazon.aws.auth_manager.cli.avp_commands.init_avp"),
+        args=(ARG_POLICY_STORE_DESCRIPTION, ARG_DRY_RUN, ARG_VERBOSE),
+    ),
+    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'",
+        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/schema.json

@@ -0,0 +1,171 @@
+{
+    "Airflow": {
+        "actions": {
+            "Connection::DELETE": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Connection"]
+                }
+            },
+            "Connection::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Connection"]
+                }
+            },
+            "Connection::POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Connection"]
+                }
+            },
+            "Connection::PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Connection"]
+                }
+            },
+            "Configuration::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Configuration"]
+                }
+            },
+            "Dag::DELETE": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Dag"],
+                    "context": {
+                        "attributes": {
+                            "dag_entity": {
+                                "type": "String",
+                                "required": false
+                            }
+                        },
+                        "type": "Record"
+                    }
+                }
+            },
+            "Dag::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Dag"],
+                    "context": {
+                        "attributes": {
+                            "dag_entity": {
+                                "required": false,
+                                "type": "String"
+                            }
+                        },
+                        "type": "Record"
+                    }
+                }
+            },
+            "Dag::POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Dag"],
+                    "context": {
+                        "attributes": {
+                            "dag_entity": {
+                                "required": false,
+                                "type": "String"
+                            }
+                        },
+                        "type": "Record"
+                    }
+                }
+            },
+            "Dag::PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Dag"],
+                    "context": {
+                        "attributes": {
+                            "dag_entity": {
+                                "required": false,
+                                "type": "String"
+                            }
+                        },
+                        "type": "Record"
+                    }
+                }
+            },
+            "Dataset::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Dataset"]
+                }
+            },
+            "Pool::DELETE": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Pool"]
+                }
+            },
+            "Pool::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Pool"]
+                }
+            },
+            "Pool::POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Pool"]
+                }
+            },
+            "Pool::PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Pool"]
+                }
+            },
+            "Variable::DELETE": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Variable"]
+                }
+            },
+            "Variable::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Variable"]
+                }
+            },
+            "Variable::POST": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Variable"]
+                }
+            },
+            "Variable::PUT": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["Variable"]
+                }
+            },
+            "View::GET": {
+                "appliesTo": {
+                    "principalTypes": ["User"],
+                    "resourceTypes": ["View"]
+                }
+            }
+        },
+        "entityTypes": {
+            "Configuration": {},
+            "Connection": {},
+            "Dag": {},
+            "Dataset": {},
+            "Pool": {},
+            "Role": {},
+            "User": {
+                "memberOfTypes": [
+                    "Role"
+                ]
+            },
+            "Variable": {},
+            "View": {}
+        }
+    }
+}

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

@@ -21,5 +21,6 @@ from __future__ import annotations
 CONF_ENABLE_KEY = "enable"
 CONF_SECTION_NAME = "aws_auth_manager"
 CONF_CONN_ID_KEY = "conn_id"
+CONF_REGION_NAME_KEY = "region_name"
 CONF_SAML_METADATA_URL_KEY = "saml_metadata_url"
 CONF_AVP_POLICY_STORE_ID_KEY = "avp_policy_store_id"