zenml-nightly 0.70.0.dev20241129__py3-none-any.whl → 0.70.0.dev20241130__py3-none-any.whl

zenml/zen_stores/migrations/versions/cc269488e5a9_separate_run_metadata.py

@@ -0,0 +1,135 @@
+"""Separate run metadata into resource link table with new UUIDs.
+
+Revision ID: cc269488e5a9
+Revises: b73bc71f1106
+Create Date: 2024-11-12 09:46:46.587478
+"""
+
+import uuid
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "cc269488e5a9"
+down_revision = "b73bc71f1106"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Creates the 'run_metadata_resource' table and migrates data."""
+    op.create_table(
+        "run_metadata_resource",
+        sa.Column(
+            "id",
+            sqlmodel.sql.sqltypes.GUID(),
+            nullable=False,
+            primary_key=True,
+        ),
+        sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+        sa.Column("resource_type", sa.String(length=255), nullable=False),
+        sa.Column(
+            "run_metadata_id",
+            sqlmodel.sql.sqltypes.GUID(),
+            sa.ForeignKey("run_metadata.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+    )
+
+    connection = op.get_bind()
+
+    run_metadata_data = connection.execute(
+        sa.text("""
+            SELECT id, resource_id, resource_type
+            FROM run_metadata
+        """)
+    ).fetchall()
+
+    # Prepare data with new UUIDs for bulk insert
+    resource_data = [
+        {
+            "id": str(uuid.uuid4()),  # Generate a new UUID for each row
+            "resource_id": row.resource_id,
+            "resource_type": row.resource_type,
+            "run_metadata_id": row.id,
+        }
+        for row in run_metadata_data
+    ]
+
+    # Perform bulk insert into `run_metadata_resource`
+    if resource_data:  # Only perform insert if there's data to migrate
+        op.bulk_insert(
+            sa.table(
+                "run_metadata_resource",
+                sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False),
+                sa.Column(
+                    "resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=False
+                ),
+                sa.Column(
+                    "resource_type", sa.String(length=255), nullable=False
+                ),
+                sa.Column(
+                    "run_metadata_id",
+                    sqlmodel.sql.sqltypes.GUID(),
+                    nullable=False,
+                ),  # Changed to BIGINT
+            ),
+            resource_data,
+        )
+
+    op.drop_column("run_metadata", "resource_id")
+    op.drop_column("run_metadata", "resource_type")
+
+    op.add_column(
+        "run_metadata",
+        sa.Column(
+            "publisher_step_id", sqlmodel.sql.sqltypes.GUID(), nullable=True
+        ),
+    )
+
+
+def downgrade() -> None:
+    """Reverts the 'run_metadata_resource' table and migrates data back."""
+    # Recreate the `resource_id` and `resource_type` columns in `run_metadata`
+    op.add_column(
+        "run_metadata",
+        sa.Column("resource_id", sqlmodel.sql.sqltypes.GUID(), nullable=True),
+    )
+    op.add_column(
+        "run_metadata",
+        sa.Column("resource_type", sa.String(length=255), nullable=True),
+    )
+
+    # Migrate data back from `run_metadata_resource` to `run_metadata`
+    connection = op.get_bind()
+
+    # Fetch data from `run_metadata_resource`
+    run_metadata_resource_data = connection.execute(
+        sa.text("""
+            SELECT resource_id, resource_type, run_metadata_id
+            FROM run_metadata_resource
+        """)
+    ).fetchall()
+
+    # Update `run_metadata` with the data from `run_metadata_resource`
+    for row in run_metadata_resource_data:
+        connection.execute(
+            sa.text("""
+                UPDATE run_metadata
+                SET resource_id = :resource_id, resource_type = :resource_type
+                WHERE id = :run_metadata_id
+            """),
+            {
+                "resource_id": row.resource_id,
+                "resource_type": row.resource_type,
+                "run_metadata_id": row.run_metadata_id,
+            },
+        )
+
+    # Drop the `run_metadata_resource` table
+    op.drop_table("run_metadata_resource")
+
+    # Drop the cached column
+    op.drop_column("run_metadata", "publisher_step_id")

zenml/zen_stores/schemas/__init__.py

@@ -39,7 +39,10 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 from zenml.zen_stores.schemas.pipeline_schemas import PipelineSchema
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
-from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+from zenml.zen_stores.schemas.run_metadata_schemas import (
+    RunMetadataResourceSchema,
+    RunMetadataSchema,
+)
 from zenml.zen_stores.schemas.schedule_schema import ScheduleSchema
 from zenml.zen_stores.schemas.secret_schemas import SecretSchema
 from zenml.zen_stores.schemas.service_schemas import ServiceSchema
@@ -90,6 +93,7 @@ __all__ = [
     "PipelineDeploymentSchema",
     "PipelineRunSchema",
     "PipelineSchema",
+    "RunMetadataResourceSchema",
     "RunMetadataSchema",
     "ScheduleSchema",
     "SecretSchema",

zenml/zen_stores/schemas/artifact_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of artifact table."""
 
