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 Pipeline Schedules."""
  
  from datetime import datetime, timedelta
  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 MetadataResourceTypes
  from zenml.models import (
      ScheduleRequest,
      ScheduleResponse,
      ScheduleResponseBody,
      ScheduleResponseMetadata,
      ScheduleResponseResources,
      ScheduleUpdate,
  )
  from zenml.utils.time_utils import utc_now
  from zenml.zen_stores.schemas.base_schemas import NamedSchema
  from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
  from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
  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 (
      RunMetadataInterface,
      jl_arg,
  )
  
  if TYPE_CHECKING:
      from zenml.zen_stores.schemas.pipeline_snapshot_schemas import (
          PipelineSnapshotSchema,
      )
      from zenml.zen_stores.schemas.run_metadata_schemas import (
          RunMetadataSchema,
      )
  
  
  class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
      """SQL Model for schedules."""
  
      __tablename__ = "schedule"
      __table_args__ = (
          UniqueConstraint(
              "name",
              "project_id",
              name="unique_schedule_name_in_project",
          ),
      )
  
      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="schedules")
  
      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="schedules")
  
      pipeline_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target=PipelineSchema.__tablename__,
          source_column="pipeline_id",
          target_column="id",
          ondelete="CASCADE",
          nullable=True,
      )
      pipeline: "PipelineSchema" = Relationship(back_populates="schedules")
      snapshot: Optional["PipelineSnapshotSchema"] = Relationship(
          back_populates="schedule"
      )
  
      orchestrator_id: Optional[UUID] = build_foreign_key_field(
          source=__tablename__,
          target=StackComponentSchema.__tablename__,
          source_column="orchestrator_id",
          target_column="id",
          ondelete="SET NULL",
          nullable=True,
      )
      orchestrator: "StackComponentSchema" = Relationship(
          back_populates="schedules"
      )
  
      run_metadata: List["RunMetadataSchema"] = Relationship(
          sa_relationship_kwargs=dict(
              secondary="run_metadata_resource",
              primaryjoin=f"and_(foreign(RunMetadataResourceSchema.resource_type)=='{MetadataResourceTypes.SCHEDULE.value}', foreign(RunMetadataResourceSchema.resource_id)==ScheduleSchema.id)",
              secondaryjoin="RunMetadataSchema.id==foreign(RunMetadataResourceSchema.run_metadata_id)",
              overlaps="run_metadata",
          ),
      )
  
      active: bool
      cron_expression: Optional[str] = Field(nullable=True)
      start_time: Optional[datetime] = Field(nullable=True)
      end_time: Optional[datetime] = Field(nullable=True)
      interval_second: Optional[float] = Field(nullable=True)
      catchup: bool
      run_once_start_time: Optional[datetime] = Field(nullable=True)
  
      @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_metadata:
          #     options.extend(
          #         [
          #             joinedload(jl_arg(ScheduleSchema.run_metadata)),
          #         ]
          #     )
  
          if include_resources:
              options.extend([joinedload(jl_arg(ScheduleSchema.user))])
  
          return options
  
      @classmethod
      def from_request(
          cls, schedule_request: ScheduleRequest
      ) -> "ScheduleSchema":
          """Create a `ScheduleSchema` from a `ScheduleRequest`.
  
          Args:
              schedule_request: The `ScheduleRequest` to create the schema from.
  
          Returns:
              The created `ScheduleSchema`.
          """
          if schedule_request.interval_second is not None:
              interval_second = schedule_request.interval_second.total_seconds()
          else:
              interval_second = None
          return cls(
              name=schedule_request.name,
              project_id=schedule_request.project,
              user_id=schedule_request.user,
              pipeline_id=schedule_request.pipeline_id,
              orchestrator_id=schedule_request.orchestrator_id,
              active=schedule_request.active,
              cron_expression=schedule_request.cron_expression,
              start_time=schedule_request.start_time,
              end_time=schedule_request.end_time,
              interval_second=interval_second,
              catchup=schedule_request.catchup,
              run_once_start_time=schedule_request.run_once_start_time,
          )
  
      def update(self, schedule_update: ScheduleUpdate) -> "ScheduleSchema":
          """Update a `ScheduleSchema` from a `ScheduleUpdateModel`.
  
          Args:
              schedule_update: The `ScheduleUpdateModel` to update the schema from.
  
          Returns:
              The updated `ScheduleSchema`.
          """
          if schedule_update.name is not None:
              self.name = schedule_update.name
  
          if schedule_update.cron_expression:
              self.cron_expression = schedule_update.cron_expression
  
          self.updated = utc_now()
          return self
  
      def to_model(
          self,
          include_metadata: bool = False,
          include_resources: bool = False,
          **kwargs: Any,
      ) -> ScheduleResponse:
          """Convert a `ScheduleSchema` to a `ScheduleResponseModel`.
  
          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:
              The created `ScheduleResponseModel`.
          """
          if self.interval_second is not None:
              interval_second = timedelta(seconds=self.interval_second)
          else:
              interval_second = None
  
          body = ScheduleResponseBody(
              user_id=self.user_id,
              project_id=self.project_id,
              active=self.active,
              cron_expression=self.cron_expression,
              start_time=self.start_time,
              end_time=self.end_time,
              interval_second=interval_second,
              catchup=self.catchup,
              updated=self.updated,
              created=self.created,
              run_once_start_time=self.run_once_start_time,
          )
          metadata = None
          if include_metadata:
              metadata = ScheduleResponseMetadata(
                  pipeline_id=self.pipeline_id,
                  orchestrator_id=self.orchestrator_id,
                  run_metadata=self.fetch_metadata(),
              )
  
          resources = None
          if include_resources:
              resources = ScheduleResponseResources(
                  user=self.user.to_model() if self.user else None,
              )
  
          return ScheduleResponse(
              id=self.id,
              name=self.name,
              body=body,
              metadata=metadata,
              resources=resources,
          )
  