zenml-nightly 0.82.1.dev20250526__py3-none-any.whl → 0.82.1.dev20250528__py3-none-any.whl

zenml/zen_stores/schemas/action_schemas.py

@@ -15,11 +15,13 @@
 
 import base64
 import json
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence
 from uuid import UUID
 
 from pydantic.json import pydantic_encoder
 from sqlalchemy import TEXT, Column, UniqueConstraint
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -35,6 +37,7 @@ 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 import TriggerSchema
@@ -99,6 +102,39 @@ class ActionSchema(NamedSchema, table=True):
 
     configuration: bytes
 
+    @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(ActionSchema.user)),
+                    joinedload(
+                        jl_arg(ActionSchema.service_account), innerjoin=True
+                    ),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(cls, request: "ActionRequest") -> "ActionSchema":
         """Convert a `ActionRequest` to a `ActionSchema`.
@@ -169,7 +205,8 @@ class ActionSchema(NamedSchema, table=True):
             The converted model.
         """
         body = ActionResponseBody(
-            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,
             flavor=self.flavor,
@@ -178,7 +215,6 @@ class ActionSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = ActionResponseMetadata(
-                project=self.project.to_model(),
                 configuration=json.loads(
                     base64.b64decode(self.configuration).decode()
                 ),
@@ -188,11 +224,11 @@ class ActionSchema(NamedSchema, table=True):
         resources = None
         if include_resources:
             resources = ActionResponseResources(
+                user=self.user.to_model() if self.user else None,
                 service_account=self.service_account.to_model(),
             )
         return ActionResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,

zenml/zen_stores/schemas/api_key_schemas.py

@@ -15,11 +15,13 @@
 
 from datetime import datetime
 from secrets import token_hex
-from typing import Any, Optional, Tuple
+from typing import Any, Optional, Sequence, Tuple
 from uuid import UUID
 
 from passlib.context import CryptContext
 from sqlalchemy import TEXT, Column, UniqueConstraint
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship
 
 from zenml.models import (
@@ -36,6 +38,7 @@ from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
 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
 
 
 class APIKeySchema(NamedSchema, table=True):
@@ -90,6 +93,31 @@ class APIKeySchema(NamedSchema, table=True):
         context = CryptContext(schemes=["bcrypt"], deprecated="auto")
         return context.hash(key)
 
+    @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(APIKeySchema.service_account), innerjoin=True),
+        ]
+
+        return options
+
     @classmethod
     def from_request(
         cls,

zenml/zen_stores/schemas/artifact_schemas.py

@@ -13,13 +13,14 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of artifact table."""
 
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple
 from uuid import UUID
 
 from pydantic import ValidationError
 from sqlalchemy import TEXT, Column, UniqueConstraint
-from sqlalchemy.orm import object_session
-from sqlmodel import Field, Relationship, desc, select
+from sqlalchemy.orm import joinedload, object_session
+from sqlalchemy.sql.base import ExecutableOption
+from sqlmodel import Field, Relationship, asc, col, desc, select
 
 from zenml.config.source import Source
 from zenml.enums import (
@@ -33,11 +34,13 @@ from zenml.models import (
     ArtifactResponse,
     ArtifactResponseBody,
     ArtifactResponseMetadata,
+    ArtifactResponseResources,
     ArtifactUpdate,
     ArtifactVersionRequest,
     ArtifactVersionResponse,
     ArtifactVersionResponseBody,
     ArtifactVersionResponseMetadata,
+    ArtifactVersionResponseResources,
     ArtifactVersionUpdate,
 )
 from zenml.models.v2.core.artifact import ArtifactRequest
@@ -46,12 +49,11 @@ from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 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.step_run_schemas import (
-    StepRunInputArtifactSchema,
-    StepRunOutputArtifactSchema,
-)
 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_visualization_schemas import (
@@ -61,6 +63,10 @@ if TYPE_CHECKING:
         ModelVersionArtifactSchema,
     )
     from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+    from zenml.zen_stores.schemas.step_run_schemas import (
+        StepRunInputArtifactSchema,
+        StepRunOutputArtifactSchema,
+    )
     from zenml.zen_stores.schemas.tag_schemas import TagSchema
 
 
@@ -114,6 +120,37 @@ class ArtifactSchema(NamedSchema, table=True):
         back_populates="artifacts",
     )
 
+    @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(ArtifactSchema.user)),
+                    # joinedload(jl_arg(ArtifactSchema.tags)),
+                ]
+            )
+
+        return options
+
     @property
     def latest_version(self) -> Optional["ArtifactVersionSchema"]:
         """Fetch the latest version for this artifact.
@@ -178,19 +215,12 @@ class ArtifactSchema(NamedSchema, table=True):
         Returns:
             The created `ArtifactResponse`.
         """
-        latest_id, latest_name = None, None
-        if latest_version := self.latest_version:
-            latest_id = latest_version.id
-            latest_name = latest_version.version
-
         # Create the body of the model
         body = ArtifactResponseBody(
+            user_id=self.user_id,
+            project_id=self.project_id,
             created=self.created,
             updated=self.updated,
-            tags=[tag.to_model() for tag in self.tags],
-            latest_version_name=latest_name,
-            latest_version_id=latest_id,
-            user=self.user.to_model() if self.user else None,
         )
 
         # Create the metadata of the model
@@ -198,15 +228,28 @@ class ArtifactSchema(NamedSchema, table=True):
         if include_metadata:
             metadata = ArtifactResponseMetadata(
                 has_custom_name=self.has_custom_name,
-                project=self.project.to_model(),
+            )
+
+        resources = None
+        if include_resources:
+            latest_id, latest_name = None, None
+            if latest_version := self.latest_version:
+                latest_id = latest_version.id
+                latest_name = latest_version.version
+
+            resources = ArtifactResponseResources(
+                user=self.user.to_model() if self.user else None,
+                tags=[tag.to_model() for tag in self.tags],
+                latest_version_id=latest_id,
+                latest_version_name=latest_name,
             )
 
         return ArtifactResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )
 
     def update(self, artifact_update: ArtifactUpdate) -> "ArtifactSchema":
@@ -305,24 +348,105 @@ class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
             overlaps="run_metadata",
         ),
     )
-    output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
-        back_populates="artifact_version",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
-    input_of_step_runs: List["StepRunInputArtifactSchema"] = Relationship(
-        back_populates="artifact_version",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
     visualizations: List["ArtifactVisualizationSchema"] = Relationship(
         back_populates="artifact_version",
         sa_relationship_kwargs={"cascade": "delete"},
     )
+
+    # Needed for cascade deletion behavior
     model_versions_artifacts_links: List["ModelVersionArtifactSchema"] = (
         Relationship(
             back_populates="artifact_version",
             sa_relationship_kwargs={"cascade": "delete"},
         )
     )
+    output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
+        back_populates="artifact_version",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
+    input_of_step_runs: List["StepRunInputArtifactSchema"] = Relationship(
+        back_populates="artifact_version",
+        sa_relationship_kwargs={"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_metadata:
+        #     options.extend(
+        #         [
+        #             joinedload(jl_arg(ArtifactVersionSchema.visualizations)),
+        #             joinedload(jl_arg(ArtifactVersionSchema.run_metadata)),
+        #         ]
+        #     )
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(ArtifactVersionSchema.user)),
+                    # joinedload(jl_arg(ArtifactVersionSchema.tags)),
+                ]
+            )
+
+        return options
+
+    @property
+    def producer_run_ids(self) -> Optional[Tuple[UUID, UUID]]:
+        """Fetch the producer run IDs for this artifact version.
+
+        Raises:
+            RuntimeError: If no session for the schema exists.
+
+        Returns:
+            The producer step run ID and pipeline run ID for this artifact
+            version.
+        """
+        from zenml.zen_stores.schemas import (
+            StepRunOutputArtifactSchema,
+            StepRunSchema,
+        )
+
+        if session := object_session(self):
+            row = session.execute(
+                select(StepRunSchema.id, StepRunSchema.pipeline_run_id)
+                .join(
+                    StepRunOutputArtifactSchema,
+                    col(StepRunOutputArtifactSchema.step_id)
+                    == col(StepRunSchema.id),
+                )
+                .where(col(StepRunOutputArtifactSchema.artifact_id) == self.id)
+                .where(
+                    col(StepRunSchema.status)
+                    == ExecutionStatus.COMPLETED.value
+                )
+                # Fetch the oldest step run
+                .order_by(asc(StepRunSchema.created))
+                .limit(1)
+            ).one_or_none()
+
+            return (row[0], row[1]) if row else None
+        else:
+            raise RuntimeError(
+                "Missing DB session to fetch producer run for artifact version."
+            )
 
     @classmethod
     def from_request(
@@ -391,35 +515,19 @@ class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
             # This is an old source which was an importable source path
             data_type = Source.from_import_path(self.data_type)
 
-        producer_step_run_id, producer_pipeline_run_id = None, None
-        if self.output_of_step_runs:
-            original_step_runs = [
-                sr
-                for sr in self.output_of_step_runs
-                if sr.step_run.status == ExecutionStatus.COMPLETED
-            ]
-            if len(original_step_runs) == 1:
-                step_run = original_step_runs[0].step_run
-                producer_step_run_id = step_run.id
-                producer_pipeline_run_id = step_run.pipeline_run_id
-            else:
-                step_run = self.output_of_step_runs[0].step_run
-                producer_step_run_id = step_run.original_step_run_id
-
         # Create the body of the model
         artifact = self.artifact.to_model()
         body = ArtifactVersionResponseBody(
+            user_id=self.user_id,
+            project_id=self.project_id,
             artifact=artifact,
             version=self.version or str(self.version_number),
-            user=self.user.to_model() if self.user else None,
             uri=self.uri,
             type=ArtifactType(self.type),
             materializer=materializer,
             data_type=data_type,
             created=self.created,
             updated=self.updated,
-            tags=[tag.to_model() for tag in self.tags],
-            producer_pipeline_run_id=producer_pipeline_run_id,
             save_type=ArtifactSaveType(self.save_type),
             artifact_store_id=self.artifact_store_id,
         )
@@ -428,17 +536,29 @@ class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
         metadata = None
         if include_metadata:
             metadata = ArtifactVersionResponseMetadata(
-                project=self.project.to_model(),
-                producer_step_run_id=producer_step_run_id,
                 visualizations=[v.to_model() for v in self.visualizations],
                 run_metadata=self.fetch_metadata(),
             )
 
         resources = None
+        if include_resources:
+            producer_step_run_id, producer_pipeline_run_id = None, None
+            if producer_run_ids := self.producer_run_ids:
+                # TODO: Why was the producer_pipeline_run_id only set for one
+                # of the cases before?
+                producer_step_run_id, producer_pipeline_run_id = (
+                    producer_run_ids
+                )
+
+            resources = ArtifactVersionResponseResources(
+                user=self.user.to_model() if self.user else None,
+                tags=[tag.to_model() for tag in self.tags],
+                producer_step_run_id=producer_step_run_id,
+                producer_pipeline_run_id=producer_pipeline_run_id,
+            )
 
         return ArtifactVersionResponse(
             id=self.id,
-            project_id=self.project_id,
             body=body,
             metadata=metadata,
             resources=resources,

zenml/zen_stores/schemas/base_schemas.py

@@ -14,9 +14,10 @@
 """Base classes for SQLModel schemas."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, Sequence, TypeVar
 from uuid import UUID, uuid4
 
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, SQLModel
 
 from zenml.utils.time_utils import utc_now
@@ -34,6 +35,30 @@ class BaseSchema(SQLModel):
     created: datetime = Field(default_factory=utc_now)
     updated: datetime = Field(default_factory=utc_now)
 
+    @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.
+
+        This method should return query options that improve the performance
+        when trying to later on converting that schema to a model.
+
+        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.
+        """
+        return []
+
     def to_model(
         self,
         include_metadata: bool = False,

zenml/zen_stores/schemas/code_repository_schemas.py

@@ -14,10 +14,12 @@
 """SQL Model Implementations for code repositories."""
 
 import json
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 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.models import (
@@ -29,6 +31,7 @@ from zenml.models import (
     CodeRepositoryResponse,
     CodeRepositoryResponseBody,
     CodeRepositoryResponseMetadata,
+    CodeRepositoryResponseResources,
     CodeRepositoryUpdate,
 )
 from zenml.utils.time_utils import utc_now
@@ -36,6 +39,7 @@ from zenml.zen_stores.schemas.base_schemas import BaseSchema, 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
 
 
 class CodeRepositorySchema(NamedSchema, table=True):
@@ -78,6 +82,36 @@ class CodeRepositorySchema(NamedSchema, table=True):
     logo_url: Optional[str] = Field()
     description: Optional[str] = Field(sa_column=Column(TEXT, 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_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(CodeRepositorySchema.user)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(
         cls, request: "CodeRepositoryRequest"
@@ -118,7 +152,8 @@ class CodeRepositorySchema(NamedSchema, table=True):
             The created CodeRepositoryResponse.
         """
         body = CodeRepositoryResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
+            project_id=self.project_id,
             source=json.loads(self.source),
             logo_url=self.logo_url,
             created=self.created,
@@ -127,16 +162,22 @@ class CodeRepositorySchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = CodeRepositoryResponseMetadata(
-                project=self.project.to_model(),
                 config=json.loads(self.config),
                 description=self.description,
             )
+
+        resources = None
+        if include_resources:
+            resources = CodeRepositoryResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
+
         return CodeRepositoryResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
-            metadata=metadata,
             body=body,
+            metadata=metadata,
+            resources=resources,
         )
 
     def update(self, update: "CodeRepositoryUpdate") -> "CodeRepositorySchema":