content: #  Copyright (c) ZenML GmbH 2024. 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.
  """SQLModel implementation of service table."""
  
  import base64
  import json
  from typing import Any, Optional, Sequence
  from uuid import UUID
  
  from pydantic import ConfigDict
  from sqlalchemy import TEXT, Column
  from sqlalchemy.orm import joinedload
  from sqlalchemy.sql.base import ExecutableOption
  from sqlmodel import Field, Relationship
  
  from zenml.models.v2.core.service import (
      ServiceRequest,
      ServiceResponse,
      ServiceResponseBody,
      ServiceResponseMetadata,
      ServiceResponseResources,
      ServiceUpdate,
  )
  from zenml.utils.dict_utils import dict_to_bytes
  from zenml.utils.time_utils import utc_now
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
  from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
  from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
  from zenml.zen_stores.schemas.project_schemas import ProjectSchema
  from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
  from zenml.zen_stores.schemas.user_schemas import UserSchema
  from zenml.zen_stores.schemas.utils import jl_arg
  
  
  class ServiceSchema(NamedSchema, table=True):
      """SQL Model for service."""
  
      __tablename__ = "service"
  
      project_id: UUID = build_foreign_key_field(
          source=__tablename__,
          target=ProjectSchema.__tablename__,
          source_column="project_id",
          target_column="id",
          ondelete="CASCADE",
          nullable=False,
      )
      project: "ProjectSchema" = Relationship(back_populates="services")
  
      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="services")
      service_source: Optional[str] = Field(
          sa_column=Column(TEXT, nullable=True)
      )
      service_type: str = Field(sa_column=Column(TEXT, nullable=False))
      type: str = Field(sa_column=Column(TEXT, nullable=False))
      flavor: str = Field(sa_column=Column(TEXT, nullable=False))
      admin_state: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
      state: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
      labels: Optional[bytes]
      config: bytes
      status: Optional[bytes]
      endpoint: Optional[bytes]
      prediction_url: Optional[str] = Field(
          sa_column=Column(TEXT, nullable=True)
      )
      health_check_url: Optional[str] = Field(
          sa_column=Column(TEXT, nullable=True)
      )
      pipeline_name: Optional[str] = Field(sa_column=Column(TEXT, nullable=True))
      pipeline_step_name: Optional[str] = Field(
          sa_column=Column(TEXT, nullable=True)
      )
      model_version_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target=ModelVersionSchema.__tablename__,
          source_column="model_version_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
      model_version: Optional["ModelVersionSchema"] = Relationship(
          back_populates="services",
      )
      pipeline_run_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target="pipeline_run",
          source_column="pipeline_run_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
      pipeline_run: Optional["PipelineRunSchema"] = Relationship(
          back_populates="services",
      )
  
      # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
      #  fields defined under base models. If not handled, this raises a warning.
      #  It is possible to suppress this warning message with the following
      #  configuration, however the ultimate solution is to rename these fields.
      #  Even though they do not cause any problems right now, if we are not
      #  careful we might overwrite some fields protected by pydantic.
      model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
  
      @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 = []
  
          if include_resources:
              options.extend(
                  [
                      joinedload(jl_arg(ServiceSchema.user)),
                      joinedload(jl_arg(ServiceSchema.model_version)),
                      joinedload(jl_arg(ServiceSchema.pipeline_run)),
                  ]
              )
  
          return options
  
      def to_model(
          self,
          include_metadata: bool = False,
          include_resources: bool = False,
          **kwargs: Any,
      ) -> ServiceResponse:
          """Convert an `ServiceSchema` to an `ServiceResponse`.
  
          Args:
              include_metadata: Whether to include metadata in the response.
              include_resources: Whether to include resources in the response.
              kwargs: Additional keyword arguments.
  
          Returns:
              The created `ServiceResponse`.
          """
          body = ServiceResponseBody(
              user_id=self.user_id,
              project_id=self.project_id,
              created=self.created,
              updated=self.updated,
              service_type=json.loads(self.service_type),
              labels=json.loads(base64.b64decode(self.labels).decode())
              if self.labels
              else None,
              state=self.state,
          )
          metadata = None
          if include_metadata:
              metadata = ServiceResponseMetadata(
                  service_source=self.service_source,
                  config=json.loads(base64.b64decode(self.config).decode()),
                  status=json.loads(base64.b64decode(self.status).decode())
                  if self.status
                  else None,
                  endpoint=json.loads(base64.b64decode(self.endpoint).decode())
                  if self.endpoint
                  else None,
                  admin_state=self.admin_state or None,
                  prediction_url=self.prediction_url or None,
                  health_check_url=self.health_check_url,
              )
          resources = None
          if include_resources:
              resources = ServiceResponseResources(
                  user=self.user.to_model() if self.user else None,
                  model_version=self.model_version.to_model()
                  if self.model_version
                  else None,
                  pipeline_run=self.pipeline_run.to_model()
                  if self.pipeline_run
                  else None,
              )
          return ServiceResponse(
              id=self.id,
              name=self.name,
              body=body,
              metadata=metadata,
              resources=resources,
          )
  
      def update(
          self,
          update: ServiceUpdate,
      ) -> "ServiceSchema":
          """Updates a `ServiceSchema` from a `ServiceUpdate`.
  
          Args:
              update: The `ServiceUpdate` to update from.
  
          Returns:
              The updated `ServiceSchema`.
          """
          for field, value in update.model_dump(
              exclude_unset=True, exclude_none=True
          ).items():
              if field == "labels":
                  self.labels = (
                      dict_to_bytes(update.labels) if update.labels else None
                  )
              elif field == "status":
                  self.status = (
                      dict_to_bytes(update.status) if update.status else None
                  )
                  self.state = (
                      update.status.get("state") if update.status else None
                  )
              elif field == "endpoint":
                  self.endpoint = (
                      dict_to_bytes(update.endpoint) if update.endpoint else None
                  )
              else:
                  setattr(self, field, value)
          self.updated = utc_now()
          return self
  
      @classmethod
      def from_request(
          cls, service_request: "ServiceRequest"
      ) -> "ServiceSchema":
          """Convert a `ServiceRequest` to a `ServiceSchema`.
  
          Args:
              service_request: The request model to convert.
  
          Returns:
              The converted schema.
          """
          return cls(
              name=service_request.name,
              project_id=service_request.project,
              user_id=service_request.user,
              service_source=service_request.service_source,
              service_type=service_request.service_type.model_dump_json(),
              type=service_request.service_type.type,
              flavor=service_request.service_type.flavor,
              admin_state=service_request.admin_state,
              config=dict_to_bytes(service_request.config),
              labels=dict_to_bytes(service_request.labels)
              if service_request.labels
              else None,
              status=dict_to_bytes(service_request.status)
              if service_request.status
              else None,
              endpoint=dict_to_bytes(service_request.endpoint)
              if service_request.endpoint
              else None,
              state=service_request.status.get("state")
              if service_request.status
              else None,
              model_version_id=service_request.model_version_id,
              pipeline_run_id=service_request.pipeline_run_id,
              prediction_url=service_request.prediction_url,
              health_check_url=service_request.health_check_url,
              pipeline_name=service_request.config.get("pipeline_name"),
              pipeline_step_name=service_request.config.get(
                  "pipeline_step_name"
              ),
          )
  