content: #  Copyright (c) ZenML GmbH 2023. 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 artifact versions."""
  
  from typing import (
      TYPE_CHECKING,
      Any,
      ClassVar,
      Dict,
      List,
      Optional,
      Type,
      TypeVar,
      Union,
  )
  from uuid import UUID
  
  from pydantic import (
      ConfigDict,
      Field,
      field_validator,
      model_validator,
  )
  
  from zenml.config.source import Source, SourceWithValidator
  from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
  from zenml.enums import ArtifactSaveType, ArtifactType
  from zenml.logger import get_logger
  from zenml.metadata.metadata_types import MetadataType
  from zenml.models.v2.base.base import BaseUpdate
  from zenml.models.v2.base.filter import FilterGenerator
  from zenml.models.v2.base.scoped import (
      ProjectScopedFilter,
      ProjectScopedRequest,
      ProjectScopedResponse,
      ProjectScopedResponseBody,
      ProjectScopedResponseMetadata,
      ProjectScopedResponseResources,
      RunMetadataFilterMixin,
      TaggableFilter,
  )
  from zenml.models.v2.core.artifact import ArtifactResponse
  from zenml.models.v2.core.tag import TagResponse
  
  if TYPE_CHECKING:
      from sqlalchemy.sql.elements import ColumnElement
  
      from zenml.models.v2.core.artifact_visualization import (
          ArtifactVisualizationRequest,
          ArtifactVisualizationResponse,
      )
      from zenml.models.v2.core.pipeline_run import PipelineRunResponse
      from zenml.models.v2.core.step_run import StepRunResponse
      from zenml.zen_stores.schemas.base_schemas import BaseSchema
  
      AnySchema = TypeVar("AnySchema", bound=BaseSchema)
  
  
  logger = get_logger(__name__)
  
  # ------------------ Request Model ------------------
  
  
  class ArtifactVersionRequest(ProjectScopedRequest):
      """Request model for artifact versions."""
  
      artifact_id: Optional[UUID] = Field(
          default=None,
          title="ID of the artifact to which this version belongs.",
      )
      artifact_name: Optional[str] = Field(
          default=None,
          title="Name of the artifact to which this version belongs.",
      )
      version: Optional[Union[int, str]] = Field(
          default=None, title="Version of the artifact."
      )
      has_custom_name: bool = Field(
          title="Whether the name is custom (True) or auto-generated (False).",
          default=False,
      )
      type: ArtifactType = Field(title="Type of the artifact.")
      artifact_store_id: Optional[UUID] = Field(
          title="ID of the artifact store in which this artifact is stored.",
          default=None,
      )
      uri: str = Field(
          title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
      )
      materializer: SourceWithValidator = Field(
          title="Materializer class to use for this artifact.",
      )
      data_type: SourceWithValidator = Field(
          title="Data type of the artifact.",
      )
      content_hash: Optional[str] = Field(
          title="The content hash of the artifact version.",
          default=None,
          max_length=STR_FIELD_MAX_LENGTH,
      )
      tags: Optional[List[str]] = Field(
          title="Tags of the artifact.",
          description="Should be a list of plain strings, e.g., ['tag1', 'tag2']",
          default=None,
      )
      visualizations: Optional[List["ArtifactVisualizationRequest"]] = Field(
          default=None, title="Visualizations of the artifact."
      )
      save_type: ArtifactSaveType = Field(
          title="The save type of the artifact version.",
      )
      metadata: Optional[Dict[str, MetadataType]] = Field(
          default=None, title="Metadata of the artifact version."
      )
      item_count: Optional[int] = Field(
          title="The number of items in the artifact version if it is "
          "sequence-like. This should only be set for artifacts that can be "
          "split into parts, like lists or arrays.",
          default=None,
      )
  
      @field_validator("version")
      @classmethod
      def str_field_max_length_check(cls, value: Any) -> Any:
          """Checks if the length of the value exceeds the maximum str length.
  
          Args:
              value: the value set in the field
  
          Returns:
              the value itself.
  
          Raises:
              AssertionError: if the length of the field is longer than the
                  maximum threshold.
          """
          assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
              "The length of the value for this field can not "
              f"exceed {STR_FIELD_MAX_LENGTH}"
          )
          return value
  
      @model_validator(mode="after")
      def _validate_request(self) -> "ArtifactVersionRequest":
          """Validate the request values.
  
          Raises:
              ValueError: If the request is invalid.
  
          Returns:
              The validated request.
          """
          if self.artifact_id and self.artifact_name:
              raise ValueError(
                  "Only one of artifact_name and artifact_id can be set."
              )
  
          if not (self.artifact_id or self.artifact_name):
              raise ValueError(
                  "Either artifact_name or artifact_id must be set."
              )
  
          return self
  
  
  # ------------------ Update Model ------------------
  
  
  class ArtifactVersionUpdate(BaseUpdate):
      """Artifact version update model."""
  
      name: Optional[str] = None
      add_tags: Optional[List[str]] = None
      remove_tags: Optional[List[str]] = None
  
  
  # ------------------ Response Model ------------------
  
  
  class ArtifactVersionResponseBody(ProjectScopedResponseBody):
      """Response body for artifact versions."""
  
      artifact: ArtifactResponse = Field(
          title="Artifact to which this version belongs."
      )
      version: str = Field(title="Version of the artifact.")
      uri: str = Field(
          title="URI of the artifact.", max_length=TEXT_FIELD_MAX_LENGTH
      )
      type: ArtifactType = Field(title="Type of the artifact.")
      materializer: SourceWithValidator = Field(
          title="Materializer class to use for this artifact.",
      )
      data_type: SourceWithValidator = Field(
          title="Data type of the artifact.",
      )
      save_type: ArtifactSaveType = Field(
          title="The save type of the artifact version.",
      )
      artifact_store_id: Optional[UUID] = Field(
          title="ID of the artifact store in which this artifact is stored.",
          default=None,
      )
      content_hash: Optional[str] = Field(
          title="The content hash of the artifact version.",
          default=None,
      )
      item_count: Optional[int] = Field(
          title="The number of items in the artifact version if it is "
          "sequence-like. This should only be set for artifacts that can be "
          "split into parts, like lists or arrays.",
          default=None,
      )
  
      @field_validator("version")
      @classmethod
      def str_field_max_length_check(cls, value: Any) -> Any:
          """Checks if the length of the value exceeds the maximum str length.
  
          Args:
              value: the value set in the field
  
          Returns:
              the value itself.
  
          Raises:
              AssertionError: if the length of the field is longer than the
                  maximum threshold.
          """
          assert len(str(value)) < STR_FIELD_MAX_LENGTH, (
              "The length of the value for this field can not "
              f"exceed {STR_FIELD_MAX_LENGTH}"
          )
          return value
  
  
  class ArtifactVersionResponseMetadata(ProjectScopedResponseMetadata):
      """Response metadata for artifact versions."""
  
      visualizations: Optional[List["ArtifactVisualizationResponse"]] = Field(
          default=None, title="Visualizations of the artifact."
      )
      run_metadata: Dict[str, MetadataType] = Field(
          default={}, title="Metadata of the artifact."
      )
  
  
  class ArtifactVersionResponseResources(ProjectScopedResponseResources):
      """Class for all resource models associated with the artifact version entity."""
  
      tags: List[TagResponse] = Field(
          title="Tags associated with the artifact version.",
      )
      producer_step_run_id: Optional[UUID] = Field(
          title="ID of the step run that produced this artifact.",
          default=None,
      )
      producer_pipeline_run_id: Optional[UUID] = Field(
          title="The ID of the pipeline run that generated this artifact version.",
          default=None,
      )
  
  
  class ArtifactVersionResponse(
      ProjectScopedResponse[
          ArtifactVersionResponseBody,
          ArtifactVersionResponseMetadata,
          ArtifactVersionResponseResources,
      ]
  ):
      """Response model for artifact versions."""
  
      def get_hydrated_version(self) -> "ArtifactVersionResponse":
          """Get the hydrated version of this artifact version.
  
          Returns:
              an instance of the same entity with the metadata field attached.
          """
          from zenml.client import Client
  
          return Client().zen_store.get_artifact_version(self.id)
  
      # Body and metadata properties
      @property
      def artifact(self) -> "ArtifactResponse":
          """The `artifact` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().artifact
  
      @property
      def version(self) -> str:
          """The `version` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().version
  
      @property
      def uri(self) -> str:
          """The `uri` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().uri
  
      @property
      def type(self) -> ArtifactType:
          """The `type` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().type
  
      @property
      def content_hash(self) -> Optional[str]:
          """The `content_hash` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().content_hash
  
      @property
      def tags(self) -> List[TagResponse]:
          """The `tags` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().tags
  
      @property
      def producer_pipeline_run_id(self) -> Optional[UUID]:
          """The `producer_pipeline_run_id` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().producer_pipeline_run_id
  
      @property
      def save_type(self) -> ArtifactSaveType:
          """The `save_type` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().save_type
  
      @property
      def artifact_store_id(self) -> Optional[UUID]:
          """The `artifact_store_id` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().artifact_store_id
  
      @property
      def producer_step_run_id(self) -> Optional[UUID]:
          """The `producer_step_run_id` property.
  
          Returns:
              the value of the property.
          """
          return self.get_resources().producer_step_run_id
  
      @property
      def visualizations(
          self,
      ) -> Optional[List["ArtifactVisualizationResponse"]]:
          """The `visualizations` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().visualizations
  
      @property
      def run_metadata(self) -> Dict[str, MetadataType]:
          """The `metadata` property.
  
          Returns:
              the value of the property.
          """
          return self.get_metadata().run_metadata
  
      @property
      def materializer(self) -> Source:
          """The `materializer` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().materializer
  
      @property
      def data_type(self) -> Source:
          """The `data_type` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().data_type
  
      @property
      def item_count(self) -> Optional[int]:
          """The `item_count` property.
  
          Returns:
              the value of the property.
          """
          return self.get_body().item_count
  
      # Helper methods
      @property
      def name(self) -> str:
          """The `name` property.
  
          Returns:
              the value of the property.
          """
          return self.artifact.name
  
      @property
      def step(self) -> "StepRunResponse":
          """Get the step that produced this artifact.
  
          Returns:
              The step that produced this artifact.
          """
          from zenml.artifacts.utils import get_producer_step_of_artifact
  
          return get_producer_step_of_artifact(self)
  
      @property
      def run(self) -> "PipelineRunResponse":
          """Get the pipeline run that produced this artifact.
  
          Returns:
              The pipeline run that produced this artifact.
          """
          from zenml.client import Client
  
          return Client().get_pipeline_run(self.step.pipeline_run_id)
  
      def load(self, disable_cache: bool = False) -> Any:
          """Materializes (loads) the data stored in this artifact.
  
          Args:
              disable_cache: Whether to disable the artifact cache.
  
          Returns:
              The materialized data.
          """
          from zenml.artifacts.in_memory_cache import InMemoryArtifactCache
          from zenml.artifacts.utils import load_artifact_from_response
  
          cache = InMemoryArtifactCache.get()
  
          if cache and (data := cache.get_artifact_data(self.id)):
              logger.debug("Returning artifact data (%s) from cache", self.id)
              return data
  
          data = load_artifact_from_response(self)
          if cache and not disable_cache:
              cache.set_artifact_data(self.id, data)
          return data
  
      def download_files(self, path: str, overwrite: bool = False) -> None:
          """Downloads data for an artifact with no materializing.
  
          Any artifacts will be saved as a zip file to the given path.
  
          Args:
              path: The path to save the binary data to.
              overwrite: Whether to overwrite the file if it already exists.
  
          Raises:
              ValueError: If the path does not end with '.zip'.
          """
          if not path.endswith(".zip"):
              raise ValueError(
                  "The path should end with '.zip' to save the binary data."
              )
          from zenml.artifacts.utils import (
              download_artifact_files_from_response,
          )
  
          download_artifact_files_from_response(
              self,
              path=path,
              overwrite=overwrite,
          )
  
      def visualize(self, title: Optional[str] = None) -> None:
          """Visualize the artifact in notebook environments.
  
          Args:
              title: Optional title to show before the visualizations.
          """
          from zenml.utils.visualization_utils import visualize_artifact
  
          visualize_artifact(self, title=title)
  
  
  # ------------------ Filter Model ------------------
  
  
  class ArtifactVersionFilter(
      ProjectScopedFilter, TaggableFilter, RunMetadataFilterMixin
  ):
      """Model to enable advanced filtering of artifact versions."""
  
      FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
          *TaggableFilter.FILTER_EXCLUDE_FIELDS,
          *RunMetadataFilterMixin.FILTER_EXCLUDE_FIELDS,
          "artifact_id",
          "artifact",
          "only_unused",
          "has_custom_name",
          "model",
          "pipeline_run",
          "model_version_id",
      ]
      CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = [
          *ProjectScopedFilter.CUSTOM_SORTING_OPTIONS,
          *TaggableFilter.CUSTOM_SORTING_OPTIONS,
          *RunMetadataFilterMixin.CUSTOM_SORTING_OPTIONS,
      ]
      CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
          *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
          *TaggableFilter.CLI_EXCLUDE_FIELDS,
          *RunMetadataFilterMixin.CLI_EXCLUDE_FIELDS,
          "artifact_id",
      ]
      API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = [
          *ProjectScopedFilter.API_MULTI_INPUT_PARAMS,
          *TaggableFilter.API_MULTI_INPUT_PARAMS,
          *RunMetadataFilterMixin.API_MULTI_INPUT_PARAMS,
      ]
  
      artifact: Optional[Union[UUID, str]] = Field(
          default=None,
          description="The name or ID of the artifact which the search is scoped "
          "to. This field must always be set and is always applied in addition "
          "to the other filters, regardless of the value of the "
          "logical_operator field.",
          union_mode="left_to_right",
      )
      artifact_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="[Deprecated] Use 'artifact' instead. ID of the artifact to which this version belongs.",
          union_mode="left_to_right",
      )
      version: Optional[str] = Field(
          default=None,
          description="Version of the artifact",
      )
      version_number: Optional[Union[int, str]] = Field(
          default=None,
          description="Version of the artifact if it is an integer",
          union_mode="left_to_right",
      )
      uri: Optional[str] = Field(
          default=None,
          description="Uri of the artifact",
      )
      materializer: Optional[str] = Field(
          default=None,
          description="Materializer used to produce the artifact",
      )
      type: Optional[str] = Field(
          default=None,
          description="Type of the artifact",
      )
      data_type: Optional[str] = Field(
          default=None,
          description="Datatype of the artifact",
      )
      artifact_store_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Artifact store for this artifact",
          union_mode="left_to_right",
      )
      model_version_id: Optional[Union[UUID, str]] = Field(
          default=None,
          description="ID of the model version that is associated with this "
          "artifact version.",
          union_mode="left_to_right",
      )
      only_unused: Optional[bool] = Field(
          default=False, description="Filter only for unused artifacts"
      )
      has_custom_name: Optional[bool] = Field(
          default=None,
          description="Filter only artifacts with/without custom names.",
      )
      model: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Name/ID of the model that is associated with this "
          "artifact version.",
      )
      pipeline_run: Optional[Union[UUID, str]] = Field(
          default=None,
          description="Name/ID of a pipeline run that is associated with this "
          "artifact version.",
      )
  
      model_config = ConfigDict(protected_namespaces=())
  
      def get_custom_filters(
          self, table: Type["AnySchema"]
      ) -> List[Union["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_, or_, select
  
          from zenml.zen_stores.schemas import (
              ArtifactSchema,
              ArtifactVersionSchema,
              ModelSchema,
              ModelVersionArtifactSchema,
              ModelVersionSchema,
              PipelineRunSchema,
              StepRunInputArtifactSchema,
              StepRunOutputArtifactSchema,
              StepRunSchema,
          )
  
          if self.artifact:
              value, operator = self._resolve_operator(self.artifact)
              artifact_filter = and_(
                  ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                  self.generate_name_or_id_query_conditions(
                      value=self.artifact, table=ArtifactSchema
                  ),
              )
              custom_filters.append(artifact_filter)
  
          if self.only_unused:
              unused_filter = and_(
                  ArtifactVersionSchema.id.notin_(  # type: ignore[attr-defined]
                      select(StepRunOutputArtifactSchema.artifact_id)
                  ),
                  ArtifactVersionSchema.id.notin_(  # type: ignore[attr-defined]
                      select(StepRunInputArtifactSchema.artifact_id)
                  ),
              )
              custom_filters.append(unused_filter)
  
          if self.model_version_id:
              value, operator = self._resolve_operator(self.model_version_id)
  
              model_version_filter = and_(
                  ArtifactVersionSchema.id
                  == ModelVersionArtifactSchema.artifact_version_id,
                  ModelVersionArtifactSchema.model_version_id
                  == ModelVersionSchema.id,
                  FilterGenerator(ModelVersionSchema)
                  .define_filter(column="id", value=value, operator=operator)
                  .generate_query_conditions(ModelVersionSchema),
              )
              custom_filters.append(model_version_filter)
  
          if self.has_custom_name is not None:
              custom_name_filter = and_(
                  ArtifactVersionSchema.artifact_id == ArtifactSchema.id,
                  ArtifactSchema.has_custom_name == self.has_custom_name,
              )
              custom_filters.append(custom_name_filter)
  
          if self.model:
              model_filter = and_(
                  ArtifactVersionSchema.id
                  == ModelVersionArtifactSchema.artifact_version_id,
                  ModelVersionArtifactSchema.model_version_id
                  == ModelVersionSchema.id,
                  ModelVersionSchema.model_id == ModelSchema.id,
                  self.generate_name_or_id_query_conditions(
                      value=self.model, table=ModelSchema
                  ),
              )
              custom_filters.append(model_filter)
  
          if self.pipeline_run:
              pipeline_run_filter = and_(
                  or_(
                      and_(
                          ArtifactVersionSchema.id
                          == StepRunOutputArtifactSchema.artifact_id,
                          StepRunOutputArtifactSchema.step_id
                          == StepRunSchema.id,
                      ),
                      and_(
                          ArtifactVersionSchema.id
                          == StepRunInputArtifactSchema.artifact_id,
                          StepRunInputArtifactSchema.step_id == StepRunSchema.id,
                      ),
                  ),
                  StepRunSchema.pipeline_run_id == PipelineRunSchema.id,
                  self.generate_name_or_id_query_conditions(
                      value=self.pipeline_run, table=PipelineRunSchema
                  ),
              )
              custom_filters.append(pipeline_run_filter)
  
          return custom_filters
  
      @model_validator(mode="after")
      def _migrate_artifact_id(self) -> "ArtifactVersionFilter":
          """Migrate value from the deprecated artifact_id attribute.
  
          Returns:
              The filter with migrated value.
          """
          # Handle deprecated artifact_id field
          if self.artifact_id is not None:
              logger.warning(
                  "The 'ArtifactVersionFilter.artifact_id' field is deprecated "
                  "and will be removed in a future version. Please use "
                  "'ArtifactVersionFilter.artifact' instead."
              )
              self.artifact = self.artifact or self.artifact_id
  
          return self
  
  
  # -------------------- Lazy Loader --------------------
  
  
  class LazyArtifactVersionResponse(ArtifactVersionResponse):
      """Lazy artifact version response.
  
      Used if the artifact version is accessed from the model in
      a pipeline context available only during pipeline compilation.
      """
  
      id: Optional[UUID] = None  # type: ignore[assignment]
      lazy_load_name: Optional[str] = None
      lazy_load_version: Optional[str] = None
      lazy_load_model_name: str
      lazy_load_model_version: Optional[str] = None
  
      def get_body(self) -> None:  # type: ignore[override]
          """Protects from misuse of the lazy loader.
  
          Raises:
              RuntimeError: always
          """
          raise RuntimeError("Cannot access artifact body before pipeline runs.")
  
      def get_metadata(self) -> None:  # type: ignore[override]
          """Protects from misuse of the lazy loader.
  
          Raises:
              RuntimeError: always
          """
          raise RuntimeError(
              "Cannot access artifact metadata before pipeline runs."
          )
  
      @property
      def run_metadata(self) -> Dict[str, MetadataType]:
          """The `metadata` property in lazy loading mode.
  
          Returns:
              getter of lazy responses for internal use.
          """
          from zenml.metadata.lazy_load import RunMetadataLazyGetter
  
          return RunMetadataLazyGetter(  # type: ignore[return-value]
              self.lazy_load_model_name,
              self.lazy_load_model_version,
              self.lazy_load_name,
              self.lazy_load_version,
          )
  