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.
  """Models representing components."""
  
  from typing import (
      TYPE_CHECKING,
      Any,
      ClassVar,
      Dict,
      List,
      Optional,
      Type,
      TypeVar,
      Union,
  )
  from uuid import UUID
  
  from pydantic import BaseModel, Field, field_validator
  
  from zenml.constants import STR_FIELD_MAX_LENGTH
  from zenml.enums import LogicalOperators, StackComponentType
  from zenml.models.v2.base.base import BaseUpdate
  from zenml.models.v2.base.scoped import (
      UserScopedFilter,
      UserScopedRequest,
      UserScopedResponse,
      UserScopedResponseBody,
      UserScopedResponseMetadata,
      UserScopedResponseResources,
  )
  from zenml.utils import secret_utils
  
  if TYPE_CHECKING:
      from sqlalchemy.sql.elements import ColumnElement
  
      from zenml.models import FlavorResponse, ServiceConnectorResponse
      from zenml.zen_stores.schemas.base_schemas import BaseSchema
  
      AnySchema = TypeVar("AnySchema", bound=BaseSchema)
  
  # ------------------ Base Model ------------------
  
  
  class ComponentBase(BaseModel):
      """Base model for components."""
  
      name: str = Field(
          title="The name of the stack component.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
  
      type: StackComponentType = Field(
          title="The type of the stack component.",
      )
  
      flavor: str = Field(
          title="The flavor of the stack component.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
      environment: Optional[Dict[str, str]] = Field(
          default=None,
          title="Environment variables to set when running on this component.",
      )
      secrets: Optional[List[Union[UUID, str]]] = Field(
          default=None,
          title="Secrets to set as environment variables when running on this component.",
      )
  
      configuration: Dict[str, Any] = Field(
          title="The stack component configuration.",
      )
  
      connector_resource_id: Optional[str] = Field(
          default=None,
          description="The ID of a specific resource instance to "
          "gain access to through the connector",
      )
  
      labels: Optional[Dict[str, Any]] = Field(
          default=None,
          title="The stack component labels.",
      )
  
  
  # ------------------ Request Model ------------------
  
  
  class ComponentRequest(ComponentBase, UserScopedRequest):
      """Request model for stack components."""
  
      ANALYTICS_FIELDS: ClassVar[List[str]] = ["type", "flavor"]
  
      connector: Optional[UUID] = Field(
          default=None,
          title="The service connector linked to this stack component.",
      )
  
      @field_validator("name")
      @classmethod
      def name_cant_be_a_secret_reference(cls, name: str) -> str:
          """Validator to ensure that the given name is not a secret reference.
  
          Args:
              name: The name to validate.
  
          Returns:
              The name if it is not a secret reference.
  
          Raises:
              ValueError: If the name is a secret reference.
          """
          if secret_utils.is_secret_reference(name):
              raise ValueError(
                  "Passing the `name` attribute of a stack component as a "
                  "secret reference is not allowed."
              )
          return name
  
  
  class DefaultComponentRequest(ComponentRequest):
      """Internal component request model used only for default stack components."""
  
  
  # ------------------ Update Model ------------------
  
  
  class ComponentUpdate(BaseUpdate):
      """Update model for stack components."""
  
      name: Optional[str] = Field(
          title="The name of the stack component.",
          max_length=STR_FIELD_MAX_LENGTH,
          default=None,
      )
      configuration: Optional[Dict[str, Any]] = Field(
          title="The stack component configuration.",
          default=None,
      )
      environment: Optional[Dict[str, str]] = Field(
          default=None,
          title="Environment variables to set when running on this component.",
      )
      connector_resource_id: Optional[str] = Field(
          description="The ID of a specific resource instance to "
          "gain access to through the connector",
          default=None,
      )
      labels: Optional[Dict[str, Any]] = Field(
          title="The stack component labels.",
          default=None,
      )
      connector: Optional[UUID] = Field(
          title="The service connector linked to this stack component.",
          default=None,
      )
      add_secrets: Optional[List[Union[UUID, str]]] = Field(
          default=None,
          title="New secrets to add to the stack component.",
      )
      remove_secrets: Optional[List[Union[UUID, str]]] = Field(
          default=None,
          title="Secrets to remove from the stack component.",
      )
  
  
  # ------------------ Response Model ------------------
  
  
  class ComponentResponseBody(UserScopedResponseBody):
      """Response body for stack components."""
  
      type: StackComponentType = Field(
          title="The type of the stack component.",
      )
      flavor_name: str = Field(
          title="The flavor of the stack component.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
      integration: Optional[str] = Field(
          default=None,
          title="The name of the integration that the component's flavor "
          "belongs to.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
      logo_url: Optional[str] = Field(
          default=None,
          title="Optionally, a url pointing to a png,"
          "svg or jpg can be attached.",
      )
  
  
  class ComponentResponseMetadata(UserScopedResponseMetadata):
      """Response metadata for stack components."""
  
      configuration: Dict[str, Any] = Field(
          title="The stack component configuration.",
      )
      environment: Dict[str, str] = Field(
          default={},
          title="Environment variables to set when running on this component.",
      )
      secrets: List[UUID] = Field(
          default=[],
          title="Secrets to set as environment variables when running on this "
          "component.",
      )
      labels: Optional[Dict[str, Any]] = Field(
          default=None,
          title="The stack component labels.",
      )
      connector_resource_id: Optional[str] = Field(
          default=None,
          description="The ID of a specific resource instance to "
          "gain access to through the connector",
      )
      connector: Optional["ServiceConnectorResponse"] = Field(
          default=None,
          title="The service connector linked to this stack component.",
      )
  
  
  class ComponentResponseResources(UserScopedResponseResources):
      """Response resources for stack components."""
  
      flavor: "FlavorResponse" = Field(
          title="The flavor of this stack component.",
      )
  
  
  class ComponentResponse(
      UserScopedResponse[
          ComponentResponseBody,
          ComponentResponseMetadata,
          ComponentResponseResources,
      ]
  ):
      """Response model for stack components."""
  
      ANALYTICS_FIELDS: ClassVar[List[str]] = ["type"]
  
      name: str = Field(
          title="The name of the stack component.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
  
      def get_analytics_metadata(self) -> Dict[str, Any]:
          """Add the component labels to analytics metadata.
  
          Returns:
              Dict of analytics metadata.
          """
          metadata = super().get_analytics_metadata()
  
          if self.labels is not None:
              metadata.update(
                  {
                      label[6:]: value
                      for label, value in self.labels.items()
                      if label.startswith("zenml:")
                  }
              )
          metadata["flavor"] = self.flavor_name
  
          return metadata
  
      def get_hydrated_version(self) -> "ComponentResponse":
          """Get the hydrated version of this component.
  
          Returns:
              an instance of the same entity with the metadata field attached.
          """
          from zenml.client import Client
  
          return Client().zen_store.get_stack_component(self.id)
  
      # Body and metadata properties
      @property
      def type(self) -> StackComponentType:
          """The `type` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().type
  
      @property
      def flavor_name(self) -> str:
          """The `flavor_name` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().flavor_name
  
      @property
      def integration(self) -> Optional[str]:
          """The `integration` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().integration
  
      @property
      def logo_url(self) -> Optional[str]:
          """The `logo_url` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().logo_url
  
      @property
      def configuration(self) -> Dict[str, Any]:
          """The `configuration` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().configuration
  
      @property
      def environment(self) -> Dict[str, str]:
          """The `environment` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().environment
  
      @property
      def secrets(self) -> List[UUID]:
          """The `secrets` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().secrets
  
      @property
      def labels(self) -> Optional[Dict[str, Any]]:
          """The `labels` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().labels
  
      @property
      def connector_resource_id(self) -> Optional[str]:
          """The `connector_resource_id` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().connector_resource_id
  
      @property
      def connector(self) -> Optional["ServiceConnectorResponse"]:
          """The `connector` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().connector
  
      @property
      def flavor(self) -> "FlavorResponse":
          """The `flavor` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().flavor
  
  
  # ------------------ Filter Model ------------------
  
  
  class ComponentFilter(UserScopedFilter):
      """Model to enable advanced stack component filtering."""
  
      FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *UserScopedFilter.FILTER_EXCLUDE_FIELDS,
          "scope_type",
          "stack_id",
      ]
      CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *UserScopedFilter.CLI_EXCLUDE_FIELDS,
          "scope_type",
      ]
      scope_type: Optional[str] = Field(
          default=None,
          description="The type to scope this query to.",
      )
      name: Optional[str] = Field(
          default=None,
          description="Name of the stack component",
      )
      flavor: Optional[str] = Field(
          default=None,
          description="Flavor of the stack component",
      )
      type: Optional[str] = Field(
          default=None,
          description="Type of the stack component",
      )
      connector_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Connector linked to the stack component",
          union_mode="left_to_right",
      )
      stack_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Stack of the stack component",
          union_mode="left_to_right",
      )
  
      def set_scope_type(self, component_type: str) -> None:
          """Set the type of component on which to perform the filtering to scope the response.
  
          Args:
              component_type: The type of component to scope the query to.
          """
          self.scope_type = component_type
  
      def generate_filter(
          self, table: Type["AnySchema"]
      ) -> Union["ColumnElement[bool]"]:
          """Generate the filter for the query.
  
          Stack components can be scoped by type to narrow the search.
  
          Args:
              table: The Table that is being queried from.
  
          Returns:
              The filter expression for the query.
          """
          from sqlmodel import and_, or_
  
          from zenml.zen_stores.schemas import (
              StackComponentSchema,
              StackCompositionSchema,
          )
  
          base_filter = super().generate_filter(table)
          if self.scope_type:
              type_filter = getattr(table, "type") == self.scope_type
              return and_(base_filter, type_filter)
  
          if self.stack_id:
              operator = (
                  or_ if self.logical_operator == LogicalOperators.OR else and_
              )
  
              stack_filter = and_(
                  StackCompositionSchema.stack_id == self.stack_id,
                  StackCompositionSchema.component_id == StackComponentSchema.id,
              )
              base_filter = operator(base_filter, stack_filter)
  
          return base_filter
  