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

zenml/zen_stores/schemas/run_template_schemas.py

@@ -13,12 +13,13 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of run template tables."""
 
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence
 from uuid import UUID
 
 from sqlalchemy import Column, String, UniqueConstraint
 from sqlalchemy.dialects.mysql import MEDIUMTEXT
-from sqlalchemy.orm import object_session
+from sqlalchemy.orm import joinedload, object_session
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship, col, desc, select
 
 from zenml.constants import MEDIUMTEXT_MAX_LENGTH
@@ -36,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.pipeline_deployment_schemas import (
@@ -116,6 +118,68 @@ class RunTemplateSchema(NamedSchema, table=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.
+        """
+        from zenml.zen_stores.schemas import PipelineDeploymentSchema
+
+        options = [
+            joinedload(jl_arg(RunTemplateSchema.source_deployment)).joinedload(
+                jl_arg(PipelineDeploymentSchema.build)
+            ),
+        ]
+
+        if include_metadata or include_resources:
+            options.extend(
+                [
+                    joinedload(
+                        jl_arg(RunTemplateSchema.source_deployment)
+                    ).joinedload(jl_arg(PipelineDeploymentSchema.pipeline)),
+                    joinedload(
+                        jl_arg(RunTemplateSchema.source_deployment)
+                    ).joinedload(
+                        jl_arg(PipelineDeploymentSchema.code_reference)
+                    ),
+                ]
+            )
+        if include_metadata:
+            options.extend(
+                [
+                    joinedload(
+                        jl_arg(RunTemplateSchema.source_deployment)
+                    ).joinedload(jl_arg(PipelineDeploymentSchema.stack)),
+                    joinedload(
+                        jl_arg(RunTemplateSchema.source_deployment)
+                    ).joinedload(jl_arg(PipelineDeploymentSchema.schedule)),
+                ]
+            )
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(RunTemplateSchema.user)),
+                    # joinedload(jl_arg(RunTemplateSchema.tags)),
+                ]
+            )
+
+        return options
+
     @property
     def latest_run(self) -> Optional["PipelineRunSchema"]:
         """Fetch the latest run for this template.
@@ -216,20 +280,17 @@ class RunTemplateSchema(NamedSchema, table=True):
             self.source_deployment
             and self.source_deployment.build
             and not self.source_deployment.build.is_local
-            and self.source_deployment.build.stack
+            and self.source_deployment.build.stack_id
         ):
             runnable = True
 
-        latest_run = self.latest_run
-
         body = RunTemplateResponseBody(
-            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,
             runnable=runnable,
             hidden=self.hidden,
-            latest_run_id=latest_run.id if latest_run else None,
-            latest_run_status=latest_run.status if latest_run else None,
         )
 
         metadata = None
@@ -247,7 +308,7 @@ class RunTemplateSchema(NamedSchema, table=True):
 
                 if (
                     self.source_deployment.build
-                    and self.source_deployment.build.stack
+                    and self.source_deployment.build.stack_id
                 ):
                     config_template = template_utils.generate_config_template(
                         deployment=self.source_deployment
@@ -257,7 +318,6 @@ class RunTemplateSchema(NamedSchema, table=True):
                     )
 
             metadata = RunTemplateResponseMetadata(
-                project=self.project.to_model(),
                 description=self.description,
                 pipeline_spec=pipeline_spec,
                 config_template=config_template,
@@ -287,7 +347,10 @@ class RunTemplateSchema(NamedSchema, table=True):
                 build = None
                 code_reference = None
 
+            latest_run = self.latest_run
+
             resources = RunTemplateResponseResources(
+                user=self.user.to_model() if self.user else None,
                 source_deployment=self.source_deployment.to_model()
                 if self.source_deployment
                 else None,
@@ -295,11 +358,12 @@ class RunTemplateSchema(NamedSchema, table=True):
                 build=build,
                 code_reference=code_reference,
                 tags=[tag.to_model() for tag in self.tags],
+                latest_run_id=latest_run.id if latest_run else None,
+                latest_run_status=latest_run.status if latest_run else None,
             )
 
         return RunTemplateResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,

zenml/zen_stores/schemas/schedule_schema.py

@@ -14,10 +14,12 @@
 """SQL Model Implementations for Pipeline Schedules."""
 
 from datetime import datetime, timedelta
-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 Field, Relationship
 
 from zenml.enums import MetadataResourceTypes
@@ -26,6 +28,7 @@ from zenml.models import (
     ScheduleResponse,
     ScheduleResponseBody,
     ScheduleResponseMetadata,
+    ScheduleResponseResources,
     ScheduleUpdate,
 )
 from zenml.utils.time_utils import utc_now
@@ -35,7 +38,10 @@ from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
 from zenml.zen_stores.schemas.project_schemas import ProjectSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.user_schemas import UserSchema
-from zenml.zen_stores.schemas.utils import RunMetadataInterface
+from zenml.zen_stores.schemas.utils import (
+    RunMetadataInterface,
+    jl_arg,
+)
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
@@ -120,6 +126,39 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
     catchup: bool
     run_once_start_time: Optional[datetime] = Field(nullable=True)
 
+    @classmethod
+    def get_query_options(
+        cls,
+        include_metadata: bool = False,
+        include_resources: bool = False,
+        **kwargs: Any,
+    ) -> Sequence[ExecutableOption]:
+        """Get the query options for the schema.
+
+        Args:
+            include_metadata: Whether metadata will be included when converting
+                the schema to a model.
+            include_resources: Whether resources will be included when
+                converting the schema to a model.
+            **kwargs: Keyword arguments to allow schema specific logic
+
+        Returns:
+            A list of query options.
+        """
+        options = []
+
+        # if include_metadata:
+        #     options.extend(
+        #         [
+        #             joinedload(jl_arg(ScheduleSchema.run_metadata)),
+        #         ]
+        #     )
+
+        if include_resources:
+            options.extend([joinedload(jl_arg(ScheduleSchema.user))])
+
+        return options
+
     @classmethod
     def from_request(
         cls, schedule_request: ScheduleRequest
@@ -189,7 +228,8 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
             interval_second = None
 
         body = ScheduleResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
+            project_id=self.project_id,
             active=self.active,
             cron_expression=self.cron_expression,
             start_time=self.start_time,
@@ -203,16 +243,21 @@ class ScheduleSchema(NamedSchema, RunMetadataInterface, table=True):
         metadata = None
         if include_metadata:
             metadata = ScheduleResponseMetadata(
-                project=self.project.to_model(),
                 pipeline_id=self.pipeline_id,
                 orchestrator_id=self.orchestrator_id,
                 run_metadata=self.fetch_metadata(),
             )
 
+        resources = None
+        if include_resources:
+            resources = ScheduleResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
+
         return ScheduleResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )

zenml/zen_stores/schemas/secret_schemas.py

@@ -15,10 +15,12 @@
 
 import base64
 import json
-from typing import Any, Dict, Optional, cast
+from typing import Any, Dict, 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 sqlalchemy_utils.types.encrypted.encrypted_type import (
     AesGcmEngine,
     InvalidCiphertextError,
@@ -31,12 +33,14 @@ from zenml.models import (
     SecretResponse,
     SecretResponseBody,
     SecretResponseMetadata,
+    SecretResponseResources,
     SecretUpdate,
 )
 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 SecretDecodeError(Exception):
@@ -75,6 +79,32 @@ class SecretSchema(NamedSchema, table=True):
     )
     user: "UserSchema" = Relationship(back_populates="secrets")
 
+    @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(SecretSchema.user))])
+
+        return options
+
     @classmethod
     def _dump_secret_values(
         cls, values: Dict[str, str], encryption_engine: Optional[AesGcmEngine]
@@ -225,11 +255,17 @@ class SecretSchema(NamedSchema, table=True):
         if include_metadata:
             metadata = SecretResponseMetadata()
 
+        resources = None
+        if include_resources:
+            resources = SecretResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
+
         # Don't load the secret values implicitly in the secret. The
         # SQL secret store will call `get_secret_values` to load the
         # values separately if SQL is used as the secrets store.
         body = SecretResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
             created=self.created,
             updated=self.updated,
             private=self.private,
@@ -239,6 +275,7 @@ class SecretSchema(NamedSchema, table=True):
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )
 
     def get_secret_values(

zenml/zen_stores/schemas/service_connector_schemas.py

@@ -16,10 +16,12 @@
 import base64
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
+from typing import TYPE_CHECKING, Any, Dict, 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.models import (
@@ -27,12 +29,14 @@ from zenml.models import (
     ServiceConnectorResponse,
     ServiceConnectorResponseBody,
     ServiceConnectorResponseMetadata,
+    ServiceConnectorResponseResources,
     ServiceConnectorUpdate,
 )
 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
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
@@ -77,6 +81,32 @@ class ServiceConnectorSchema(NamedSchema, table=True):
         back_populates="connector",
     )
 
+    @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(ServiceConnectorSchema.user))])
+
+        return options
+
     @property
     def resource_types_list(self) -> List[str]:
         """Returns the resource types as a list.
