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 run template tables."""
  
  from typing import TYPE_CHECKING, Any, List, Optional, Sequence
  from uuid import UUID
  
  from sqlalchemy import Column, String, UniqueConstraint
  from sqlalchemy.dialects.mysql import MEDIUMTEXT
  from sqlalchemy.orm import joinedload, object_session
  from sqlalchemy.sql.base import ExecutableOption
  from sqlmodel import Field, Relationship, col, desc, select
  
  from zenml.constants import MEDIUMTEXT_MAX_LENGTH
  from zenml.enums import TaggableResourceTypes
  from zenml.models import (
      RunTemplateRequest,
      RunTemplateResponse,
      RunTemplateResponseBody,
      RunTemplateResponseMetadata,
      RunTemplateResponseResources,
      RunTemplateUpdate,
  )
  from zenml.utils.time_utils import utc_now
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
  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
  
  if TYPE_CHECKING:
      from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
      from zenml.zen_stores.schemas.pipeline_snapshot_schemas import (
          PipelineSnapshotSchema,
      )
      from zenml.zen_stores.schemas.tag_schemas import TagSchema
  
  
  class RunTemplateSchema(NamedSchema, table=True):
      """SQL Model for run templates."""
  
      __tablename__ = "run_template"
      __table_args__ = (
          UniqueConstraint(
              "name",
              "project_id",
              name="unique_template_name_in_project",
          ),
      )
  
      hidden: bool = Field(
          default=False,
          title="Whether the run template is hidden.",
      )
  
      description: Optional[str] = Field(
          sa_column=Column(
              String(length=MEDIUMTEXT_MAX_LENGTH).with_variant(
                  MEDIUMTEXT, "mysql"
              ),
              nullable=True,
          )
      )
  
      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,
      )
      project_id: UUID = build_foreign_key_field(
          source=__tablename__,
          target=ProjectSchema.__tablename__,
          source_column="project_id",
          target_column="id",
          ondelete="CASCADE",
          nullable=False,
      )
      source_snapshot_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target="pipeline_snapshot",
          source_column="source_snapshot_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
  
      user: Optional["UserSchema"] = Relationship(
          back_populates="run_templates",
      )
      project: "ProjectSchema" = Relationship()
      source_snapshot: Optional["PipelineSnapshotSchema"] = Relationship(
          sa_relationship_kwargs={
              "foreign_keys": "RunTemplateSchema.source_snapshot_id",
          }
      )
  
      tags: List["TagSchema"] = Relationship(
          sa_relationship_kwargs=dict(
              primaryjoin=f"and_(foreign(TagResourceSchema.resource_type)=='{TaggableResourceTypes.RUN_TEMPLATE.value}', foreign(TagResourceSchema.resource_id)==RunTemplateSchema.id)",
              secondary="tag_resource",
              secondaryjoin="TagSchema.id == foreign(TagResourceSchema.tag_id)",
              order_by="TagSchema.name",
              overlaps="tags",
          ),
      )
  
      @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.
          """
          from zenml.zen_stores.schemas import PipelineSnapshotSchema
  
          options = [
              joinedload(jl_arg(RunTemplateSchema.source_snapshot)).joinedload(
                  jl_arg(PipelineSnapshotSchema.build)
              ),
          ]
  
          if include_metadata or include_resources:
              options.extend(
                  [
                      joinedload(
                          jl_arg(RunTemplateSchema.source_snapshot)
                      ).joinedload(jl_arg(PipelineSnapshotSchema.pipeline)),
                      joinedload(
                          jl_arg(RunTemplateSchema.source_snapshot)
                      ).joinedload(
                          jl_arg(PipelineSnapshotSchema.code_reference)
                      ),
                  ]
              )
          if include_metadata:
              options.extend(
                  [
                      joinedload(
                          jl_arg(RunTemplateSchema.source_snapshot)
                      ).joinedload(jl_arg(PipelineSnapshotSchema.stack)),
                      joinedload(
                          jl_arg(RunTemplateSchema.source_snapshot)
                      ).joinedload(jl_arg(PipelineSnapshotSchema.schedule)),
                  ]
              )
  
          if include_resources:
              options.extend(
                  [
                      joinedload(jl_arg(RunTemplateSchema.user)),
                      # joinedload(jl_arg(RunTemplateSchema.tags)),
                  ]
              )
  
          return options
  
      @property
      def latest_run(self) -> Optional["PipelineRunSchema"]:
          """Fetch the latest run for this template.
  
          Raises:
              RuntimeError: If no session for the schema exists.
  
          Returns:
              The latest run for this template.
          """
          from zenml.zen_stores.schemas import (
              PipelineRunSchema,
              PipelineSnapshotSchema,
          )
  
          if session := object_session(self):
              return (
                  session.execute(
                      select(PipelineRunSchema)
                      .join(
                          PipelineSnapshotSchema,
                          col(PipelineSnapshotSchema.id)
                          == col(PipelineRunSchema.snapshot_id),
                      )
                      .where(PipelineSnapshotSchema.template_id == self.id)
                      .order_by(desc(PipelineRunSchema.created))
                      .limit(1)
                  )
                  .scalars()
                  .one_or_none()
              )
          else:
              raise RuntimeError(
                  "Missing DB session to fetch latest run for template."
              )
  
      @classmethod
      def from_request(
          cls,
          request: RunTemplateRequest,
      ) -> "RunTemplateSchema":
          """Create a schema from a request.
  
          Args:
              request: The request to convert.
  
  
          Returns:
              The created schema.
          """
          return cls(
              user_id=request.user,
              project_id=request.project,
              name=request.name,
              description=request.description,
              hidden=request.hidden,
              source_snapshot_id=request.source_snapshot_id,
          )
  
      def update(self, update: RunTemplateUpdate) -> "RunTemplateSchema":
          """Update the schema.
  
          Args:
              update: The update model.
  
          Returns:
              The updated schema.
          """
          for field, value in update.model_dump(
              exclude_unset=True, exclude_none=True
          ).items():
              if field in ["add_tags", "remove_tags"]:
                  # Tags are handled separately
                  continue
              setattr(self, field, value)
  
          self.updated = utc_now()
          return self
  
      def to_model(
          self,
          include_metadata: bool = False,
          include_resources: bool = False,
          **kwargs: Any,
      ) -> RunTemplateResponse:
          """Convert the schema to a response model.
  
          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
  
          Returns:
              Model representing this schema.
          """
          runnable = False
          if (
              self.source_snapshot
              and self.source_snapshot.build
              and not self.source_snapshot.build.is_local
              and self.source_snapshot.build.stack_id
          ):
              runnable = True
  
          body = RunTemplateResponseBody(
              user_id=self.user_id,
              project_id=self.project_id,
              created=self.created,
              updated=self.updated,
              runnable=runnable,
              hidden=self.hidden,
          )
  
          metadata = None
          if include_metadata:
              pipeline_spec = None
              config_template = None
              config_schema = None
  
              if self.source_snapshot:
                  from zenml.zen_stores import template_utils
  
                  source_snapshot_model = self.source_snapshot.to_model(
                      include_metadata=True
                  )
                  pipeline_spec = source_snapshot_model.pipeline_spec
  
                  if (
                      self.source_snapshot.build
                      and self.source_snapshot.build.stack_id
                  ):
                      config_template = template_utils.generate_config_template(
                          snapshot=self.source_snapshot,
                          pipeline_configuration=source_snapshot_model.pipeline_configuration,
                          step_configurations=source_snapshot_model.step_configurations,
                      )
                      config_schema = template_utils.generate_config_schema(
                          snapshot=self.source_snapshot,
                          pipeline_configuration=source_snapshot_model.pipeline_configuration,
                          step_configurations=source_snapshot_model.step_configurations,
                      )
  
              metadata = RunTemplateResponseMetadata(
                  description=self.description,
                  pipeline_spec=pipeline_spec,
                  config_template=config_template,
                  config_schema=config_schema,
              )
  
          resources = None
          if include_resources:
              if self.source_snapshot:
                  pipeline = (
                      self.source_snapshot.pipeline.to_model()
                      if self.source_snapshot.pipeline
                      else None
                  )
                  build = (
                      self.source_snapshot.build.to_model()
                      if self.source_snapshot.build
                      else None
                  )
                  code_reference = (
                      self.source_snapshot.code_reference.to_model()
                      if self.source_snapshot.code_reference
                      else None
                  )
              else:
                  pipeline = None
                  build = None
                  code_reference = None
  
              latest_run = self.latest_run
  
              resources = RunTemplateResponseResources(
                  user=self.user.to_model() if self.user else None,
                  source_snapshot=self.source_snapshot.to_model()
                  if self.source_snapshot
                  else None,
                  pipeline=pipeline,
                  build=build,
                  code_reference=code_reference,
                  tags=[tag.to_model() for tag in self.tags],
                  latest_run_id=latest_run.id if latest_run else None,
                  latest_run_status=latest_run.status if latest_run else None,
              )
  
          return RunTemplateResponse(
              id=self.id,
              name=self.name,
              body=body,
              metadata=metadata,
              resources=resources,
          )
  