zenml-nightly 0.82.1.dev20250527__py3-none-any.whl → 0.83.0.dev20250529__py3-none-any.whl

zenml/zen_stores/schemas/step_run_schemas.py

@@ -15,12 +15,14 @@
 
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
 from uuid import UUID
 
 from pydantic import ConfigDict
 from sqlalchemy import TEXT, Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship, SQLModel
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
@@ -53,7 +55,10 @@ 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 RunMetadataInterface
+from zenml.zen_stores.schemas.utils import (
+    RunMetadataInterface,
+    jl_arg,
+)
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
@@ -185,6 +190,55 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
 
     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.
+        """
+        from zenml.zen_stores.schemas import ModelVersionSchema
+
+        options = [
+            joinedload(jl_arg(StepRunSchema.deployment)),
+            joinedload(jl_arg(StepRunSchema.pipeline_run)),
+        ]
+
+        if include_metadata:
+            options.extend(
+                [
+                    joinedload(jl_arg(StepRunSchema.logs)),
+                    # joinedload(jl_arg(StepRunSchema.parents)),
+                    # joinedload(jl_arg(StepRunSchema.run_metadata)),
+                ]
+            )
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(StepRunSchema.model_version)).joinedload(
+                        jl_arg(ModelVersionSchema.model), innerjoin=True
+                    ),
+                    joinedload(jl_arg(StepRunSchema.user)),
+                    # joinedload(jl_arg(StepRunSchema.input_artifacts)),
+                    # joinedload(jl_arg(StepRunSchema.output_artifacts)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(
         cls, request: StepRunRequest, deployment_id: Optional[UUID]
@@ -234,22 +288,6 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
         Raises:
             ValueError: In case the step run configuration is missing.
         """
-        input_artifacts = {
-            artifact.name: StepRunInputResponse(
-                input_type=StepRunInputArtifactType(artifact.type),
-                **artifact.artifact_version.to_model().model_dump(),
-            )
-            for artifact in self.input_artifacts
-        }
-
-        output_artifacts: Dict[str, List["ArtifactVersionResponse"]] = {}
-        for artifact in self.output_artifacts:
-            if artifact.name not in output_artifacts:
-                output_artifacts[artifact.name] = []
-            output_artifacts[artifact.name].append(
-                artifact.artifact_version.to_model()
-            )
-
         step = None
         if self.deployment is not None:
             step_configurations = json.loads(
@@ -284,20 +322,19 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
             )
 
         body = StepRunResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
+            project_id=self.project_id,
             status=ExecutionStatus(self.status),
             start_time=self.start_time,
             end_time=self.end_time,
-            inputs=input_artifacts,
-            outputs=output_artifacts,
             created=self.created,
             updated=self.updated,
             model_version_id=self.model_version_id,