@@ -241,7 +271,7 @@ class ServiceConnectorSchema(NamedSchema, table=True):
             A `ServiceConnectorModel`
         """
         body = ServiceConnectorResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
             created=self.created,
             updated=self.updated,
             description=self.description,
@@ -265,9 +295,16 @@ class ServiceConnectorSchema(NamedSchema, table=True):
                 expiration_seconds=self.expiration_seconds,
                 labels=self.labels_dict,
             )
+        resources = None
+        if include_resources:
+            resources = ServiceConnectorResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
+
         return ServiceConnectorResponse(
             id=self.id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )

zenml/zen_stores/schemas/service_schemas.py

@@ -15,11 +15,13 @@
 
 import base64
 import json
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 from uuid import UUID
 
 from pydantic import ConfigDict
 from sqlalchemy import TEXT, Column
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql.base import ExecutableOption
 from sqlmodel import Field, Relationship
 
 from zenml.models.v2.core.service import (
@@ -38,6 +40,7 @@ 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 jl_arg
 
 
 class ServiceSchema(NamedSchema, table=True):
@@ -117,6 +120,38 @@ class ServiceSchema(NamedSchema, 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 = []
+
+        if include_resources:
+            options.extend(
+                [
+                    joinedload(jl_arg(ServiceSchema.user)),
+                    joinedload(jl_arg(ServiceSchema.model_version)),
+                    joinedload(jl_arg(ServiceSchema.pipeline_run)),
+                ]
+            )
+
+        return options
+
     def to_model(
         self,
         include_metadata: bool = False,
@@ -134,7 +169,8 @@ class ServiceSchema(NamedSchema, table=True):
             The created `ServiceResponse`.
         """
         body = ServiceResponseBody(
-            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,
             service_type=json.loads(self.service_type),
@@ -146,7 +182,6 @@ class ServiceSchema(NamedSchema, table=True):
         metadata = None
         if include_metadata:
             metadata = ServiceResponseMetadata(
-                project=self.project.to_model(),
                 service_source=self.service_source,
                 config=json.loads(base64.b64decode(self.config).decode()),
                 status=json.loads(base64.b64decode(self.status).decode())
@@ -162,6 +197,7 @@ class ServiceSchema(NamedSchema, table=True):
         resources = None
         if include_resources:
             resources = ServiceResponseResources(
+                user=self.user.to_model() if self.user else None,
                 model_version=self.model_version.to_model()
                 if self.model_version
                 else None,
@@ -171,7 +207,6 @@ class ServiceSchema(NamedSchema, table=True):
             )
         return ServiceResponse(
             id=self.id,
-            project_id=self.project_id,
             name=self.name,
             body=body,
             metadata=metadata,

zenml/zen_stores/schemas/stack_schemas.py

@@ -15,22 +15,26 @@
 
 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 Field, Relationship, SQLModel
 
 from zenml.models import (
     StackResponse,
     StackResponseBody,
     StackResponseMetadata,
+    StackResponseResources,
     StackUpdate,
 )
 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
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.component_schemas import (
@@ -106,6 +110,41 @@ class StackSchema(NamedSchema, table=True):
         back_populates="stack",
     )
 
+    @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(StackSchema.components)).joinedload(
+        #                 jl_arg(StackComponentSchema.flavor_schema)
+        #             ),
+        #         ]
+        #     )
+
+        if include_resources:
+            options.extend([joinedload(jl_arg(StackSchema.user))])
+
+        return options
+
     def update(
         self,
         stack_update: "StackUpdate",
@@ -152,7 +191,7 @@ class StackSchema(NamedSchema, table=True):
             The converted model.
         """
         body = StackResponseBody(
-            user=self.user.to_model() if self.user else None,
+            user_id=self.user_id,
             created=self.created,
             updated=self.updated,
         )
@@ -166,10 +205,16 @@ class StackSchema(NamedSchema, table=True):
                 else None,
                 description=self.description,
             )
+        resources = None
+        if include_resources:
+            resources = StackResponseResources(
+                user=self.user.to_model() if self.user else None,
+            )
 
         return StackResponse(
             id=self.id,
             name=self.name,
             body=body,
             metadata=metadata,
+            resources=resources,
         )