content: #  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
  #
  #  Licensed 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:
  #
  #       https://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.
  """High-level helper functions to write endpoints with RBAC."""
  
  from typing import Any, Callable, List, Optional, Tuple, TypeVar, Union
  from uuid import UUID
  
  from zenml.models import (
      BaseFilter,
      BaseIdentifiedResponse,
      BaseRequest,
      BaseUpdate,
      Page,
      ProjectScopedFilter,
      UserScopedRequest,
  )
  from zenml.zen_server.feature_gate.endpoint_utils import (
      check_entitlement,
      report_usage,
  )
  from zenml.zen_server.rbac.models import Action, ResourceType
  from zenml.zen_server.rbac.utils import (
      batch_verify_permissions_for_models,
      dehydrate_page,
      dehydrate_response_model,
      dehydrate_response_model_batch,
      delete_model_resource,
      get_allowed_resource_ids,
      get_resource_type_for_model,
      verify_permission,
      verify_permission_for_model,
  )
  from zenml.zen_server.utils import (
      get_auth_context,
      server_config,
      set_filter_project_scope,
  )
  
  AnyRequest = TypeVar("AnyRequest", bound=BaseRequest)
  AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
  AnyOtherResponse = TypeVar("AnyOtherResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
  AnyFilter = TypeVar("AnyFilter", bound=BaseFilter)
  AnyUpdate = TypeVar("AnyUpdate", bound=BaseUpdate)
  UUIDOrStr = TypeVar("UUIDOrStr", UUID, Union[UUID, str])
  
  
  def verify_permissions_and_create_entity(
      request_model: AnyRequest,
      create_method: Callable[[AnyRequest], AnyResponse],
      surrogate_models: Optional[List[AnyOtherResponse]] = None,
      skip_entitlements: bool = False,
  ) -> AnyResponse:
      """Verify permissions and create the entity if authorized.
  
      Args:
          request_model: The entity request model.
          create_method: The method to create the entity.
          surrogate_models: Optional list of surrogate models to verify
              UPDATE permissions for instead of verifying CREATE permissions for
              the request model.
          skip_entitlements: Whether to skip the entitlement check and usage
              increment.
  
      Returns:
          A model of the created entity.
      """
      if isinstance(request_model, UserScopedRequest):
          auth_context = get_auth_context()
          assert auth_context
  
          # Ignore the user field set in the request model, if any, and set it to
          # the current user's ID instead. This is just a precaution, given that
          # the SQLZenStore also does this same validation on all request models.
          request_model.user = auth_context.user.id
  
      if surrogate_models:
          batch_verify_permissions_for_models(
              models=surrogate_models, action=Action.UPDATE
          )
      else:
          verify_permission_for_model(model=request_model, action=Action.CREATE)
  
      resource_type = get_resource_type_for_model(request_model)
  
      if resource_type:
          needs_usage_increment = (
              not skip_entitlements
              and resource_type in server_config().reportable_resources
          )
          if needs_usage_increment:
              check_entitlement(resource_type)
  
      created = create_method(request_model)
  
      if resource_type and needs_usage_increment:
          report_usage(resource_type, resource_id=created.id)
  
      return dehydrate_response_model(created)
  
  
  def verify_permissions_and_batch_create_entity(
      batch: List[AnyRequest],
      create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
  ) -> List[AnyResponse]:
      """Verify permissions and create a batch of entities if authorized.
  
      Args:
          batch: The batch to create.
          create_method: The method to create the entities.
  
      Raises:
          RuntimeError: If the resource type is usage-tracked.
  
      Returns:
          The created entities.
      """
      auth_context = get_auth_context()
      assert auth_context
  
      resource_types = set()
      for request_model in batch:
          resource_type = get_resource_type_for_model(request_model)
          if resource_type:
              resource_types.add(resource_type)
  
          if isinstance(request_model, UserScopedRequest):
              # Ignore the user field set in the request model, if any, and set it
              # to the current user's ID instead. This is just a precaution, given
              # that the SQLZenStore also does this same validation on all request
              # models.
              request_model.user = auth_context.user.id
  
      batch_verify_permissions_for_models(models=batch, action=Action.CREATE)
  
      if resource_types & set(server_config().reportable_resources):
          raise RuntimeError(
              "Batch requests are currently not possible with usage-tracked "
              "features."
          )
  
      created = create_method(batch)
      return dehydrate_response_model_batch(created)
  
  
  def verify_permissions_and_get_or_create_entity(
      request_model: AnyRequest,
      get_or_create_method: Callable[
          [AnyRequest, Optional[Callable[[], None]]], Tuple[AnyResponse, bool]
      ],
  ) -> Tuple[AnyResponse, bool]:
      """Verify permissions and create the entity if authorized.
  
      Args:
          request_model: The entity request model.
          get_or_create_method: The method to get or create the entity.
  
      Returns:
          The entity and a boolean indicating whether the entity was created.
      """
      if isinstance(request_model, UserScopedRequest):
          auth_context = get_auth_context()
          assert auth_context
  
          # Ignore the user field set in the request model, if any, and set it to
          # the current user's ID instead. This is just a precaution, given that
          # the SQLZenStore also does this same validation on all request models.
          request_model.user = auth_context.user.id
  
      resource_type = get_resource_type_for_model(request_model)
      needs_usage_increment = (
          resource_type and resource_type in server_config().reportable_resources
      )
  
      def _pre_creation_hook() -> None:
          verify_permission_for_model(model=request_model, action=Action.CREATE)
          if resource_type and needs_usage_increment:
              check_entitlement(feature=resource_type)
  
      model, created = get_or_create_method(request_model, _pre_creation_hook)
  
      if not created:
          verify_permission_for_model(model=model, action=Action.READ)
      elif resource_type and needs_usage_increment:
          report_usage(resource_type, resource_id=model.id)
  
      return dehydrate_response_model(model), created
  
  
  def verify_permissions_and_get_entity(
      id: UUIDOrStr,
      get_method: Callable[[UUIDOrStr], AnyResponse],
      **get_method_kwargs: Any,
  ) -> AnyResponse:
      """Verify permissions and fetch an entity.
  
      Args:
          id: The ID of the entity to fetch.
          get_method: The method to fetch the entity.
          get_method_kwargs: Keyword arguments to pass to the get method.
  
      Returns:
          A model of the fetched entity.
      """
      model = get_method(id, **get_method_kwargs)
      verify_permission_for_model(model, action=Action.READ)
      return dehydrate_response_model(model)
  
  
  def verify_permissions_and_list_entities(
      filter_model: AnyFilter,
      resource_type: ResourceType,
      list_method: Callable[[AnyFilter], Page[AnyResponse]],
      **list_method_kwargs: Any,
  ) -> Page[AnyResponse]:
      """Verify permissions and list entities.
  
      Args:
          filter_model: The entity filter model.
          resource_type: The resource type of the entities to list.
          list_method: The method to list the entities.
          list_method_kwargs: Keyword arguments to pass to the list method.
  
      Returns:
          A page of entity models.
  
      Raises:
          ValueError: If the filter's project scope is not set or is not a UUID.
      """
      auth_context = get_auth_context()
      assert auth_context
  
      project_id: Optional[UUID] = None
      if isinstance(filter_model, ProjectScopedFilter):
          # A project scoped filter must always be scoped to a specific
          # project. This is required for the RBAC check to work.
          set_filter_project_scope(filter_model)
          if not filter_model.project or not isinstance(
              filter_model.project, UUID
          ):
              raise ValueError(
                  "Project scope must be a UUID, got "
                  f"{type(filter_model.project)}."
              )
          project_id = filter_model.project
  
      allowed_ids = get_allowed_resource_ids(
          resource_type=resource_type, project_id=project_id
      )
      filter_model.configure_rbac(
          authenticated_user_id=auth_context.user.id, id=allowed_ids
      )
      page = list_method(filter_model, **list_method_kwargs)
      return dehydrate_page(page)
  
  
  def verify_permissions_and_update_entity(
      id: UUIDOrStr,
      update_model: AnyUpdate,
      get_method: Callable[[UUIDOrStr, bool], AnyResponse],
      update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
      **update_method_kwargs: Any,
  ) -> AnyResponse:
      """Verify permissions and update an entity.
  
      Args:
          id: The ID of the entity to update.
          update_model: The entity update model.
          get_method: The method to fetch the entity.
          update_method: The method to update the entity.
          update_method_kwargs: Keyword arguments to pass to the update method.
  
      Returns:
          A model of the updated entity.
      """
      # We don't need the hydrated version here
      model = get_method(id, False)
      verify_permission_for_model(model, action=Action.UPDATE)
      updated_model = update_method(
          model.id, update_model, **update_method_kwargs
      )
      return dehydrate_response_model(updated_model)
  
  
  def verify_permissions_and_delete_entity(
      id: UUIDOrStr,
      get_method: Callable[[UUIDOrStr, bool], AnyResponse],
      delete_method: Callable[[UUIDOrStr], None],
      **delete_method_kwargs: Any,
  ) -> AnyResponse:
      """Verify permissions and delete an entity.
  
      Args:
          id: The ID of the entity to delete.
          get_method: The method to fetch the entity.
          delete_method: The method to delete the entity.
          delete_method_kwargs: Keyword arguments to pass to the delete method.
  
      Returns:
          The deleted entity.
      """
      model = get_method(id, True)
      verify_permission_for_model(model, action=Action.DELETE)
      delete_method(model.id, **delete_method_kwargs)
      delete_model_resource(model)
  
      return model
  
  
  def verify_permissions_and_prune_entities(
      resource_type: ResourceType,
      prune_method: Callable[..., None],
      project_id: Optional[UUID] = None,
      **kwargs: Any,
  ) -> None:
      """Verify permissions and prune entities of certain type.
  
      Args:
          resource_type: The resource type of the entities to prune.
          prune_method: The method to prune the entities.
          project_id: The project ID to prune the entities for.
          kwargs: Keyword arguments to pass to the prune method.
      """
      verify_permission(
          resource_type=resource_type,
          action=Action.PRUNE,
          project_id=project_id,
      )
      prune_method(**kwargs)
  