-import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, List, Optional
 from uuid import UUID
@@ -50,6 +49,7 @@ from zenml.zen_stores.schemas.step_run_schemas import (
     StepRunOutputArtifactSchema,
 )
 from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import RunMetadataInterface
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
@@ -59,7 +59,9 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.model_schemas import (
         ModelVersionArtifactSchema,
     )
-    from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataResourceSchema,
+    )
     from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 
 
@@ -171,7 +173,7 @@ class ArtifactSchema(NamedSchema, table=True):
         return self
 
 
-class ArtifactVersionSchema(BaseSchema, table=True):
+class ArtifactVersionSchema(BaseSchema, RunMetadataInterface, table=True):
     """SQL Model for artifact versions."""
 
     __tablename__ = "artifact_version"
@@ -242,12 +244,12 @@ class ArtifactVersionSchema(BaseSchema, table=True):
     workspace: "WorkspaceSchema" = Relationship(
         back_populates="artifact_versions"
     )
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="artifact_version",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="artifact_versions",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ArtifactVersionSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     output_of_step_runs: List["StepRunOutputArtifactSchema"] = Relationship(
@@ -376,9 +378,7 @@ class ArtifactVersionSchema(BaseSchema, table=True):
                 workspace=self.workspace.to_model(),
                 producer_step_run_id=producer_step_run_id,
                 visualizations=[v.to_model() for v in self.visualizations],
-                run_metadata={
-                    m.key: json.loads(m.value) for m in self.run_metadata
-                },
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None

zenml/zen_stores/schemas/component_schemas.py

@@ -56,7 +56,6 @@ class StackComponentSchema(NamedSchema, table=True):
     flavor: str
     configuration: bytes
     labels: Optional[bytes]
-    component_spec_path: Optional[str]
 
     workspace_id: UUID = build_foreign_key_field(
         source=__tablename__,
@@ -135,7 +134,6 @@ class StackComponentSchema(NamedSchema, table=True):
             name=request.name,
             workspace_id=request.workspace,
             user_id=request.user,
-            component_spec_path=request.component_spec_path,
             type=request.type,
             flavor=request.flavor,
             configuration=base64.b64encode(
@@ -218,7 +216,6 @@ class StackComponentSchema(NamedSchema, table=True):
                 labels=json.loads(base64.b64decode(self.labels).decode())
                 if self.labels
                 else None,
-                component_spec_path=self.component_spec_path,
                 connector_resource_id=self.connector_resource_id,
                 connector=self.connector.to_model()
                 if self.connector

zenml/zen_stores/schemas/model_schemas.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """SQLModel implementation of model tables."""
 
-import json
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 from uuid import UUID
@@ -51,11 +50,16 @@ from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
 from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
 from zenml.zen_stores.schemas.constants import MODEL_VERSION_TABLENAME
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
-from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+from zenml.zen_stores.schemas.run_metadata_schemas import (
+    RunMetadataResourceSchema,
+)
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 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 (
+    RunMetadataInterface,
+    get_page_from_list,
+)
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
@@ -219,7 +223,7 @@ class ModelSchema(NamedSchema, table=True):
         return self
 
 
-class ModelVersionSchema(NamedSchema, table=True):
+class ModelVersionSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for model version."""
 
     __tablename__ = MODEL_VERSION_TABLENAME
@@ -299,12 +303,12 @@ class ModelVersionSchema(NamedSchema, table=True):
     description: str = Field(sa_column=Column(TEXT, nullable=True))
     stage: str = Field(sa_column=Column(TEXT, nullable=True))
 
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="model_version",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="model_versions",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ModelVersionSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     pipeline_runs: List["PipelineRunSchema"] = Relationship(
@@ -402,9 +406,7 @@ class ModelVersionSchema(NamedSchema, table=True):
             metadata = ModelVersionResponseMetadata(
                 workspace=self.workspace.to_model(),
                 description=self.description,
-                run_metadata={
-                    rm.key: json.loads(rm.value) for rm in self.run_metadata
-                },
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None

zenml/zen_stores/schemas/pipeline_run_schemas.py

@@ -15,7 +15,7 @@
 
 import json
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 from uuid import UUID
 
 from pydantic import ConfigDict
@@ -34,6 +34,7 @@ from zenml.models import (
     PipelineRunResponseBody,
     PipelineRunResponseMetadata,
     PipelineRunUpdate,
+    RunMetadataEntry,
 )
 from zenml.models.v2.core.pipeline_run import PipelineRunResponseResources
 from zenml.zen_stores.schemas.base_schemas import NamedSchema
@@ -48,6 +49,7 @@ from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
 from zenml.zen_stores.schemas.stack_schemas import StackSchema
 from zenml.zen_stores.schemas.trigger_schemas import TriggerExecutionSchema
 from zenml.zen_stores.schemas.user_schemas import UserSchema
+from zenml.zen_stores.schemas.utils import RunMetadataInterface
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
@@ -56,13 +58,15 @@ if TYPE_CHECKING:
         ModelVersionPipelineRunSchema,
         ModelVersionSchema,
     )
-    from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataResourceSchema,
+    )
     from zenml.zen_stores.schemas.service_schemas import ServiceSchema
     from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
     from zenml.zen_stores.schemas.tag_schemas import TagResourceSchema
 
 
-class PipelineRunSchema(NamedSchema, table=True):
+class PipelineRunSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for pipeline runs."""
 
     __tablename__ = "pipeline_run"
@@ -136,12 +140,12 @@ class PipelineRunSchema(NamedSchema, table=True):
     )
     workspace: "WorkspaceSchema" = Relationship(back_populates="runs")
     user: Optional["UserSchema"] = Relationship(back_populates="runs")
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="pipeline_run",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="pipeline_runs",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataSchema.resource_id)==PipelineRunSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==PipelineRunSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     logs: Optional["LogsSchema"] = Relationship(
@@ -249,6 +253,24 @@ class PipelineRunSchema(NamedSchema, table=True):
             model_version_id=request.model_version_id,
         )
 
+    def fetch_metadata_collection(self) -> Dict[str, List[RunMetadataEntry]]:
+        """Fetches all the metadata entries related to the pipeline run.
+
+        Returns:
+            a dictionary, where the key is the key of the metadata entry
+                and the values represent the list of entries with this key.
+        """
+        # Fetch the metadata related to this run
+        metadata_collection = super().fetch_metadata_collection()
+
+        # Fetch the metadata related to the steps of this run
+        for s in self.step_runs:
+            step_metadata = s.fetch_metadata_collection()
+            for k, v in step_metadata.items():
+                metadata_collection[f"{s.name}::{k}"] = v
+
+        return metadata_collection
+
     def to_model(
         self,
         include_metadata: bool = False,
@@ -275,11 +297,6 @@ class PipelineRunSchema(NamedSchema, table=True):
             else {}
         )
 
-        run_metadata = {
-            metadata_schema.key: json.loads(metadata_schema.value)
-            for metadata_schema in self.run_metadata
-        }
-
         if self.deployment is not None:
             deployment = self.deployment.to_model()
 
@@ -350,13 +367,13 @@ class PipelineRunSchema(NamedSchema, table=True):
 
             steps = {step.name: step.to_model() for step in self.step_runs}
 
-            steps_substitutions = {
+            step_substitutions = {
                 step_name: step.config.substitutions
                 for step_name, step in steps.items()
             }
             metadata = PipelineRunResponseMetadata(
                 workspace=self.workspace.to_model(),
-                run_metadata=run_metadata,
+                run_metadata=self.fetch_metadata(),
                 config=config,
                 steps=steps,
                 start_time=self.start_time,
@@ -371,7 +388,7 @@ class PipelineRunSchema(NamedSchema, table=True):
                 if self.deployment
                 else None,
                 is_templatable=is_templatable,
-                steps_substitutions=steps_substitutions,
+                step_substitutions=step_substitutions,
             )
 
         resources = None

zenml/zen_stores/schemas/run_metadata_schemas.py

@@ -14,15 +14,16 @@
 """SQLModel implementation of pipeline run metadata tables."""
 
 from typing import TYPE_CHECKING, List, Optional
-from uuid import UUID
+from uuid import UUID, uuid4
 
 from sqlalchemy import TEXT, VARCHAR, Column
-from sqlmodel import Field, Relationship
+from sqlmodel import Field, Relationship, SQLModel
 
 from zenml.enums import MetadataResourceTypes
 from zenml.zen_stores.schemas.base_schemas import BaseSchema
 from zenml.zen_stores.schemas.component_schemas import StackComponentSchema
 from zenml.zen_stores.schemas.schema_utils import build_foreign_key_field
+from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
 from zenml.zen_stores.schemas.user_schemas import UserSchema
 from zenml.zen_stores.schemas.workspace_schemas import WorkspaceSchema
 
@@ -30,7 +31,6 @@ if TYPE_CHECKING:
     from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
     from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
     from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
-    from zenml.zen_stores.schemas.step_run_schemas import StepRunSchema
 
 
 class RunMetadataSchema(BaseSchema, table=True):
@@ -38,35 +38,10 @@ class RunMetadataSchema(BaseSchema, table=True):
 
     __tablename__ = "run_metadata"
 
-    resource_id: UUID
-    resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
-    pipeline_run: List["PipelineRunSchema"] = Relationship(
-        back_populates="run_metadata",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataSchema.resource_id)==PipelineRunSchema.id)",
-            overlaps="run_metadata,step_run,artifact_version,model_version",
-        ),
-    )
-    step_run: List["StepRunSchema"] = Relationship(
-        back_populates="run_metadata",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataSchema.resource_id)==StepRunSchema.id)",
-            overlaps="run_metadata,pipeline_run,artifact_version,model_version",
-        ),
-    )
-    artifact_version: List["ArtifactVersionSchema"] = Relationship(
-        back_populates="run_metadata",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ArtifactVersionSchema.id)",
-            overlaps="run_metadata,pipeline_run,step_run,model_version",
-        ),
-    )
-    model_version: List["ModelVersionSchema"] = Relationship(
+    # Relationship to link to resources
+    resources: List["RunMetadataResourceSchema"] = Relationship(
         back_populates="run_metadata",
-        sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataSchema.resource_id)==ModelVersionSchema.id)",
-            overlaps="run_metadata,pipeline_run,step_run,artifact_version",
-        ),
+        sa_relationship_kwargs={"cascade": "delete"},
     )
     stack_component_id: Optional[UUID] = build_foreign_key_field(
         source=__tablename__,
@@ -103,3 +78,63 @@ class RunMetadataSchema(BaseSchema, table=True):
     key: str
     value: str = Field(sa_column=Column(TEXT, nullable=False))
     type: str
+
+    publisher_step_id: Optional[UUID] = build_foreign_key_field(
+        source=__tablename__,
+        target=StepRunSchema.__tablename__,
+        source_column="publisher_step_id",
+        target_column="id",
+        ondelete="SET NULL",
+        nullable=True,
+    )
+
+
+class RunMetadataResourceSchema(SQLModel, table=True):
+    """Table for linking resources to run metadata entries."""
+
+    __tablename__ = "run_metadata_resource"
+
+    id: UUID = Field(default_factory=uuid4, primary_key=True)
+    resource_id: UUID
+    resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
+    run_metadata_id: UUID = build_foreign_key_field(
+        source=__tablename__,
+        target=RunMetadataSchema.__tablename__,
+        source_column="run_metadata_id",
+        target_column="id",
+        ondelete="CASCADE",
+        nullable=False,
+    )
+
+    # Relationship back to the base metadata table
+    run_metadata: RunMetadataSchema = Relationship(back_populates="resources")
+
+    # Relationship to link specific resource types
+    pipeline_runs: List["PipelineRunSchema"] = Relationship(
+        back_populates="run_metadata_resources",
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.PIPELINE_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==PipelineRunSchema.id)",
+            overlaps="run_metadata_resources,step_runs,artifact_versions,model_versions",
+        ),
+    )
+    step_runs: List["StepRunSchema"] = Relationship(
+        back_populates="run_metadata_resources",
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
+            overlaps="run_metadata_resources,pipeline_runs,artifact_versions,model_versions",
+        ),
+    )
+    artifact_versions: List["ArtifactVersionSchema"] = Relationship(
+        back_populates="run_metadata_resources",
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.ARTIFACT_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ArtifactVersionSchema.id)",
+            overlaps="run_metadata_resources,pipeline_runs,step_runs,model_versions",
+        ),
+    )
+    model_versions: List["ModelVersionSchema"] = Relationship(
+        back_populates="run_metadata_resources",
+        sa_relationship_kwargs=dict(
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.MODEL_VERSION.value}', foreign(RunMetadataResourceSchema.resource_id)==ModelVersionSchema.id)",
+            overlaps="run_metadata_resources,pipeline_runs,step_runs,artifact_versions",
+        ),
+    )