+            substitutions=step.config.substitutions,
         )
         metadata = None
         if include_metadata:
             metadata = StepRunResponseMetadata(
-                project=self.project.to_model(),
                 config=step.config,
                 spec=step.spec,
                 cache_key=self.cache_key,
@@ -318,11 +355,33 @@ class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
             if self.model_version:
                 model_version = self.model_version.to_model()
 
-            resources = StepRunResponseResources(model_version=model_version)
+            input_artifacts: Dict[str, List[StepRunInputResponse]] = {}
+            for input_artifact in self.input_artifacts:
+                if input_artifact.name not in input_artifacts:
+                    input_artifacts[input_artifact.name] = []
+                step_run_input = StepRunInputResponse(
+                    input_type=StepRunInputArtifactType(input_artifact.type),
+                    **input_artifact.artifact_version.to_model().model_dump(),
+                )
+                input_artifacts[input_artifact.name].append(step_run_input)
+
+            output_artifacts: Dict[str, List["ArtifactVersionResponse"]] = {}
+            for output_artifact in self.output_artifacts:
+                if output_artifact.name not in output_artifacts:
+                    output_artifacts[output_artifact.name] = []
+                output_artifacts[output_artifact.name].append(
+                    output_artifact.artifact_version.to_model()
+                )
+
+            resources = StepRunResponseResources(
+                user=self.user.to_model() if self.user else None,
+                model_version=model_version,
+                inputs=input_artifacts,
+                outputs=output_artifacts,
+            )
 
         return StepRunResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,
@@ -410,7 +469,10 @@ class StepRunInputArtifactSchema(SQLModel, table=True):
 
     # Relationships
     step_run: "StepRunSchema" = Relationship(back_populates="input_artifacts")
-    artifact_version: "ArtifactVersionSchema" = Relationship()
+    artifact_version: "ArtifactVersionSchema" = Relationship(
+        back_populates="input_of_step_runs",
+        sa_relationship_kwargs={"lazy": "joined"},
+    )
 
 
 class StepRunOutputArtifactSchema(SQLModel, table=True):
@@ -446,5 +508,6 @@ class StepRunOutputArtifactSchema(SQLModel, table=True):
     # Relationship
     step_run: "StepRunSchema" = Relationship(back_populates="output_artifacts")
     artifact_version: "ArtifactVersionSchema" = Relationship(
-        back_populates="output_of_step_runs"
+        back_populates="output_of_step_runs",
+        sa_relationship_kwargs={"lazy": "joined"},
     )

zenml/zen_stores/schemas/tag_schemas.py

@@ -13,11 +13,13 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of tag tables."""
 
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Sequence
 from uuid import UUID
 
 from sqlalchemy import VARCHAR, Column, UniqueConstraint
-from sqlmodel import Field, Relationship
+from sqlalchemy.orm import joinedload, noload, object_session
+from sqlalchemy.sql.base import ExecutableOption
+from sqlmodel import Field, Relationship, col, func, select
 
 from zenml.enums import ColorVariants, TaggableResourceTypes
 from zenml.models import (
@@ -28,6 +30,7 @@ from zenml.models import (
     TagResponse,
     TagResponseBody,
     TagResponseMetadata,
+    TagResponseResources,
     TagUpdate,
 )
 from zenml.utils.time_utils import utc_now
@@ -37,6 +40,7 @@ from zenml.zen_stores.schemas.schema_utils import (
     build_index,
 )
 from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import jl_arg
 
 
 class TagSchema(NamedSchema, table=True):
@@ -68,6 +72,57 @@ class TagSchema(NamedSchema, table=True):
         sa_relationship_kwargs={"overlaps": "tags", "cascade": "delete"},
     )
 
+    @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(TagSchema.user))])
+
+        return options
+
+    @property
+    def tagged_count(self) -> int:
+        """Fetch the number of resources tagged with this tag.
+
+        Raises:
+            RuntimeError: If no session for the schema exists.
+
+        Returns:
+            The number of resources tagged with this tag.
+        """
+        from zenml.zen_stores.schemas import TagResourceSchema
+
+        if session := object_session(self):
+            count = session.scalar(
+                select(func.count(col(TagResourceSchema.id)))
+                .where(TagResourceSchema.tag_id == self.id)
+                .options(noload("*"))
+            )
+
+            return int(count) if count else 0
+        else:
+            raise RuntimeError(
+                "Missing DB session to fetch tagged count for tag."
+            )
+
     @classmethod
     def from_request(cls, request: TagRequest) -> "TagSchema":
         """Convert an `TagRequest` to an `TagSchema`.
@@ -104,19 +159,28 @@ class TagSchema(NamedSchema, table=True):
         """
         metadata = None
         if include_metadata:
-            metadata = TagResponseMetadata()
+            metadata = TagResponseMetadata(
+                tagged_count=self.tagged_count,
+            )
+
+        resources = None
+        if include_resources:
+            resources = TagResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
+
         return TagResponse(
             id=self.id,
             name=self.name,
             body=TagResponseBody(
-                user=self.user.to_model() if self.user else None,
+                user_id=self.user_id,
                 created=self.created,
                 updated=self.updated,
                 color=ColorVariants(self.color),
                 exclusive=self.exclusive,
-                tagged_count=len(self.links),
             ),
             metadata=metadata,
+            resources=resources,
         )
 
     def update(self, update: TagUpdate) -> "TagSchema":

zenml/zen_stores/schemas/trigger_schemas.py

@@ -15,10 +15,12 @@
 
 import base64
 import json
-from typing import Any, List, Optional, cast
+from typing import Any, List, Optional, Sequence, cast
 from uuid import UUID
 
 from sqlalchemy import TEXT, Column, UniqueConstraint
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship
 
 from zenml.config.schedule import Schedule
@@ -44,7 +46,10 @@ from zenml.zen_stores.schemas.event_source_schemas import EventSourceSchema
 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 get_page_from_list
+from zenml.zen_stores.schemas.utils import (
+    get_page_from_list,
+    jl_arg,
+)
 
 
 class TriggerSchema(NamedSchema, table=True):
@@ -118,6 +123,40 @@ class TriggerSchema(NamedSchema, table=True):
     description: str = Field(sa_column=Column(TEXT, nullable=True))
     is_active: bool = Field(nullable=False)
 
+    @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 = [
+            joinedload(jl_arg(TriggerSchema.action), innerjoin=True),
+            joinedload(jl_arg(TriggerSchema.event_source), innerjoin=True),
+        ]
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(TriggerSchema.user)),
+                    # joinedload(jl_arg(TriggerSchema.executions)),
+                ]
+            )
+
+        return options
+
     def update(self, trigger_update: "TriggerUpdate") -> "TriggerSchema":
         """Updates a trigger schema with a trigger update model.
 
@@ -192,7 +231,8 @@ class TriggerSchema(NamedSchema, table=True):
         from zenml.models import TriggerExecutionResponse
 
         body = TriggerResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
+            project_id=self.project_id,
             created=self.created,
             updated=self.updated,
             action_flavor=self.action.flavor,
@@ -208,7 +248,6 @@ class TriggerSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = TriggerResponseMetadata(
-                project=self.project.to_model(),
                 event_filter=json.loads(
                     base64.b64decode(self.event_filter).decode()
                 ),
@@ -231,6 +270,7 @@ class TriggerSchema(NamedSchema, table=True):
                 ),
             )
             resources = TriggerResponseResources(
+                user=self.user.to_model() if self.user else None,
                 action=self.action.to_model(),
                 event_source=self.event_source.to_model()
                 if self.event_source
@@ -239,7 +279,6 @@ class TriggerSchema(NamedSchema, table=True):
             )
         return TriggerResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,

