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

zenml/zen_stores/schemas/component_schemas.py

@@ -15,10 +15,12 @@
 
 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 sqlalchemy import UniqueConstraint
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Relationship
 
 from zenml.enums import StackComponentType
@@ -38,6 +40,7 @@ from zenml.zen_stores.schemas.service_connector_schemas import (
 )
 from zenml.zen_stores.schemas.stack_schemas import StackCompositionSchema
 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.flavor_schemas import FlavorSchema
@@ -111,6 +114,43 @@ class StackComponentSchema(NamedSchema, table=True):
 
     connector_resource_id: Optional[str]
 
+    @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(StackComponentSchema.flavor_schema)),
+        ]
+
+        if include_metadata:
+            options.extend(
+                [joinedload(jl_arg(StackComponentSchema.connector))]
+            )
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(StackComponentSchema.user)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(
         cls,
@@ -190,9 +230,9 @@ class StackComponentSchema(NamedSchema, table=True):
             A `ComponentModel`
         """
         body = ComponentResponseBody(
+            user_id=self.user_id,
             type=StackComponentType(self.type),
             flavor_name=self.flavor,
-            user=self.user.to_model() if self.user else None,
             created=self.created,
             updated=self.updated,
             logo_url=self.flavor_schema.logo_url
@@ -224,7 +264,8 @@ class StackComponentSchema(NamedSchema, table=True):
                 )
 
             resources = ComponentResponseResources(
-                flavor=self.flavor_schema.to_model()
+                user=self.user.to_model() if self.user else None,
+                flavor=self.flavor_schema.to_model(),
             )
         return ComponentResponse(
             id=self.id,

zenml/zen_stores/schemas/device_schemas.py

@@ -15,10 +15,12 @@
 
 from datetime import datetime, timedelta
 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.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Relationship
 
 from zenml.enums import OAuthDeviceStatus
@@ -29,12 +31,14 @@ from zenml.models import (
     OAuthDeviceResponse,
     OAuthDeviceResponseBody,
     OAuthDeviceResponseMetadata,
+    OAuthDeviceResponseResources,
     OAuthDeviceUpdate,
 )
 from zenml.utils.time_utils import utc_now
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 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 OAuthDeviceSchema(BaseSchema, table=True):
@@ -69,6 +73,36 @@ class OAuthDeviceSchema(BaseSchema, table=True):
     )
     user: Optional["UserSchema"] = Relationship(back_populates="auth_devices")
 
+    @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(OAuthDeviceSchema.user)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def _generate_user_code(cls) -> str:
         """Generate a user code for an OAuth2 device.
@@ -230,7 +264,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             )
 
         body = OAuthDeviceResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
             created=self.created,
             updated=self.updated,
             client_id=self.client_id,
@@ -241,10 +275,16 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             ip_address=self.ip_address,
             hostname=self.hostname,
         )
+        resources = None
+        if include_resources:
+            resources = OAuthDeviceResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
         return OAuthDeviceResponse(
             id=self.id,
             body=body,
             metadata=metadata,
+            resources=resources,
         )
 
     def to_internal_model(
@@ -269,6 +309,7 @@ class OAuthDeviceSchema(BaseSchema, table=True):
             id=device_model.id,
             body=device_model.body,
             metadata=device_model.metadata,
+            resources=device_model.resources,
             user_code=self.user_code,
             device_code=self.device_code,
         )

zenml/zen_stores/schemas/event_source_schemas.py

@@ -15,10 +15,12 @@
 
 import base64
 import json
-from typing import TYPE_CHECKING, Any, List, Optional, cast
+from typing import TYPE_CHECKING, 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 import EventSourceResponseMetadata
@@ -36,7 +38,10 @@ 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 get_page_from_list
+from zenml.zen_stores.schemas.utils import (
+    get_page_from_list,
+    jl_arg,
+)
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import TriggerSchema
@@ -85,6 +90,37 @@ class EventSourceSchema(NamedSchema, table=True):
     configuration: bytes
     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 = []
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(EventSourceSchema.user)),
+                    # joinedload(jl_arg(EventSourceSchema.triggers)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(cls, request: EventSourceRequest) -> "EventSourceSchema":
         """Convert an `EventSourceRequest` to an `EventSourceSchema`.
@@ -133,9 +169,10 @@ class EventSourceSchema(NamedSchema, table=True):
         from zenml.models import TriggerResponse
 
         body = EventSourceResponseBody(
+            user_id=self.user_id,
+            project_id=self.project_id,
             created=self.created,
             updated=self.updated,
-            user=self.user.to_model() if self.user else None,
             flavor=self.flavor,
             plugin_subtype=self.plugin_subtype,
             is_active=self.is_active,
@@ -152,12 +189,12 @@ class EventSourceSchema(NamedSchema, table=True):
                 ),
             )
             resources = EventSourceResponseResources(
+                user=self.user.to_model() if self.user else None,
                 triggers=triggers,
             )
         metadata = None
         if include_metadata:
             metadata = EventSourceResponseMetadata(
-                project=self.project.to_model(),
                 description=self.description,
                 configuration=json.loads(
                     base64.b64decode(self.configuration).decode()
@@ -165,7 +202,6 @@ class EventSourceSchema(NamedSchema, table=True):
             )
         return EventSourceResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,

zenml/zen_stores/schemas/flavor_schemas.py

@@ -14,10 +14,12 @@
 """SQL Model Implementations for Flavors."""
 
 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.enums import StackComponentType
@@ -25,12 +27,14 @@ from zenml.models import (
     FlavorResponse,
     FlavorResponseBody,
     FlavorResponseMetadata,
+    FlavorResponseResources,
     FlavorUpdate,
 )
 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 FlavorSchema(NamedSchema, table=True):
@@ -78,6 +82,36 @@ class FlavorSchema(NamedSchema, table=True):
 
     is_custom: bool = Field(default=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(FlavorSchema.user)),
+                ]
+            )
+
+        return options
+
     def update(
         self,
         flavor_update: "FlavorUpdate",
@@ -121,7 +155,7 @@ class FlavorSchema(NamedSchema, table=True):
             The flavor model.
         """
         body = FlavorResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
             type=StackComponentType(self.type),
             integration=self.integration,
             source=self.source,