zenml/zen_stores/schemas/step_run_schemas.py

@@ -51,16 +51,19 @@ from zenml.zen_stores.schemas.pipeline_deployment_schemas import (
 from zenml.zen_stores.schemas.pipeline_run_schemas import PipelineRunSchema
 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.workspace_schemas import WorkspaceSchema
 
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas.artifact_schemas import ArtifactVersionSchema
     from zenml.zen_stores.schemas.logs_schemas import LogsSchema
     from zenml.zen_stores.schemas.model_schemas import ModelVersionSchema
-    from zenml.zen_stores.schemas.run_metadata_schemas import RunMetadataSchema
+    from zenml.zen_stores.schemas.run_metadata_schemas import (
+        RunMetadataResourceSchema,
+    )
 
 
-class StepRunSchema(NamedSchema, table=True):
+class StepRunSchema(NamedSchema, RunMetadataInterface, table=True):
     """SQL Model for steps of pipeline runs."""
 
     __tablename__ = "step_run"
@@ -140,12 +143,12 @@ class StepRunSchema(NamedSchema, table=True):
     deployment: Optional["PipelineDeploymentSchema"] = Relationship(
         back_populates="step_runs"
     )
-    run_metadata: List["RunMetadataSchema"] = Relationship(
-        back_populates="step_run",
+    run_metadata_resources: List["RunMetadataResourceSchema"] = Relationship(
+        back_populates="step_runs",
         sa_relationship_kwargs=dict(
-            primaryjoin=f"and_(RunMetadataSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataSchema.resource_id)==StepRunSchema.id)",
+            primaryjoin=f"and_(RunMetadataResourceSchema.resource_type=='{MetadataResourceTypes.STEP_RUN.value}', foreign(RunMetadataResourceSchema.resource_id)==StepRunSchema.id)",
             cascade="delete",
-            overlaps="run_metadata",
+            overlaps="run_metadata_resources",
         ),
     )
     input_artifacts: List["StepRunInputArtifactSchema"] = Relationship(
@@ -170,6 +173,9 @@ class StepRunSchema(NamedSchema, table=True):
     model_version: "ModelVersionSchema" = Relationship(
         back_populates="step_runs",
     )
+    original_step_run: Optional["StepRunSchema"] = Relationship(
+        sa_relationship_kwargs={"remote_side": "StepRunSchema.id"}
+    )
 
     model_config = ConfigDict(protected_namespaces=())  # type: ignore[assignment]
 
@@ -222,11 +228,6 @@ class StepRunSchema(NamedSchema, table=True):
             RuntimeError: If the step run schema does not have a deployment_id
                 or a step_configuration.
         """
-        run_metadata = {
-            metadata_schema.key: json.loads(metadata_schema.value)
-            for metadata_schema in self.run_metadata
-        }
-
         input_artifacts = {
             artifact.name: StepRunInputResponse(
                 input_type=StepRunInputArtifactType(artifact.type),
@@ -313,7 +314,7 @@ class StepRunSchema(NamedSchema, table=True):
                 pipeline_run_id=self.pipeline_run_id,
                 original_step_run_id=self.original_step_run_id,
                 parent_step_ids=[p.parent_id for p in self.parents],
-                run_metadata=run_metadata,
+                run_metadata=self.fetch_metadata(),
             )
 
         resources = None