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.
  """CLI for managing ZenML server deployments."""
  
  import ipaddress
  import re
  from typing import List, Optional, Union
  
  import click
  from rich.errors import MarkupError
  
  import zenml
  from zenml.cli import utils as cli_utils
  from zenml.cli.cli import TagGroup, cli
  from zenml.cli.login import login, logout
  from zenml.client import Client
  from zenml.config.global_config import GlobalConfiguration
  from zenml.console import console
  from zenml.enums import CliCategories, ServiceState, StoreType
  from zenml.exceptions import AuthorizationException
  from zenml.logger import get_logger
  from zenml.login.credentials import ServerCredentials, ServerType
  from zenml.utils.server_utils import (
      connected_to_local_server,
      get_local_server,
  )
  
  logger = get_logger(__name__)
  
  
  @cli.command(
      "up",
      help="""Start the ZenML dashboard locally.
  
  DEPRECATED: Please use `zenml login --local` instead.             
  """,
  )
  @click.option(
      "--docker",
      is_flag=True,
      help="Start the ZenML dashboard as a Docker container instead of a local "
      "process.",
      default=False,
      type=click.BOOL,
  )
  @click.option(
      "--port",
      type=int,
      default=None,
      help="Use a custom TCP port value for the ZenML dashboard.",
  )
  @click.option(
      "--ip-address",
      type=ipaddress.ip_address,
      default=None,
      help="Have the ZenML dashboard listen on an IP address different than the "
      "localhost.",
  )
  @click.option(
      "--blocking",
      is_flag=True,
      help="Run the ZenML dashboard in blocking mode. The CLI will not return "
      "until the dashboard is stopped.",
      default=False,
      type=click.BOOL,
  )
  @click.option(
      "--image",
      type=str,
      default=None,
      help="Use a custom Docker image for the ZenML server. Only used when "
      "`--docker` is set.",
  )
  @click.option(
      "--ngrok-token",
      type=str,
      default=None,
      help="Specify an ngrok auth token to use for exposing the ZenML server.",
  )
  def up(
      docker: bool = False,
      ip_address: Union[
          ipaddress.IPv4Address, ipaddress.IPv6Address, None
      ] = None,
      port: Optional[int] = None,
      blocking: bool = False,
      image: Optional[str] = None,
      ngrok_token: Optional[str] = None,
  ) -> None:
      """Start the ZenML dashboard locally and connect the client to it.
  
      Args:
          docker: Use a docker deployment instead of the local process.
          ip_address: The IP address to bind the server to.
          port: The port to bind the server to.
          blocking: Block the CLI while the server is running.
          image: A custom Docker image to use for the server, when the
              `--docker` flag is set.
          ngrok_token: An ngrok auth token to use for exposing the ZenML dashboard
              on a public domain. Primarily used for accessing the dashboard in
              Colab.
      """
      cli_utils.warning(
          "The `zenml up` command is deprecated and will be removed in a "
          "future release. Please use the `zenml login --local` command instead."
      )
  
      # Calling the `zenml login` command
      cli_utils.declare("Calling `zenml login --local`...")
      login.callback(  # type: ignore[misc]
          local=True,
          docker=docker,
          ip_address=ip_address,
          port=port,
          blocking=blocking,
          image=image,
          ngrok_token=ngrok_token,
      )
  
  
  @cli.command(
      "show",
      help="""Show the ZenML dashboard.
  
  DEPRECATED: Please use `zenml server show` instead.             
  """,
  )
  @click.option(
      "--ngrok-token",
      type=str,
      default=None,
      help="Specify an ngrok auth token to use for exposing the ZenML server.",
  )
  def legacy_show(ngrok_token: Optional[str] = None) -> None:
      """Show the ZenML dashboard.
  
      Args:
          ngrok_token: An ngrok auth token to use for exposing the ZenML dashboard
              on a public domain. Primarily used for accessing the dashboard in
              Colab.
      """
      cli_utils.warning(
          "The `zenml show` command is deprecated and will be removed in a "
          "future release. Please use the `zenml server show` command "
          "instead."
      )
  
      # Calling the `zenml server show` command
      cli_utils.declare("Calling `zenml server show`...")
      show(local=False, ngrok_token=ngrok_token)
  
  
  @cli.command(
      "down",
      help="""Shut down the local ZenML dashboard.
  
  DEPRECATED: Please use `zenml logout local` instead.
  """,
  )
  def down() -> None:
      """Shut down the local ZenML dashboard."""
      cli_utils.warning(
          "The `zenml down` command is deprecated and will be removed in a "
          "future release. Please use the `zenml logout --local` command instead."
      )
  
      # Calling the `zenml logout` command
      cli_utils.declare("Calling `zenml logout --local`...")
      logout.callback(  # type: ignore[misc]
          local=True
      )
  
  
  @cli.command(
      "status", help="Show information about the current configuration."
  )
  def status() -> None:
      """Show details about the current configuration."""
      from zenml.login.credentials_store import get_credentials_store
      from zenml.login.pro.client import ZenMLProClient
      from zenml.login.pro.constants import ZENML_PRO_API_URL
  
      gc = GlobalConfiguration()
      client = Client()
      _ = client.zen_store
  
      store_cfg = gc.store_configuration
  
      # Write about the current ZenML client
      cli_utils.declare("-----ZenML Client Status-----")
      if gc.uses_local_store:
          cli_utils.declare(
              f"Connected to the local ZenML database: '{store_cfg.url}'"
          )
      elif connected_to_local_server():
          cli_utils.declare(
              f"Connected to the local ZenML server: {store_cfg.url}"
          )
      elif re.match(r"^mysql://", store_cfg.url):
          cli_utils.declare(
              f"Connected directly to a SQL database: '{store_cfg.url}'"
          )
      else:
          credentials_store = get_credentials_store()
          server = credentials_store.get_credentials(store_cfg.url)
          if server:
              if server.type == ServerType.PRO:
                  # If connected to a ZenML Pro server, refresh the server info
                  if credentials_store.can_login(
                      server_url=server.pro_api_url or ZENML_PRO_API_URL,
                  ):
                      pro_client = ZenMLProClient(
                          server.pro_api_url or ZENML_PRO_API_URL
                      )
                      pro_servers = pro_client.workspace.list(
                          url=store_cfg.url, member_only=True
                      )
                      if pro_servers:
                          credentials_store.update_server_info(
                              server_url=store_cfg.url,
                              server_info=pro_servers[0],
                          )
  
                  cli_utils.declare(
                      f"Connected to a ZenML Pro server: `{server.server_name_hyperlink}`"
                      f" [{server.server_id_hyperlink}]"
                  )
  
                  cli_utils.declare(
                      f"  ZenML Pro Organization: {server.organization_hyperlink}"
                  )
              else:
                  cli_utils.declare(
                      f"Connected to a remote ZenML server: `{server.dashboard_hyperlink}`"
                  )
  
              cli_utils.declare(f"  Dashboard: {server.dashboard_hyperlink}")
              cli_utils.declare(f"  API: {server.api_hyperlink}")
              cli_utils.declare(f"  Server status: '{server.status}'")
              cli_utils.declare(f"  Server authentication: {server.auth_status}")
  
          else:
              cli_utils.declare(
                  f"Connected to a remote ZenML server: [link={store_cfg.url}]"
                  f"{store_cfg.url}[/link]"
              )
  
      try:
          client.zen_store.get_store_info()
      except Exception as e:
          cli_utils.warning(f"Error while initializing client: {e}")
      else:
          # Write about the active entities
          scope = "repository" if client.uses_local_configuration else "global"
          cli_utils.declare(f"  The active user is: '{client.active_user.name}'")
  
          try:
              cli_utils.declare(
                  f"  The active project is: '{client.active_project.name}' ({scope})"
              )
          except RuntimeError:
              cli_utils.declare("  No active project set.")
  
          cli_utils.declare(
              f"  The active stack is: '{client.active_stack_model.name}' ({scope})"
          )
  
      if client.root:
          cli_utils.declare(f"Active repository root: {client.root}")
  
      # Write about the configuration files
      cli_utils.declare(f"Using configuration from: '{gc.config_directory}'")
      cli_utils.declare(
          f"Local store files are located at: '{gc.local_stores_path}'"
      )
  
      cli_utils.declare("\n-----Local ZenML Server Status-----")
      local_server = get_local_server()
      if local_server:
          if local_server.status:
              if local_server.status.status == ServiceState.ACTIVE:
                  cli_utils.declare(
                      f"The local {local_server.config.provider} server is "
                      f"running at: {local_server.status.url}"
                  )
              else:
                  cli_utils.declare(
                      f"The local {local_server.config.provider} server is not "
                      "available."
                  )
                  cli_utils.declare(
                      f"  Server state: {local_server.status.status}"
                  )
                  if local_server.status.status_message:
                      cli_utils.declare(
                          f"  Status message: {local_server.status.status_message}"
                      )
          else:
              cli_utils.declare(
                  f"The local {local_server.config.provider} server is not "
                  "running."
              )
      else:
          cli_utils.declare("The local server has not been started.")
  
  
  @cli.command(
      "connect",
      help=(
          """Connect to a remote ZenML server.
  
      DEPRECATED: Please use `zenml login` instead.
  
      Examples:
  
        * to re-login to the current ZenML server or connect to a ZenML Pro server:
  
          zenml connect
  
        * to log in to a particular ZenML server:
  
          zenml connect --url=http://zenml.example.com:8080
      """
      ),
  )
  @click.option(
      "--url",
      "-u",
      help="The URL where the ZenML server is running.",
      required=False,
      type=str,
  )
  @click.option(
      "--username",
      help="(Deprecated) The username that is used to authenticate with a ZenML "
      "server. If omitted, the web login will be used.",
      required=False,
      type=str,
  )
  @click.option(
      "--password",
      help="(Deprecated) The password that is used to authenticate with a ZenML "
      "server. If omitted, a prompt will be shown to enter the password.",
      required=False,
      type=str,
  )
  @click.option(
      "--api-key",
      help="Use an API key to authenticate with a ZenML server. If "
      "omitted, the web login will be used.",
      required=False,
      type=str,
  )
  @click.option(
      "--no-verify-ssl",
      is_flag=True,
      help="Whether to verify the server's TLS certificate",
      default=False,
  )
  @click.option(
      "--ssl-ca-cert",
      help="A path to a CA bundle file to use to verify the server's TLS "
      "certificate or the CA bundle value itself",
      required=False,
      type=str,
  )
  def connect(
      url: Optional[str] = None,
      username: Optional[str] = None,
      password: Optional[str] = None,
      api_key: Optional[str] = None,
      no_verify_ssl: bool = False,
      ssl_ca_cert: Optional[str] = None,
  ) -> None:
      """Connect to a remote ZenML server.
  
      Args:
          url: The URL where the ZenML server is reachable.
          username: The username that is used to authenticate with the ZenML
              server.
          password: The password that is used to authenticate with the ZenML
              server.
          api_key: The API key that is used to authenticate with the ZenML
              server.
          no_verify_ssl: Whether to verify the server's TLS certificate.
          ssl_ca_cert: A path to a CA bundle to use to verify the server's TLS
              certificate or the CA bundle value itself.
      """
      cli_utils.warning(
          "The `zenml connect` command is deprecated and will be removed in a "
          "future release. Please use the `zenml login` command instead. "
      )
  
      if password is not None or username is not None:
          cli_utils.warning(
              "Connecting to a ZenML server using a username and password is "
              "insecure because the password is locally stored on your "
              "filesystem and is no longer supported. The web login workflow will "
              "be used instead. An alternative for non-interactive environments "
              "is to create and use a service account API key (see "
              "https://docs.zenml.io/deploying-zenml/connecting-to-zenml/connect-with-a-service-account "
              "for more information)."
          )
  
      # Calling the `zenml login` command
      cli_utils.declare("Calling `zenml login`...")
      login.callback(  # type: ignore[misc]
          server=url,
          api_key=api_key,
          no_verify_ssl=no_verify_ssl,
          ssl_ca_cert=ssl_ca_cert,
      )
  
  
  @cli.command(
      "disconnect",
      help="""Disconnect from a ZenML server.
  
  DEPRECATED: Please use `zenml logout` instead.
  """,
  )
  def disconnect_server() -> None:
      """Disconnect from a ZenML server."""
      cli_utils.warning(
          "The `zenml disconnect` command is deprecated and will be removed in a "
          "future release. Please use the `zenml logout` command instead."
      )
  
      # Calling the `zenml logout` command
      cli_utils.declare("Calling `zenml logout`...")
      logout.callback()  # type: ignore[misc]
  
  
  @cli.command("logs", help="Show the logs for the local ZenML server.")
  @click.option(
      "--follow",
      "-f",
      is_flag=True,
      help="Continue to output new log data as it becomes available.",
  )
  @click.option(
      "--tail",
      "-t",
      type=click.INT,
      default=None,
      help="Only show the last NUM lines of log output.",
  )
  @click.option(
      "--raw",
      "-r",
      is_flag=True,
      help="Show raw log contents (don't pretty-print logs).",
  )
  def logs(
      follow: bool = False,
      raw: bool = False,
      tail: Optional[int] = None,
  ) -> None:
      """Display the logs for a ZenML server.
  
      Args:
          follow: Continue to output new log data as it becomes available.
          tail: Only show the last NUM lines of log output.
          raw: Show raw log contents (don't pretty-print logs).
      """
      server = get_local_server()
      if server is None:
          cli_utils.error(
              "The local ZenML dashboard is not running. Please call `zenml "
              "login --local` first to start the ZenML dashboard locally."
          )
  
      from zenml.zen_server.deploy.deployer import LocalServerDeployer
  
      deployer = LocalServerDeployer()
  
      cli_utils.declare(
          f"Showing logs for the local {server.config.provider} server"
      )
  
      from zenml.zen_server.deploy.exceptions import (
          ServerDeploymentNotFoundError,
      )
  
      try:
          logs = deployer.get_server_logs(follow=follow, tail=tail)
      except ServerDeploymentNotFoundError as e:
          cli_utils.error(f"Server not found: {e}")
  
      for line in logs:
          # don't pretty-print log lines that are already pretty-printed
          if raw or line.startswith("\x1b["):
              console.print(line, markup=False)
          else:
              try:
                  console.print(line)
              except MarkupError:
                  console.print(line, markup=False)
  
  
  @cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
  def server() -> None:
      """Commands for managing ZenML servers."""
  
  
  @server.command(
      "list",
      help="""List all ZenML servers that this client is authenticated to.
  
      The CLI can be authenticated to multiple ZenML servers at the same time,
      even though it can only be connected to one server at a time. You can list
      all the ZenML servers that the client is currently authenticated to by
      using this command.
  
      When logged in to ZenML Pro, this list will also include all ZenML Pro
      servers that the authenticated user can access or could potentially access,
      including details such as their current state and the organization they
      belong to.
  
      The complete list of servers displayed by this command includes the
      following:
  
        * ZenML Pro servers that the authenticated ZenML Pro user can or could
          access. The client needs to be logged to ZenML Pro via
          `zenml login --pro` to access these servers.
  
        * ZenML servers that the client has logged in to via
          `zenml login <server-url>` in the past.
  
        * the local ZenML server started with `zenml login --local`, if one is
          running.
  
      By default, this command does not display ZenML servers that are not
      accessible: servers that are not running, are no longer accessible due to
      an expired authentication and ZenML Pro servers where the user is not a
      member. To include these servers in the list, use the `--all` flag.
      """,
  )
  @click.option(
      "--verbose",
      "-v",
      is_flag=True,
      help="Show verbose output.",
  )
  @click.option(
      "--all",
      "-a",
      is_flag=True,
      help="Show all ZenML servers, including those that are not running "
      "and those with an expired authentication.",
  )
  @click.option(
      "--pro-api-url",
      type=str,
      default=None,
      help="Custom URL for the ZenML Pro API. Useful when disconnecting "
      "from a self-hosted ZenML Pro deployment.",
  )
  def server_list(
      verbose: bool = False,
      all: bool = False,
      pro_api_url: Optional[str] = None,
  ) -> None:
      """List all ZenML servers that this client is authorized to access.
  
      Args:
          verbose: Whether to show verbose output.
          all: Whether to show all ZenML servers.
          pro_api_url: Custom URL for the ZenML Pro API.
      """
      from zenml.login.credentials_store import get_credentials_store
      from zenml.login.pro.client import ZenMLProClient
      from zenml.login.pro.constants import ZENML_PRO_API_URL
      from zenml.login.pro.workspace.models import WorkspaceRead, WorkspaceStatus
  
      pro_api_url = pro_api_url or ZENML_PRO_API_URL
      pro_api_url = pro_api_url.rstrip("/")
  
      credentials_store = get_credentials_store()
      pro_token = credentials_store.get_pro_token(
          allow_expired=True, pro_api_url=pro_api_url
      )
      current_store_config = GlobalConfiguration().store_configuration
  
      # The list of ZenML Pro servers kept in the credentials store
      pro_servers = credentials_store.list_credentials(type=ServerType.PRO)
      # The list of regular remote ZenML servers kept in the credentials store
      servers = list(credentials_store.list_credentials(type=ServerType.REMOTE))
      # The list of local ZenML servers kept in the credentials store
      local_servers = list(
          credentials_store.list_credentials(type=ServerType.LOCAL)
      )
  
      if pro_token and not pro_token.expired:
          # If the ZenML Pro authentication is still valid, we include all ZenML
          # Pro servers that the current ZenML Pro user can access, even those
          # that the user has never connected to (and are therefore not stored in
          # the credentials store).
  
          accessible_pro_servers: List[WorkspaceRead] = []
          try:
              client = ZenMLProClient(pro_api_url)
              accessible_pro_servers = client.workspace.list(member_only=not all)
          except AuthorizationException as e:
              cli_utils.warning(f"ZenML Pro authorization error: {e}")
  
          # We update the list of stored ZenML Pro servers with the ones that the
          # client is a member of
          for accessible_server in accessible_pro_servers:
              for idx, stored_server in enumerate(pro_servers):
                  if stored_server.server_id == accessible_server.id:
                      # All ZenML Pro servers accessible by the current ZenML Pro
                      # user have an authentication that is valid at least until
                      # the current ZenML Pro authentication token expires.
                      stored_server.update_server_info(
                          accessible_server,
                      )
                      updated_server = stored_server.model_copy()
                      # Replace the current server API token with the current
                      # ZenML Pro API token to reflect the current authentication
                      # status.
                      updated_server.api_token = pro_token
                      pro_servers[idx] = updated_server
                      break
              else:
                  stored_server = ServerCredentials(
                      url=accessible_server.url or "",
                      api_token=pro_token,
                  )
                  stored_server.update_server_info(accessible_server)
                  pro_servers.append(stored_server)
  
          if not all:
              accessible_pro_servers = [
                  s
                  for s in accessible_pro_servers
                  if s.status == WorkspaceStatus.AVAILABLE
              ]
  
          if not accessible_pro_servers:
              cli_utils.declare(
                  "No ZenML Pro servers that are accessible to the current "
                  "user could be found."
              )
              if not all:
                  cli_utils.declare(
                      "Hint: use the `--all` flag to show all ZenML servers, "
                      "including those that the client is not currently "
                      "authorized to access or are not running."
                  )
  
      elif pro_servers:
          cli_utils.warning(
              "The ZenML Pro authentication has expired. Please re-login "
              "to ZenML Pro using `zenml login` to include all ZenML Pro servers "
              "that you are a member of in the list."
          )
  
      # We add the local server to the list of servers, if it is running
      local_server = get_local_server()
      if local_server:
          url = (
              local_server.status.url if local_server.status else None
          ) or local_server.config.url
          status = local_server.status.status if local_server.status else ""
          local_servers.append(
              ServerCredentials(
                  url=url or "",
                  status=status,
                  version=zenml.__version__,
                  server_id=GlobalConfiguration().user_id,
                  server_name=f"local {local_server.config.provider} server",
              )
          )
  
      all_servers = pro_servers + local_servers + servers
  
      if not all:
          # Filter out servers that are expired or not running
          all_servers = [s for s in all_servers if s.is_available]
  
      if verbose:
          columns = [
              "type",
              "server_id_hyperlink",
              "server_name_hyperlink",
              "organization_hyperlink" if pro_servers else "",
              "version",
              "status",
              "dashboard_url",
              "api_hyperlink",
              "auth_status",
          ]
      elif all:
          columns = [
              "type",
              "server_id_hyperlink",
              "server_name_hyperlink",
              "organization_hyperlink" if pro_servers else "",
              "version",
              "status",
              "api_hyperlink",
          ]
      else:
          columns = [
              "type",
              "server_id_hyperlink" if pro_servers else "",
              "server_name_hyperlink",
              "organization_hyperlink" if pro_servers else "",
              "version",
              "api_hyperlink" if servers else "",
          ]
  
      # Remove empty columns
      columns = [c for c in columns if c]
  
      # Figure out if the client is already connected to one of the
      # servers in the list
      current_server: List[ServerCredentials] = []
      if current_store_config.type == StoreType.REST:
          current_server = [
              s for s in all_servers if s.url == current_store_config.url
          ]
  
      cli_utils.print_pydantic_models(  # type: ignore[type-var]
          all_servers,
          columns=columns,
          active_models=current_server,
          show_active=True,
          column_aliases={
              "server_id_hyperlink": "server_id",
              "server_name_hyperlink": "server_name",
              "organization_hyperlink": "organization",
          },
      )
  
  
  @server.command(
      "show",
      help="Show the ZenML dashboard for the server that the client is connected to.",
  )
  @click.option(
      "--local",
      is_flag=True,
      help="Show the ZenML dashboard for the local server.",
      default=False,
      type=click.BOOL,
  )
  @click.option(
      "--ngrok-token",
      type=str,
      default=None,
      help="Specify an ngrok auth token to use for exposing the local ZenML "
      "server. Only used when `--local` is set. Primarily used for accessing the "
      "local dashboard in Colab.",
  )
  def show(local: bool = False, ngrok_token: Optional[str] = None) -> None:
      """Show the ZenML dashboard.
  
      Args:
          local: Whether to show the ZenML dashboard for the local server.
          ngrok_token: An ngrok auth token to use for exposing the ZenML dashboard
              on a public domain. Primarily used for accessing the local dashboard
              in Colab.
      """
      try:
          zenml.show(ngrok_token=ngrok_token)
      except RuntimeError as e:
          cli_utils.exception(e)
  