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.
  """Models representing pipeline templates."""
  
  from typing import (
      TYPE_CHECKING,
      Any,
      ClassVar,
      Dict,
      List,
      Optional,
      Type,
      TypeVar,
      Union,
  )
  from uuid import UUID
  
  from pydantic import Field
  
  from zenml.config.pipeline_spec import PipelineSpec
  from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
  from zenml.enums import ExecutionStatus
  from zenml.models.v2.base.base import BaseUpdate
  from zenml.models.v2.base.scoped import (
      ProjectScopedFilter,
      ProjectScopedRequest,
      ProjectScopedResponse,
      ProjectScopedResponseBody,
      ProjectScopedResponseMetadata,
      ProjectScopedResponseResources,
      TaggableFilter,
  )
  from zenml.models.v2.core.code_reference import (
      CodeReferenceResponse,
  )
  from zenml.models.v2.core.pipeline import PipelineResponse
  from zenml.models.v2.core.pipeline_build import (
      PipelineBuildResponse,
  )
  from zenml.models.v2.core.pipeline_snapshot import (
      PipelineSnapshotResponse,
  )
  from zenml.models.v2.core.tag import TagResponse
  
  if TYPE_CHECKING:
      from sqlalchemy.sql.elements import ColumnElement
  
      from zenml.zen_stores.schemas.base_schemas import BaseSchema
  
      AnySchema = TypeVar("AnySchema", bound=BaseSchema)
  
  
  # ------------------ Request Model ------------------
  
  
  class RunTemplateRequest(ProjectScopedRequest):
      """Request model for run templates."""
  
      name: str = Field(
          title="The name of the run template.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
      description: Optional[str] = Field(
          default=None,
          title="The description of the run template.",
          max_length=TEXT_FIELD_MAX_LENGTH,
      )
      source_snapshot_id: UUID = Field(
          title="The snapshot that should be the base of the created template."
      )
      hidden: bool = Field(
          default=False,
          title="Whether the run template is hidden.",
      )
      tags: Optional[List[str]] = Field(
          default=None,
          title="Tags of the run template.",
      )
  
  
  # ------------------ Update Model ------------------
  
  
  class RunTemplateUpdate(BaseUpdate):
      """Run template update model."""
  
      name: Optional[str] = Field(
          default=None,
          title="The name of the run template.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
      description: Optional[str] = Field(
          default=None,
          title="The description of the run template.",
          max_length=TEXT_FIELD_MAX_LENGTH,
      )
      hidden: Optional[bool] = Field(
          default=None,
          title="Whether the run template is hidden.",
      )
      add_tags: Optional[List[str]] = Field(
          default=None, title="New tags to add to the run template."
      )
      remove_tags: Optional[List[str]] = Field(
          default=None, title="Tags to remove from the run template."
      )
  
  
  # ------------------ Response Model ------------------
  
  
  class RunTemplateResponseBody(ProjectScopedResponseBody):
      """Response body for run templates."""
  
      runnable: bool = Field(
          title="If a run can be started from the template.",
      )
      hidden: bool = Field(
          default=False,
          title="Whether the run template is hidden.",
      )
  
  
  class RunTemplateResponseMetadata(ProjectScopedResponseMetadata):
      """Response metadata for run templates."""
  
      description: Optional[str] = Field(
          default=None,
          title="The description of the run template.",
      )
      pipeline_spec: Optional[PipelineSpec] = Field(
          default=None, title="The spec of the pipeline."
      )
      config_template: Optional[Dict[str, Any]] = Field(
          default=None, title="Run configuration template."
      )
      config_schema: Optional[Dict[str, Any]] = Field(
          default=None, title="Run configuration schema."
      )
  
  
  class RunTemplateResponseResources(ProjectScopedResponseResources):
      """All resource models associated with the run template."""
  
      source_snapshot: Optional[PipelineSnapshotResponse] = Field(
          default=None,
          title="The snapshot that is the source of the template.",
      )
      pipeline: Optional[PipelineResponse] = Field(
          default=None, title="The pipeline associated with the template."
      )
      build: Optional[PipelineBuildResponse] = Field(
          default=None,
          title="The pipeline build associated with the template.",
      )
      code_reference: Optional[CodeReferenceResponse] = Field(
          default=None,
          title="The code reference associated with the template.",
      )
      tags: List[TagResponse] = Field(
          title="Tags associated with the run template.",
      )
      latest_run_id: Optional[UUID] = Field(
          default=None,
          title="The ID of the latest run of the run template.",
      )
      latest_run_status: Optional[ExecutionStatus] = Field(
          default=None,
          title="The status of the latest run of the run template.",
      )
  
  
  class RunTemplateResponse(
      ProjectScopedResponse[
          RunTemplateResponseBody,
          RunTemplateResponseMetadata,
          RunTemplateResponseResources,
      ]
  ):
      """Response model for run templates."""
  
      name: str = Field(
          title="The name of the run template.",
          max_length=STR_FIELD_MAX_LENGTH,
      )
  
      def get_hydrated_version(self) -> "RunTemplateResponse":
          """Return the hydrated version of this run template.
  
          Returns:
              The hydrated run template.
          """
          from zenml.client import Client
  
          return Client().zen_store.get_run_template(
              template_id=self.id, hydrate=True
          )
  
      # Body and metadata properties
      @property
      def runnable(self) -> bool:
          """The `runnable` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().runnable
  
      @property
      def hidden(self) -> bool:
          """The `hidden` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().hidden
  
      @property
      def latest_run_id(self) -> Optional[UUID]:
          """The `latest_run_id` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().latest_run_id
  
      @property
      def latest_run_status(self) -> Optional[ExecutionStatus]:
          """The `latest_run_status` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().latest_run_status
  
      @property
      def description(self) -> Optional[str]:
          """The `description` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().description
  
      @property
      def pipeline_spec(self) -> Optional[PipelineSpec]:
          """The `pipeline_spec` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().pipeline_spec
  
      @property
      def config_template(self) -> Optional[Dict[str, Any]]:
          """The `config_template` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().config_template
  
      @property
      def config_schema(self) -> Optional[Dict[str, Any]]:
          """The `config_schema` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().config_schema
  
      @property
      def source_snapshot(self) -> Optional[PipelineSnapshotResponse]:
          """The `source_snapshot` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().source_snapshot
  
      @property
      def pipeline(self) -> Optional[PipelineResponse]:
          """The `pipeline` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().pipeline
  
      @property
      def build(self) -> Optional[PipelineBuildResponse]:
          """The `build` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().build
  
      @property
      def code_reference(self) -> Optional[CodeReferenceResponse]:
          """The `code_reference` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().code_reference
  
      @property
      def tags(self) -> List[TagResponse]:
          """The `tags` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().tags
  
  
  # ------------------ Filter Model ------------------
  
  
  class RunTemplateFilter(ProjectScopedFilter, TaggableFilter):
      """Model for filtering of run templates."""
  
      FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
          *TaggableFilter.FILTER_EXCLUDE_FIELDS,
          "code_repository_id",
          "stack_id",
          "build_id",
          "pipeline_id",
          "pipeline",
          "stack",
          "hidden",
      ]
      CUSTOM_SORTING_OPTIONS = [
          *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
          *TaggableFilter.CUSTOM_SORTING_OPTIONS,
      ]
      CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
          *TaggableFilter.CLI_EXCLUDE_FIELDS,
      ]
  
      name: Optional[str] = Field(
          default=None,
          description="Name of the run template.",
      )
      hidden: Optional[bool] = Field(
          default=None,
          description="Whether the run template is hidden.",
      )
      pipeline_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Pipeline associated with the template.",
          union_mode="left_to_right",
      )
      build_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Build associated with the template.",
          union_mode="left_to_right",
      )
      stack_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Stack associated with the template.",
          union_mode="left_to_right",
      )
      code_repository_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Code repository associated with the template.",
          union_mode="left_to_right",
      )
      pipeline: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Name/ID of the pipeline associated with the template.",
      )
      stack: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Name/ID of the stack associated with the template.",
      )
  
      def get_custom_filters(
          self, table: Type["AnySchema"]
      ) -> List["ColumnElement[bool]"]:
          """Get custom filters.
  
          Args:
              table: The query table.
  
          Returns:
              A list of custom filters.
          """
          custom_filters = super().get_custom_filters(table)
  
          from sqlmodel import and_, col
  
          from zenml.zen_stores.schemas import (
              CodeReferenceSchema,
              PipelineSchema,
              PipelineSnapshotSchema,
              RunTemplateSchema,
              StackSchema,
          )
  
          if self.hidden is not None:
              custom_filters.append(
                  col(RunTemplateSchema.hidden).is_(self.hidden)
              )
  
          if self.code_repository_id:
              code_repo_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.code_reference_id
                  == CodeReferenceSchema.id,
                  CodeReferenceSchema.code_repository_id
                  == self.code_repository_id,
              )
              custom_filters.append(code_repo_filter)
  
          if self.stack_id:
              stack_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.stack_id == self.stack_id,
              )
              custom_filters.append(stack_filter)
  
          if self.build_id:
              build_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.build_id == self.build_id,
              )
              custom_filters.append(build_filter)
  
          if self.pipeline_id:
              pipeline_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.pipeline_id == self.pipeline_id,
              )
              custom_filters.append(pipeline_filter)
  
          if self.pipeline:
              pipeline_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.pipeline_id == PipelineSchema.id,
                  self.generate_name_or_id_query_conditions(
                      value=self.pipeline,
                      table=PipelineSchema,
                  ),
              )
              custom_filters.append(pipeline_filter)
  
          if self.stack:
              stack_filter = and_(
                  RunTemplateSchema.source_snapshot_id
                  == PipelineSnapshotSchema.id,
                  PipelineSnapshotSchema.stack_id == StackSchema.id,
                  self.generate_name_or_id_query_conditions(
                      value=self.stack,
                      table=StackSchema,
                  ),
              )
              custom_filters.append(stack_filter)
  
          return custom_filters
  