zenml/zen_stores/schemas/utils.py

@@ -15,8 +15,9 @@
 
 import json
 import math
-from typing import Dict, List, Type, TypeVar
+from typing import Any, Dict, List, Type, TypeVar, cast
 
+from sqlalchemy.orm import InstrumentedAttribute
 from sqlmodel import Relationship
 
 from zenml.metadata.metadata_types import MetadataType
@@ -26,6 +27,18 @@ from zenml.zen_stores.schemas.base_schemas import BaseSchema
 S = TypeVar("S", bound=BaseSchema)
 
 
+def jl_arg(column: Any) -> InstrumentedAttribute[Any]:
+    """Cast a SQLModel column to a joinedload argument.
+
+    Args:
+        column: The column.
+
+    Returns:
+        The column cast to a joinedload argument.
+    """
+    return cast(InstrumentedAttribute[Any], column)
+
+
 def get_page_from_list(
     items_list: List[S],
     response_model: Type[BaseResponse],  # type: ignore[type-arg]
@@ -77,9 +90,14 @@ class RunMetadataInterface:
 
     run_metadata = Relationship()
 
-    def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
+    def fetch_metadata_collection(
+        self, **kwargs: Any
+    ) -> Dict[str, List[RunMetadataEntry]]:
         """Fetches all the metadata entries related to the entity.
 
+        Args:
+            **kwargs: Keyword arguments.
+
         Returns:
             A dictionary, where the key is the key of the metadata entry
                 and the values represent the list of entries with this key.
@@ -98,14 +116,17 @@ class RunMetadataInterface:
 
         return metadata_collection
 
-    def fetch_metadata(self) -> Dict[str, MetadataType]:
+    def fetch_metadata(self, **kwargs: Any) -> Dict[str, MetadataType]:
         """Fetches the latest metadata entry related to the entity.
 
+        Args:
+            **kwargs: Keyword arguments to pass to the metadata collection.
+
         Returns:
             A dictionary, where the key is the key of the metadata entry
                 and the values represent the latest entry with this key.
         """
-        metadata_collection = self.fetch_metadata_collection()
+        metadata_collection = self.fetch_metadata_collection(**kwargs)
         return {
             k: sorted(v, key=lambda x: x.created, reverse=True)[0].value
             for k, v in metadata_collection.items()