content: #  Copyright (c) ZenML GmbH 2022. 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.
  """Endpoint definitions for authentication (login)."""
  
  from typing import List, Optional
  
  from fastapi import APIRouter, Security
  
  import zenml
  from zenml.constants import (
      ACTIVATE,
      API,
      INFO,
      LOAD_INFO,
      ONBOARDING_STATE,
      SERVER_SETTINGS,
      STATISTICS,
      VERSION_1,
  )
  from zenml.enums import AuthScheme
  from zenml.exceptions import IllegalOperationError
  from zenml.models import (
      ComponentFilter,
      ProjectFilter,
      ServerActivationRequest,
      ServerLoadInfo,
      ServerModel,
      ServerSettingsResponse,
      ServerSettingsUpdate,
      ServerStatistics,
      StackFilter,
      UserResponse,
  )
  from zenml.zen_server.auth import AuthContext, authorize
  from zenml.zen_server.exceptions import error_response
  from zenml.zen_server.rbac.models import ResourceType
  from zenml.zen_server.rbac.utils import get_allowed_resource_ids
  from zenml.zen_server.utils import (
      async_fastapi_endpoint_wrapper,
      server_config,
      zen_store,
  )
  
  router = APIRouter(
      prefix=API + VERSION_1,
      tags=["server"],
      responses={401: error_response},
  )
  
  
  @router.get("/version")
  async def version() -> str:
      """Get version of the server.
  
      Returns:
          String representing the version of the server.
      """
      return zenml.__version__
  
  
  @router.get(
      INFO,
      responses={401: error_response, 404: error_response, 422: error_response},
  )
  @async_fastapi_endpoint_wrapper
  def server_info() -> ServerModel:
      """Get information about the server.
  
      Returns:
          Information about the server.
      """
      return zen_store().get_store_info()
  
  
  @router.get(
      LOAD_INFO,
      response_model=ServerLoadInfo,
  )
  @async_fastapi_endpoint_wrapper
  def server_load_info(_: AuthContext = Security(authorize)) -> ServerLoadInfo:
      """Get information about the server load.
  
      Returns:
          Information about the server load.
      """
      import threading
  
      # Get the current number of threads
      num_threads = len(threading.enumerate())
  
      store = zen_store()
  
      if store.config.driver == "sqlite":
          # SQLite doesn't have a connection pool
          return ServerLoadInfo(
              threads=num_threads,
              db_connections_total=0,
              db_connections_active=0,
              db_connections_overflow=0,
          )
  
      from sqlalchemy.pool import QueuePool
  
      # Get the number of connections
      pool = store.engine.pool
      assert isinstance(pool, QueuePool)
      idle_conn = pool.checkedin()
      active_conn = pool.checkedout()
      overflow_conn = max(0, pool.overflow())
      total_conn = idle_conn + active_conn
  
      return ServerLoadInfo(
          threads=num_threads,
          db_connections_total=total_conn,
          db_connections_active=active_conn,
          db_connections_overflow=overflow_conn,
      )
  
  
  @router.get(
      ONBOARDING_STATE,
      responses={
          401: error_response,
          404: error_response,
          422: error_response,
      },
  )
  @async_fastapi_endpoint_wrapper
  def get_onboarding_state(
      _: AuthContext = Security(authorize),
  ) -> List[str]:
      """Get the onboarding state of the server.
  
      Returns:
          The onboarding state of the server.
      """
      return zen_store().get_onboarding_state()
  
  
  # We don't have any concrete value that tells us whether a server is a cloud
  # workspace, so we use `external_server_id` as the best proxy option.
  # For cloud workspaces, we don't add these endpoints as the server settings don't
  # have any effect and even allow users to disable functionality that is
  # necessary for the cloud onboarding to work.
  if server_config().external_server_id is None:
  
      @router.get(
          SERVER_SETTINGS,
          responses={
              401: error_response,
              404: error_response,
              422: error_response,
          },
      )
      @async_fastapi_endpoint_wrapper
      def get_settings(
          _: AuthContext = Security(authorize),
          hydrate: bool = True,
      ) -> ServerSettingsResponse:
          """Get settings of the server.
  
          Args:
              hydrate: Whether to hydrate the response.
  
          Returns:
              Settings of the server.
          """
          return zen_store().get_server_settings(hydrate=hydrate)
  
      @router.put(
          SERVER_SETTINGS,
          responses={
              401: error_response,
              404: error_response,
              422: error_response,
          },
      )
      @async_fastapi_endpoint_wrapper
      def update_server_settings(
          settings_update: ServerSettingsUpdate,
          auth_context: AuthContext = Security(authorize),
      ) -> ServerSettingsResponse:
          """Updates the settings of the server.
  
          Args:
              settings_update: Settings update.
              auth_context: Authentication context.
  
          Raises:
              IllegalOperationError: If trying to update admin properties without
                  admin permissions.
  
          Returns:
              The updated settings.
          """
          if not server_config().rbac_enabled:
              will_update_admin_properties = bool(
                  settings_update.model_dump(
                      exclude_none=True, exclude={"onboarding_state"}
                  )
              )
  
              if not auth_context.user.is_admin and will_update_admin_properties:
                  raise IllegalOperationError(
                      "Only admins can update server settings."
                  )
  
          return zen_store().update_server_settings(settings_update)
  
  
  # When the auth scheme is set to EXTERNAL, users cannot be managed via the
  # API and the server is activated on deployment
  if server_config().auth_scheme != AuthScheme.EXTERNAL:
  
      @router.put(
          ACTIVATE,
          responses={
              401: error_response,
              404: error_response,
              422: error_response,
          },
      )
      @async_fastapi_endpoint_wrapper
      def activate_server(
          activate_request: ServerActivationRequest,
      ) -> Optional[UserResponse]:
          """Updates a stack.
  
          Args:
              activate_request: The request to activate the server.
  
          Returns:
              The default admin user that was created during activation, if any.
          """
          return zen_store().activate_server(activate_request)
  
  
  @router.get(
      STATISTICS,
      responses={401: error_response, 404: error_response, 422: error_response},
  )
  @async_fastapi_endpoint_wrapper
  def get_server_statistics(
      auth_context: AuthContext = Security(authorize),
  ) -> ServerStatistics:
      """Gets server statistics.
  
      Args:
          auth_context: Authentication context.
  
      Returns:
          Statistics of the server.
      """
      user_id = auth_context.user.id
      component_filter = ComponentFilter()
      component_filter.configure_rbac(
          authenticated_user_id=user_id,
          id=get_allowed_resource_ids(
              resource_type=ResourceType.STACK_COMPONENT
          ),
      )
  
      project_filter = ProjectFilter()
      project_filter.configure_rbac(
          authenticated_user_id=user_id,
          id=get_allowed_resource_ids(resource_type=ResourceType.PROJECT),
      )
  
      stack_filter = StackFilter()
      stack_filter.configure_rbac(
          authenticated_user_id=user_id,
          id=get_allowed_resource_ids(resource_type=ResourceType.STACK),
      )
  
      return ServerStatistics(
          stacks=zen_store().count_stacks(filter_model=stack_filter),
          components=zen_store().count_stack_components(
              filter_model=component_filter
          ),
          projects=zen_store().count_projects(filter_model=project_filter),
      )
  