zenml-nightly 0.75.0.dev20250311__py3-none-any.whl → 0.75.0.dev20250313__py3-none-any.whl

zenml/zen_server/template_execution/utils.py

@@ -149,9 +149,10 @@ def run_template(
     )
 
     def _task() -> None:
-        pypi_requirements, apt_packages = (
-            requirements_utils.get_requirements_for_stack(stack=stack)
-        )
+        (
+            pypi_requirements,
+            apt_packages,
+        ) = requirements_utils.get_requirements_for_stack(stack=stack)
 
         if build.python_version:
             version_info = version.parse(build.python_version)
@@ -406,7 +407,6 @@ def deployment_request_from_template(
     assert deployment.stack
     assert deployment.build
     deployment_request = PipelineDeploymentRequest(
-        user=user_id,
         workspace=deployment.workspace.id,
         run_name_template=config.run_name
         or get_default_run_name(pipeline_name=pipeline_configuration.name),
@@ -463,6 +463,7 @@ def get_pipeline_run_analytics_metadata(
     }
 
     return {
+        "workspace_id": deployment.workspace.id,
         "store_type": "rest",  # This method is called from within a REST endpoint
         **stack_metadata,
         "total_steps": len(deployment.step_configurations),

zenml/zen_server/utils.py

@@ -26,9 +26,11 @@ from typing import (
     Tuple,
     Type,
     TypeVar,
+    Union,
     cast,
 )
 from urllib.parse import urlparse
+from uuid import UUID
 
 from pydantic import BaseModel, ValidationError
 
@@ -44,6 +46,7 @@ from zenml.constants import (
 from zenml.enums import StoreType
 from zenml.exceptions import IllegalOperationError, OAuthError
 from zenml.logger import get_logger
+from zenml.models.v2.base.scoped import WorkspaceScopedFilter
 from zenml.plugins.plugin_flavor_registry import PluginFlavorRegistry
 from zenml.zen_server.cache import MemoryCache
 from zenml.zen_server.deploy.deployment import (
@@ -635,3 +638,21 @@ def get_zenml_headers() -> Dict[str, str]:
         headers["zenml-server-url"] = config.server_url
 
     return headers
+
+
+def set_filter_workspace_scope(
+    filter_model: WorkspaceScopedFilter,
+    workspace_name_or_id: Optional[Union[UUID, str]] = None,
+) -> None:
+    """Set the workspace scope of the filter model.
+
+    Args:
+        filter_model: The filter model to set the scope for.
+        workspace_name_or_id: The workspace to set the scope for. If not
+            provided, the workspace scope is determined from the request
+            workspace filter or the default workspace, in that order.
+    """
+    zen_store().set_filter_workspace_id(
+        filter_model=filter_model,
+        workspace_name_or_id=workspace_name_or_id,
+    )

zenml/zen_server/zen_server_api.py

@@ -73,6 +73,7 @@ from zenml.zen_server.routers import (
     pipeline_deployments_endpoints,
     pipelines_endpoints,
     plugin_endpoints,
+    run_metadata_endpoints,
     run_templates_endpoints,
     runs_endpoints,
     schedule_endpoints,
@@ -85,6 +86,7 @@ from zenml.zen_server.routers import (
     stack_deployment_endpoints,
     stacks_endpoints,
     steps_endpoints,
+    tag_resource_endpoints,
     tags_endpoints,
     triggers_endpoints,
     users_endpoints,
@@ -472,6 +474,7 @@ app.include_router(pipelines_endpoints.router)
 app.include_router(pipeline_builds_endpoints.router)
 app.include_router(pipeline_deployments_endpoints.router)
 app.include_router(runs_endpoints.router)
+app.include_router(run_metadata_endpoints.router)
 app.include_router(run_templates_endpoints.router)
 app.include_router(schedule_endpoints.router)
 app.include_router(secrets_endpoints.router)
@@ -487,6 +490,7 @@ app.include_router(stack_components_endpoints.router)
 app.include_router(stack_components_endpoints.types_router)
 app.include_router(steps_endpoints.router)
 app.include_router(tags_endpoints.router)
+app.include_router(tag_resource_endpoints.router)
 app.include_router(triggers_endpoints.router)
 app.include_router(users_endpoints.router)
 app.include_router(users_endpoints.current_user_router)

zenml/zen_stores/base_zen_store.py

@@ -295,16 +295,16 @@ class BaseZenStore(
         active_workspace_name_or_id: Optional[Union[str, UUID]] = None,
         active_stack_id: Optional[UUID] = None,
         config_name: str = "",
-    ) -> Tuple[WorkspaceResponse, StackResponse]:
+    ) -> Tuple[Optional[WorkspaceResponse], StackResponse]:
         """Validate the active configuration.
 
         Call this method to validate the supplied active workspace and active
         stack values.
 
-        This method is guaranteed to return valid workspace ID and stack ID
-        values. If the supplied workspace and stack are not set or are not valid
-        (e.g. they do not exist or are not accessible), the default workspace and
-        default workspace stack will be returned in their stead.
+        This method returns a valid workspace and stack values. If the
+        supplied workspace and stack are not set or are not valid (e.g. they
+        do not exist or are not accessible), the default workspace and default
+        stack will be returned in their stead.
 
         Args:
             active_workspace_name_or_id: The name or ID of the active workspace.
@@ -315,28 +315,34 @@ class BaseZenStore(
         Returns:
             A tuple containing the active workspace and active stack.
         """
-        active_workspace: WorkspaceResponse
+        active_workspace: Optional[WorkspaceResponse] = None
 
         if active_workspace_name_or_id:
             try:
                 active_workspace = self.get_workspace(
                     active_workspace_name_or_id
                 )
-            except KeyError:
-                active_workspace = self._get_default_workspace()
-
+            except (KeyError, IllegalOperationError):
+                active_workspace_name_or_id = None
                 logger.warning(
                     f"The current {config_name} active workspace is no longer "
-                    f"available. Resetting the active workspace to "
-                    f"'{active_workspace.name}'."
+                    f"available."
                 )
-        else:
-            active_workspace = self._get_default_workspace()
 
-            logger.info(
-                f"Setting the {config_name} active workspace "
-                f"to '{active_workspace.name}'."
-            )
+        if active_workspace is None:
+            try:
+                active_workspace = self._get_default_workspace()
+            except (KeyError, IllegalOperationError):
+                logger.warning(
+                    "An active workspace is not set. Please set the active "
+                    "workspace by running `zenml workspace set "
+                    "<workspace-name>`."
+                )
+            else:
+                logger.info(
+                    f"Setting the {config_name} active workspace "
+                    f"to '{active_workspace.name}'."
+                )
 
         active_stack: StackResponse
 
@@ -351,28 +357,14 @@ class BaseZenStore(
                     "Resetting the active stack to default.",
                     config_name,
                 )
-                active_stack = self._get_default_stack(
-                    workspace_id=active_workspace.id
-                )
-            else:
-                if active_stack.workspace.id != active_workspace.id:
-                    logger.warning(
-                        "The current %s active stack is not part of the active "
-                        "workspace. Resetting the active stack to default.",
-                        config_name,
-                    )
-                    active_stack = self._get_default_stack(
-                        workspace_id=active_workspace.id
-                    )
+                active_stack = self._get_default_stack()
 
         else:
             logger.warning(
                 "Setting the %s active stack to default.",
                 config_name,
             )
-            active_stack = self._get_default_stack(
-                workspace_id=active_workspace.id
-            )
+            active_stack = self._get_default_stack()
 
         return active_workspace, active_stack
 
@@ -462,29 +454,22 @@ class BaseZenStore(
 
     def _get_default_stack(
         self,
-        workspace_id: UUID,
     ) -> StackResponse:
-        """Get the default stack for a user in a workspace.
-
-        Args:
-            workspace_id: ID of the workspace.
+        """Get the default stack.
 
         Returns:
-            The default stack in the workspace.
+            The default stack.
 
         Raises:
-            KeyError: if the workspace or default stack doesn't exist.
+            KeyError: if the default stack doesn't exist.
         """
         default_stacks = self.list_stacks(
             StackFilter(
-                workspace_id=workspace_id,
                 name=DEFAULT_STACK_AND_COMPONENT_NAME,
             )
         )
         if default_stacks.total == 0:
-            raise KeyError(
-                f"No default stack found in workspace {workspace_id}."
-            )
+            raise KeyError("No default stack found.")
         return default_stacks.items[0]
 
     def get_external_user(self, user_id: UUID) -> UserResponse:

zenml/zen_stores/migrations/versions/0392807467dc_add_build_duration.py

@@ -0,0 +1,34 @@
+"""Add build duration [0392807467dc].
+
+Revision ID: 0392807467dc
+Revises: 0.75.0
+Create Date: 2025-02-28 15:50:06.960529
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "0392807467dc"
+down_revision = "0.75.0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_build", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("duration", sa.Integer(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("pipeline_build", schema=None) as batch_op:
+        batch_op.drop_column("duration")
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/1cb6477f72d6_move_artifact_save_type.py

@@ -23,7 +23,8 @@ def upgrade() -> None:
         batch_op.add_column(sa.Column("save_type", sa.TEXT(), nullable=True))
 
     # Step 2: Move data from step_run_output_artifact.type to artifact_version.save_type
-    op.execute("""
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = (
             SELECT max(step_run_output_artifact.type)
@@ -31,17 +32,22 @@ def upgrade() -> None:
             WHERE step_run_output_artifact.artifact_id = artifact_version.id
             GROUP BY artifact_id
         )
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = 'step_output'
         WHERE artifact_version.save_type = 'default'
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE artifact_version
         SET save_type = 'external'
         WHERE save_type is NULL
-    """)
+    """
+    )
 
     # # Step 3: Set save_type to non-nullable
     with op.batch_alter_table("artifact_version", schema=None) as batch_op:
@@ -69,7 +75,8 @@ def downgrade() -> None:
         )
 
     # Move data back from artifact_version.save_type to step_run_output_artifact.type
-    op.execute("""
+    op.execute(
+        """
         UPDATE step_run_output_artifact
         SET type = (
             SELECT max(artifact_version.save_type)
@@ -77,12 +84,15 @@ def downgrade() -> None:
             WHERE step_run_output_artifact.artifact_id = artifact_version.id
             GROUP BY artifact_id
         )
-    """)
-    op.execute("""
+    """
+    )
+    op.execute(
+        """
         UPDATE step_run_output_artifact
         SET type = 'default'
         WHERE step_run_output_artifact.type = 'step_output'
-    """)
+    """
+    )
 
     # Set type to non-nullable
     with op.batch_alter_table(

zenml/zen_stores/migrations/versions/1f9d1cd00b90_add_unique_name_constraints.py

@@ -0,0 +1,231 @@
+"""add unique name constraints [1f9d1cd00b90].
+
+Revision ID: 1f9d1cd00b90
+Revises: f76a368a25a5
+Create Date: 2025-02-22 20:18:34.258987
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.orm import Session
+
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+# revision identifiers, used by Alembic.
+revision = "1f9d1cd00b90"
+down_revision = "f76a368a25a5"
+branch_labels = None
+depends_on = None
+
+
+def resolve_duplicate_names(
+    table_name: str, other_columns: list[str], session: Session
+) -> None:
+    """Resolve duplicate entities.
+
+    Args:
+        table_name: The name of the table to resolve duplicate entities for.
+        other_columns: The columns that are part of the unique constraint,
+            excluding the name column.
+        session: The SQLAlchemy session to use.
+    """
+    columns = ["name"] + other_columns
+    duplicates = session.execute(
+        sa.text(
+            f"""
+                SELECT id, name
+                FROM `{table_name}`
+                WHERE ({", ".join(columns)}) IN (
+                    SELECT {", ".join(columns)}
+                    FROM `{table_name}`
+                    GROUP BY {", ".join(columns)}
+                    HAVING COUNT(*) > 1
+                )
+            """  # nosec B608
+        )
+    )
+    for id_, name in list(duplicates)[1:]:
+        logger.warning(f"Duplicate {table_name}: {name} (id: {id_})")
+        session.execute(
+            sa.text(
+                f"""
+                    UPDATE {table_name}
+                    SET name = :new_name
+                    WHERE id = :id_
+                """  # nosec B608
+            ),
+            params={"id_": id_, "new_name": f"{name}_{id_[:6]}"},
+        )
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    resolve_duplicate_names("action", ["workspace_id"], session)
+
+    with op.batch_alter_table("action", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_action_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("api_key", ["service_account_id"], session)
+
+    with op.batch_alter_table("api_key", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_api_key_name_in_service_account",
+            ["name", "service_account_id"],
+        )
+
+    resolve_duplicate_names("artifact", ["workspace_id"], session)
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_artifact_name", type_="unique")
+        batch_op.create_unique_constraint(
+            "unique_artifact_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("code_repository", ["workspace_id"], session)
+
+    with op.batch_alter_table("code_repository", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_code_repository_name_in_workspace",
+            ["name", "workspace_id"],
+        )
+
+    resolve_duplicate_names("event_source", ["workspace_id"], session)
+
+    with op.batch_alter_table("event_source", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_event_source_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("flavor", ["type"], session)
+
+    with op.batch_alter_table("flavor", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_flavor_name_and_type", ["name", "type"]
+        )
+
+    resolve_duplicate_names("schedule", ["workspace_id"], session)
+
+    with op.batch_alter_table("schedule", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_schedule_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("secret", ["private", "user_id"], session)
+
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_secret_name_private_scope_user",
+            ["name", "private", "user_id"],
+        )
+
+    resolve_duplicate_names("service_connector", [], session)
+
+    with op.batch_alter_table("service_connector", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_service_connector_name", ["name"]
+        )
+
+    resolve_duplicate_names("stack", [], session)
+
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_stack_name", ["name"])
+
+    resolve_duplicate_names("stack_component", ["type"], session)
+
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_component_name_and_type", ["name", "type"]
+        )
+
+    resolve_duplicate_names("tag", [], session)
+
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_tag_name", ["name"])
+
+    resolve_duplicate_names("trigger", ["workspace_id"], session)
+
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_trigger_name_in_workspace", ["name", "workspace_id"]
+        )
+
+    resolve_duplicate_names("workspace", [], session)
+
+    with op.batch_alter_table("workspace", schema=None) as batch_op:
+        batch_op.create_unique_constraint("unique_workspace_name", ["name"])
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("workspace", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_workspace_name", type_="unique")
+
+    with op.batch_alter_table("trigger", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_trigger_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_tag_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_component_name_and_type", type_="unique"
+        )
+
+    with op.batch_alter_table("stack", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_stack_name", type_="unique")
+
+    with op.batch_alter_table("service_connector", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_service_connector_name", type_="unique"
+        )
+
+    with op.batch_alter_table("secret", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_secret_name_and_private_scope", type_="unique"
+        )
+
+    with op.batch_alter_table("schedule", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_schedule_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("flavor", schema=None) as batch_op:
+        batch_op.drop_constraint("unique_flavor_name_and_type", type_="unique")
+
+    with op.batch_alter_table("event_source", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_event_source_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("code_repository", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_code_repository_name_in_workspace", type_="unique"
+        )
+
+    with op.batch_alter_table("artifact", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_artifact_name_in_workspace", type_="unique"
+        )
+        batch_op.create_unique_constraint("unique_artifact_name", ["name"])
+
+    with op.batch_alter_table("api_key", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_api_key_name_in_service_account", type_="unique"
+        )
+
+    with op.batch_alter_table("action", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_action_name_in_workspace", type_="unique"
+        )

zenml/zen_stores/migrations/versions/288f4fb6e112_make_tags_user_scoped.py

@@ -0,0 +1,74 @@
+"""make tags user scoped [288f4fb6e112].
+
+Revision ID: 288f4fb6e112
+Revises: 3b1776345020
+Create Date: 2025-02-19 15:16:42.954792
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+from sqlalchemy.orm import Session
+
+# revision identifiers, used by Alembic.
+revision = "288f4fb6e112"
+down_revision = "3b1776345020"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        # First add columns as nullable
+        batch_op.add_column(
+            sa.Column("user_id", sqlmodel.sql.sqltypes.GUID(), nullable=True)
+        )
+
+        # Add foreign key constraints
+        batch_op.create_foreign_key(
+            "fk_tag_user_id_user",
+            "user",
+            ["user_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+    bind = op.get_bind()
+    session = Session(bind=bind)
+
+    tags = session.execute(
+        sa.text("""
+            SELECT t.id, tr.resource_id, tr.resource_type
+            FROM tag t
+            JOIN tag_resource tr ON t.id = tr.tag_id
+        """)
+    )
+
+    tag_ids = []
+    for tag_id, resource_id, resource_type in tags:
+        if tag_id in tag_ids:
+            continue
+        tag_ids.append(tag_id)
+        session.execute(
+            sa.text(
+                f"""
+                    UPDATE tag
+                    SET user_id = (
+                        SELECT r.user_id
+                        FROM {resource_type} r
+                        WHERE r.id = :resource_id
+                    )
+                    WHERE id = :tag_id
+                """  # nosec B608
+            ),
+            params={"resource_id": resource_id, "tag_id": tag_id},
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("tag", schema=None) as batch_op:
+        batch_op.drop_constraint("fk_tag_user_id_user", type_="foreignkey")
+        batch_op.drop_column("user_id")

zenml/zen_stores/migrations/versions/2e695a26fe7a_add_user_default_workspace.py

@@ -0,0 +1,45 @@
+"""add user default workspace [2e695a26fe7a].
+
+Revision ID: 2e695a26fe7a
+Revises: 1f9d1cd00b90
+Create Date: 2025-02-24 18:19:43.121393
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "2e695a26fe7a"
+down_revision = "1f9d1cd00b90"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "default_workspace_id",
+                sqlmodel.sql.sqltypes.GUID(),
+                nullable=True,
+            )
+        )
+        batch_op.create_foreign_key(
+            "fk_user_default_workspace_id_workspace",
+            "workspace",
+            ["default_workspace_id"],
+            ["id"],
+            ondelete="SET NULL",
+        )
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    with op.batch_alter_table("user", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "fk_user_default_workspace_id_workspace", type_="foreignkey"
+        )
+        batch_op.drop_column("default_workspace_id")