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.
  """SQL Model Implementations for Stack Components."""
  
  import base64
  import json
  from typing import TYPE_CHECKING, Any, List, Optional, Sequence
  from uuid import UUID
  
  from sqlalchemy import UniqueConstraint
  from sqlalchemy.orm import joinedload
  from sqlalchemy.sql.base import ExecutableOption
  from sqlmodel import Field, Relationship
  
  from zenml.enums import SecretResourceTypes, StackComponentType
  from zenml.models import (
      ComponentRequest,
      ComponentResponse,
      ComponentResponseBody,
      ComponentResponseMetadata,
      ComponentResponseResources,
      ComponentUpdate,
  )
  from zenml.utils.time_utils import utc_now
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
  from zenml.zen_stores.schemas.service_connector_schemas import (
      ServiceConnectorSchema,
  )
  from zenml.zen_stores.schemas.stack_schemas import StackCompositionSchema
  from zenml.zen_stores.schemas.user_schemas import UserSchema
  from zenml.zen_stores.schemas.utils import jl_arg
  
  if TYPE_CHECKING:
      from zenml.zen_stores.schemas.flavor_schemas import FlavorSchema
      from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
      from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
      from zenml.zen_stores.schemas.secret_schemas import SecretSchema
      from zenml.zen_stores.schemas.stack_schemas import StackSchema
  
  
  class StackComponentSchema(NamedSchema, table=True):
      """SQL Model for stack components."""
  
      __tablename__ = "stack_component"
      __table_args__ = (
          UniqueConstraint(
              "name",
              "type",
              name="unique_component_name_and_type",
          ),
      )
  
      type: str
      flavor: str
      configuration: bytes
      labels: Optional[bytes]
      environment: Optional[bytes] = Field(default=None)
  
      user_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target=UserSchema.__tablename__,
          source_column="user_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
      user: Optional["UserSchema"] = Relationship(back_populates="components")
  
      stacks: List["StackSchema"] = Relationship(
          back_populates="components", link_model=StackCompositionSchema
      )
      schedules: List["ScheduleSchema"] = Relationship(
          back_populates="orchestrator",
      )
  
      run_metadata: List["RunMetadataSchema"] = Relationship(
          back_populates="stack_component",
      )
      flavor_schema: Optional["FlavorSchema"] = Relationship(
          sa_relationship_kwargs={
              "primaryjoin": "and_(foreign(StackComponentSchema.flavor) == FlavorSchema.name, foreign(StackComponentSchema.type) == FlavorSchema.type)",
              "lazy": "joined",
              "uselist": False,
          },
      )
  
      connector_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target=ServiceConnectorSchema.__tablename__,
          source_column="connector_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
      connector: Optional["ServiceConnectorSchema"] = Relationship(
          back_populates="components"
      )
  
      connector_resource_id: Optional[str]
      secrets: List["SecretSchema"] = Relationship(
          sa_relationship_kwargs=dict(
              primaryjoin=f"and_(foreign(SecretResourceSchema.resource_type)=='{SecretResourceTypes.STACK_COMPONENT.value}', foreign(SecretResourceSchema.resource_id)==StackComponentSchema.id)",
              secondary="secret_resource",
              secondaryjoin="SecretSchema.id == foreign(SecretResourceSchema.secret_id)",
              order_by="SecretSchema.name",
              overlaps="secrets",
          ),
      )
  
      @classmethod
      def get_query_options(
          cls,
          include_metadata: bool = False,
          include_resources: bool = False,
          **kwargs: Any,
      ) -> Sequence[ExecutableOption]:
          """Get the query options for the schema.
  
          Args:
              include_metadata: Whether metadata will be included when converting
                  the schema to a model.
              include_resources: Whether resources will be included when
                  converting the schema to a model.
              **kwargs: Keyword arguments to allow schema specific logic
  
          Returns:
              A list of query options.
          """
          options = [
              joinedload(jl_arg(StackComponentSchema.flavor_schema)),
          ]
  
          if include_metadata:
              options.extend(
                  [joinedload(jl_arg(StackComponentSchema.connector))]
              )
  
          if include_resources:
              options.extend(
                  [
                      joinedload(jl_arg(StackComponentSchema.user)),
                  ]
              )
  
          return options
  
      @classmethod
      def from_request(
          cls,
          request: "ComponentRequest",
          service_connector: Optional[ServiceConnectorSchema] = None,
      ) -> "StackComponentSchema":
          """Create a component schema from a request.
  
          Args:
              request: The request from which to create the component.
              service_connector: Optional service connector to link to the
                  component.
  
          Returns:
              The component schema.
          """
          return cls(
              name=request.name,
              user_id=request.user,
              type=request.type,
              flavor=request.flavor,
              configuration=base64.b64encode(
                  json.dumps(request.configuration).encode("utf-8")
              ),
              labels=base64.b64encode(
                  json.dumps(request.labels).encode("utf-8")
              ),
              environment=base64.b64encode(
                  json.dumps(request.environment).encode("utf-8")
              ),
              connector=service_connector,
              connector_resource_id=request.connector_resource_id,
          )
  
      def update(
          self, component_update: "ComponentUpdate"
      ) -> "StackComponentSchema":
          """Updates a `StackComponentSchema` from a `ComponentUpdate`.
  
          Args:
              component_update: The `ComponentUpdate` to update from.
  
          Returns:
              The updated `StackComponentSchema`.
          """
          for field, value in component_update.model_dump(
              exclude_unset=True,
              exclude={"user", "connector", "add_secrets", "remove_secrets"},
          ).items():
              if field == "configuration":
                  self.configuration = base64.b64encode(
                      json.dumps(component_update.configuration).encode("utf-8")
                  )
              elif field == "labels":
                  self.labels = base64.b64encode(
                      json.dumps(component_update.labels).encode("utf-8")
                  )
              elif field == "environment":
                  self.environment = base64.b64encode(
                      json.dumps(component_update.environment).encode("utf-8")
                  )
              else:
                  setattr(self, field, value)
  
          self.updated = utc_now()
          return self
  
      def to_model(
          self,
          include_metadata: bool = False,
          include_resources: bool = False,
          **kwargs: Any,
      ) -> "ComponentResponse":
          """Creates a `ComponentModel` from an instance of a `StackComponentSchema`.
  
          Args:
              include_metadata: Whether the metadata will be filled.
              include_resources: Whether the resources will be filled.
              **kwargs: Keyword arguments to allow schema specific logic
  
          Raises:
              RuntimeError: If the flavor for the component is missing in the DB.
  
          Returns:
              A `ComponentModel`
          """
          body = ComponentResponseBody(
              user_id=self.user_id,
              type=StackComponentType(self.type),
              flavor_name=self.flavor,
              created=self.created,
              updated=self.updated,
              logo_url=self.flavor_schema.logo_url
              if self.flavor_schema
              else None,
              integration=self.flavor_schema.integration
              if self.flavor_schema
              else None,
          )
          metadata = None
          if include_metadata:
              environment = None
              if self.environment:
                  environment = json.loads(
                      base64.b64decode(self.environment).decode()
                  )
              metadata = ComponentResponseMetadata(
                  configuration=json.loads(
                      base64.b64decode(self.configuration).decode()
                  ),
                  labels=json.loads(base64.b64decode(self.labels).decode())
                  if self.labels
                  else None,
                  environment=environment or {},
                  connector_resource_id=self.connector_resource_id,
                  connector=self.connector.to_model()
                  if self.connector
                  else None,
                  secrets=[secret.id for secret in self.secrets],
              )
          resources = None
          if include_resources:
              if not self.flavor_schema:
                  raise RuntimeError(
                      f"Missing flavor {self.flavor} for component {self.name}."
                  )
  
              resources = ComponentResponseResources(
                  user=self.user.to_model() if self.user else None,
                  flavor=self.flavor_schema.to_model(),
              )
          return ComponentResponse(
              id=self.id,
              name=self.name,
              body=body,
              metadata=metadata,
              resources=resources,
          )
  