@@ -140,9 +174,15 @@ class FlavorSchema(NamedSchema, table=True):
                 docs_url=self.docs_url,
                 sdk_docs_url=self.sdk_docs_url,
             )
+        resources = None
+        if include_resources:
+            resources = FlavorResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
         return FlavorResponse(
             id=self.id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )

zenml/zen_stores/schemas/model_schemas.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, cast
 from uuid import UUID, uuid4
 
 from pydantic import ConfigDict
@@ -24,11 +24,11 @@ from sqlalchemy import (
     Column,
     UniqueConstraint,
 )
-from sqlalchemy.orm import object_session
+from sqlalchemy.orm import joinedload, object_session
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship, desc, select
 
 from zenml.enums import (
-    ArtifactType,
     MetadataResourceTypes,
     TaggableResourceTypes,
 )
@@ -38,6 +38,7 @@ from zenml.models import (
     ModelResponse,
     ModelResponseBody,
     ModelResponseMetadata,
+    ModelResponseResources,
     ModelUpdate,
     ModelVersionArtifactRequest,
     ModelVersionArtifactResponse,
@@ -65,6 +66,7 @@ from zenml.zen_stores.schemas.user_schemas import UserSchema
 from zenml.zen_stores.schemas.utils import (
     RunMetadataInterface,
     get_page_from_list,
+    jl_arg,
 )
 
 if TYPE_CHECKING:
@@ -127,6 +129,37 @@ class ModelSchema(NamedSchema, table=True):
         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_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(ModelSchema.user)),
+                    # joinedload(jl_arg(ModelSchema.tags)),
+                ]
+            )
+
+        return options
+
     @property
     def latest_version(self) -> Optional["ModelVersionSchema"]:
         """Fetch the latest version for this model.
@@ -194,19 +227,9 @@ class ModelSchema(NamedSchema, table=True):
         Returns:
             The created `ModelResponse`.
         """
-        tags = [tag.to_model() for tag in self.tags]
-
-        if latest_version := self.latest_version:
-            latest_version_name = latest_version.name
-            latest_version_id = latest_version.id
-        else:
-            latest_version_name = None
-            latest_version_id = None
-
         metadata = None
         if include_metadata:
             metadata = ModelResponseMetadata(
-                project=self.project.to_model(),
                 license=self.license,
                 description=self.description,
                 audience=self.audience,
@@ -217,21 +240,35 @@ class ModelSchema(NamedSchema, table=True):
                 save_models_to_registry=self.save_models_to_registry,
             )
 
+        resources = None
+        if include_resources:
+            if latest_version := self.latest_version:
+                latest_version_name = latest_version.name
+                latest_version_id = latest_version.id
+            else:
+                latest_version_name = None
+                latest_version_id = None
+
+            resources = ModelResponseResources(
+                user=self.user.to_model() if self.user else None,
+                tags=[tag.to_model() for tag in self.tags],
+                latest_version_name=latest_version_name,
+                latest_version_id=latest_version_id,
+            )
+
         body = ModelResponseBody(
-            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,
-            tags=tags,
-            latest_version_name=latest_version_name,
-            latest_version_id=latest_version_id,
         )
 
         return ModelResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )
 
     def update(
@@ -317,14 +354,6 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
         nullable=False,
     )
     model: "ModelSchema" = Relationship(back_populates="model_versions")
