content: #  Copyright (c) ZenML GmbH 2023. 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.
  """Util functions for handling stacks, components, and flavors."""
  
  from typing import Any, Dict, Optional, Union
  
  from zenml.client import Client
  from zenml.enums import StackComponentType, StoreType
  from zenml.logger import get_logger
  from zenml.models import FlavorFilter, FlavorResponse
  from zenml.stack.flavor import Flavor
  from zenml.stack.stack_component import StackComponentConfig
  from zenml.zen_stores.base_zen_store import BaseZenStore
  
  logger = get_logger(__name__)
  
  
  def validate_stack_component_config(
      configuration_dict: Dict[str, Any],
      flavor: Union[FlavorResponse, str],
      component_type: StackComponentType,
      zen_store: Optional[BaseZenStore] = None,
      validate_custom_flavors: bool = True,
      existing_config: Optional[Dict[str, Any]] = None,
  ) -> Optional[StackComponentConfig]:
      """Validate the configuration of a stack component.
  
      Args:
          configuration_dict: The stack component configuration to validate.
          flavor: The model or name of the flavor of the stack component.
          component_type: The type of the stack component.
          zen_store: An optional ZenStore in which to look for the flavor. If not
              provided, the flavor will be fetched via the regular ZenML Client.
              This is mainly useful for checks running inside the ZenML server.
          validate_custom_flavors: When loading custom flavors from the local
              environment, this flag decides whether the import failures are
              raised or an empty value is returned.
          existing_config: The existing stack component configuration.
  
      Returns:
          The validated stack component configuration or None, if the
          flavor is a custom flavor that could not be imported from the local
          environment and the `validate_custom_flavors` flag is set to False.
  
      Raises:
          ValueError: If the configuration is invalid.
          ImportError: If the flavor class could not be imported.
          ModuleNotFoundError: If the flavor class could not be imported.
      """
      if isinstance(flavor, FlavorResponse):
          flavor_model = flavor
      else:
          if zen_store:
              flavor_model = get_flavor_by_name_and_type_from_zen_store(
                  zen_store=zen_store,
                  flavor_name=flavor,
                  component_type=component_type,
              )
          else:
              flavor_model = Client().get_flavor_by_name_and_type(
                  name=flavor,
                  component_type=component_type,
              )
  
      try:
          flavor_class = Flavor.from_model(flavor_model)
      except (ImportError, ModuleNotFoundError):
          # The flavor class couldn't be loaded.
          if flavor_model.is_custom and not validate_custom_flavors:
              return None
          raise
  
      config_class = flavor_class.config_class
  
      extra_keys = set(configuration_dict) - set(config_class.model_fields)
      if existing_config:
          # If there is an existing configuration, don't fail because of keys
          # that were already set. This can be the case if the config stored in
          # the DB contains attributes that have been removed.
          extra_keys = extra_keys - set(existing_config)
  
      if extra_keys:
          raise ValueError(
              f"Invalid configuration keys {list(extra_keys)} for "
              f"{flavor_model.name} {component_type}."
          )
  
      configuration = config_class(**configuration_dict)
  
      if not configuration.is_valid:
          raise ValueError(
              f"Invalid stack component configuration. Please verify "
              f"the configurations set for {component_type}."
          )
      return configuration
  
  
  def warn_if_config_server_mismatch(
      configuration: StackComponentConfig,
  ) -> None:
      """Warn if a component configuration is mismatched with the ZenML server.
  
      Args:
          configuration: The component configuration to check.
      """
      zen_store = Client().zen_store
      if configuration.is_remote and zen_store.is_local_store():
          if zen_store.type == StoreType.REST:
              logger.warning(
                  "You are configuring a stack component that is running "
                  "remotely while using a local ZenML server. The component "
                  "may not be able to reach the local ZenML server and will "
                  "therefore not be functional. Please consider deploying "
                  "and/or using a remote ZenML server instead."
              )
          else:
              logger.warning(
                  "You are configuring a stack component that is running "
                  "remotely while connected to the local database. The component "
                  "will not be able to reach the local database and will "
                  "therefore not be functional. Please consider deploying "
                  "and/or using a remote ZenML server instead."
              )
      elif configuration.is_local and not zen_store.is_local_store():
          logger.warning(
              "You are configuring a stack component that is using "
              "local resources while connected to a remote ZenML server. The "
              "stack component may not be usable from other hosts or by "
              "other users. You should consider using a non-local stack "
              "component alternative instead."
          )
  
  
  def get_flavor_by_name_and_type_from_zen_store(
      zen_store: BaseZenStore,
      flavor_name: str,
      component_type: StackComponentType,
  ) -> FlavorResponse:
      """Get a stack component flavor by name and type from a ZenStore.
  
      Args:
          zen_store: The ZenStore to query.
          flavor_name: The name of a stack component flavor.
          component_type: The type of the stack component.
  
      Returns:
          The flavor model.
  
      Raises:
          KeyError: If no flavor with the given name and type exists.
      """
      flavors = zen_store.list_flavors(
          FlavorFilter(name=flavor_name, type=component_type)
      )
      if not flavors:
          raise KeyError(
              f"No flavor with name '{flavor_name}' and type "
              f"'{component_type}' exists."
          )
      return flavors[0]
  