-    artifact_links: List["ModelVersionArtifactSchema"] = Relationship(
-        back_populates="model_version",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
-    pipeline_run_links: List["ModelVersionPipelineRunSchema"] = Relationship(
-        back_populates="model_version",
-        sa_relationship_kwargs={"cascade": "delete"},
-    )
     tags: List["TagSchema"] = Relationship(
         sa_relationship_kwargs=dict(
             primaryjoin=f"and_(foreign(TagResourceSchema.resource_type)=='{TaggableResourceTypes.MODEL_VERSION.value}', foreign(TagResourceSchema.resource_id)==ModelVersionSchema.id)",
@@ -369,6 +398,16 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
     # numeric.
     producer_run_id_if_numeric: UUID
 
+    # Needed for cascade deletion behavior
+    artifact_links: List["ModelVersionArtifactSchema"] = Relationship(
+        back_populates="model_version",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
+    pipeline_run_links: List["ModelVersionPipelineRunSchema"] = Relationship(
+        back_populates="model_version",
+        sa_relationship_kwargs={"cascade": "delete"},
+    )
+
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.
     #  It is possible to suppress this warning message with the following
@@ -377,6 +416,47 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
     #  careful we might overwrite some fields protected by pydantic.
     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.
+        """
+        options = [
+            joinedload(jl_arg(ModelVersionSchema.model), innerjoin=True),
+        ]
+
+        # if include_metadata:
+        #     options.extend(
+        #         [
+        #             joinedload(jl_arg(ModelVersionSchema.run_metadata)),
+        #         ]
+        #     )
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(ModelVersionSchema.user)),
+                    # joinedload(jl_arg(ModelVersionSchema.services)),
+                    # joinedload(jl_arg(ModelVersionSchema.tags)),
+                ]
+            )
+
+        return options
+
     @classmethod
     def from_request(
         cls,
@@ -430,44 +510,9 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
         """
         from zenml.models import ServiceResponse
 
-        # Construct {name: {version: id}} dicts for all linked artifacts
-        model_artifact_ids: Dict[str, Dict[str, UUID]] = {}
-        deployment_artifact_ids: Dict[str, Dict[str, UUID]] = {}
-        data_artifact_ids: Dict[str, Dict[str, UUID]] = {}
-        for artifact_link in self.artifact_links:
-            if not artifact_link.artifact_version:
-                continue
-            artifact_name = artifact_link.artifact_version.artifact.name
-            artifact_version = str(artifact_link.artifact_version.version)
-            artifact_version_id = artifact_link.artifact_version.id
-            if artifact_link.artifact_version.type == ArtifactType.MODEL.value:
-                model_artifact_ids.setdefault(artifact_name, {}).update(
-                    {str(artifact_version): artifact_version_id}
-                )
-            elif (
-                artifact_link.artifact_version.type
-                == ArtifactType.SERVICE.value
-            ):
-                deployment_artifact_ids.setdefault(artifact_name, {}).update(
-                    {str(artifact_version): artifact_version_id}
-                )
-            else:
-                data_artifact_ids.setdefault(artifact_name, {}).update(
-                    {str(artifact_version): artifact_version_id}
-                )
-
-        # Construct {name: id} dict for all linked pipeline runs
-        pipeline_run_ids: Dict[str, UUID] = {}
-        for pipeline_run_link in self.pipeline_run_links:
-            if not pipeline_run_link.pipeline_run:
-                continue
-            pipeline_run = pipeline_run_link.pipeline_run
-            pipeline_run_ids[pipeline_run.name] = pipeline_run.id
-
         metadata = None
         if include_metadata:
             metadata = ModelVersionResponseMetadata(
-                project=self.project.to_model(),
                 description=self.description,
                 run_metadata=self.fetch_metadata(),
             )
@@ -484,26 +529,23 @@ class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
                 ),
             )
             resources = ModelVersionResponseResources(
+                user=self.user.to_model() if self.user else None,
                 services=services,
+                tags=[tag.to_model() for tag in self.tags],
             )
 
         body = ModelVersionResponseBody(
-            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,
             stage=self.stage,
             number=self.number,
             model=self.model.to_model(),
-            model_artifact_ids=model_artifact_ids,
-            data_artifact_ids=data_artifact_ids,
-            deployment_artifact_ids=deployment_artifact_ids,
-            pipeline_run_ids=pipeline_run_ids,
-            tags=[tag.to_model() for tag in self.tags],
         )
 
         return ModelVersionResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,
@@ -549,9 +591,7 @@ class ModelVersionArtifactSchema(BaseSchema, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
-    model_version: "ModelVersionSchema" = Relationship(
-        back_populates="artifact_links"
-    )
+    model_version: "ModelVersionSchema" = Relationship()
     artifact_version_id: UUID = build_foreign_key_field(
         source=__tablename__,
         target=ArtifactVersionSchema.__tablename__,
@@ -632,9 +672,7 @@ class ModelVersionPipelineRunSchema(BaseSchema, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
-    model_version: "ModelVersionSchema" = Relationship(
-        back_populates="pipeline_run_links"
-    )
+    model_version: "ModelVersionSchema" = Relationship()
     pipeline_run_id: UUID = build_foreign_key_field(
         source=__tablename__,
         target=PipelineRunSchema.__tablename__,
@@ -643,9 +681,7 @@ class ModelVersionPipelineRunSchema(BaseSchema, table=True):
         ondelete="CASCADE",
         nullable=False,
     )
-    pipeline_run: "PipelineRunSchema" = Relationship(
-        back_populates="model_versions_pipeline_runs_links"
-    )
+    pipeline_run: "PipelineRunSchema" = Relationship()
 
     # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
     #  fields defined under base models. If not handled, this raises a warning.