zenml-nightly 0.70.0.dev20241125__py3-none-any.whl → 0.71.0.dev20241220__py3-none-any.whl

zenml/zen_server/routers/auth_endpoints.py

@@ -41,7 +41,6 @@ from zenml.constants import (
 from zenml.enums import (
     APITokenType,
     AuthScheme,
-    ExecutionStatus,
     OAuthDeviceStatus,
     OAuthGrantTypes,
 )
@@ -452,6 +451,7 @@ def device_authorization(
 @handle_exceptions
 def api_token(
     token_type: APITokenType = APITokenType.GENERIC,
+    expires_in: Optional[int] = None,
     schedule_id: Optional[UUID] = None,
     pipeline_run_id: Optional[UUID] = None,
     step_run_id: Optional[UUID] = None,
@@ -463,7 +463,8 @@ def api_token(
     of API tokens are supported:
 
     * Generic API token: This token is short-lived and can be used for
-    generic automation tasks.
+    generic automation tasks. The expiration can be set by the user, but the
+    server will impose a maximum expiration time.
     * Workload API token: This token is scoped to a specific pipeline run, step
     run or schedule and is used by pipeline workloads to authenticate with the
     server. A pipeline run ID, step run ID or schedule ID must be provided and
@@ -475,6 +476,10 @@ def api_token(
 
     Args:
         token_type: The type of API token to generate.
+        expires_in: The expiration time of the generic API token in seconds.
+            If not set, the server will use the default expiration time for
+            generic API tokens. The server also imposes a maximum expiration
+            time.
         schedule_id: The ID of the schedule to scope the workload API token to.
         pipeline_run_id: The ID of the pipeline run to scope the workload API
             token to.
@@ -502,9 +507,19 @@ def api_token(
 
         config = server_config()
 
+        if not expires_in:
+            expires_in = config.generic_api_token_lifetime
+
+        if expires_in > config.generic_api_token_max_lifetime:
+            raise ValueError(
+                f"The maximum expiration time for generic API tokens allowed "
+                f"by this server is {config.generic_api_token_max_lifetime} "
+                "seconds."
+            )
+
         return generate_access_token(
             user_id=token.user_id,
-            expires_in=config.generic_api_token_lifetime,
+            expires_in=expires_in,
         ).access_token
 
     verify_permission(
@@ -573,10 +588,7 @@ def api_token(
                 "security reasons."
             )
 
-        if pipeline_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if pipeline_run.status.is_finished:
             raise ValueError(
                 f"The execution of pipeline run {pipeline_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "
@@ -593,10 +605,7 @@ def api_token(
                 "be generated for non-existent step runs for security reasons."
             )
 
-        if step_run.status in [
-            ExecutionStatus.FAILED,
-            ExecutionStatus.COMPLETED,
-        ]:
+        if step_run.status.is_finished:
             raise ValueError(
                 f"The execution of step run {step_run_id} has already "
                 "concluded and API tokens can no longer be generated for it "
@@ -611,4 +620,6 @@ def api_token(
         schedule_id=schedule_id,
         pipeline_run_id=pipeline_run_id,
         step_run_id=step_run_id,
+        # Never expire the token
+        expires_in=0,
     ).access_token

zenml/zen_server/routers/steps_endpoints.py

@@ -142,10 +142,16 @@ def get_step(
     Returns:
         The step.
     """
-    step = zen_store().get_run_step(step_id, hydrate=hydrate)
+    # We always fetch the step hydrated because we need the pipeline_run_id
+    # for the permission checks. If the user requested an unhydrated response,
+    # we later remove the metadata
+    step = zen_store().get_run_step(step_id, hydrate=True)
     pipeline_run = zen_store().get_run(step.pipeline_run_id)
     verify_permission_for_model(pipeline_run, action=Action.READ)
 
+    if hydrate is False:
+        step.metadata = None
+
     return dehydrate_response_model(step)
 
 

zenml/zen_server/routers/users_endpoints.py

@@ -46,14 +46,10 @@ from zenml.zen_server.auth import (
 )
 from zenml.zen_server.exceptions import error_response
 from zenml.zen_server.rate_limit import RequestLimiter
-from zenml.zen_server.rbac.endpoint_utils import (
-    verify_permissions_and_create_entity,
-)
 from zenml.zen_server.rbac.models import Action, Resource, ResourceType
 from zenml.zen_server.rbac.utils import (
     dehydrate_page,
     dehydrate_response_model,
-    get_allowed_resource_ids,
     get_schema_for_resource_type,
     update_resource_membership,
     verify_permission_for_model,
@@ -112,17 +108,18 @@ def list_users(
     Returns:
         A list of all users.
     """
-    allowed_ids = get_allowed_resource_ids(resource_type=ResourceType.USER)
-    if allowed_ids is not None:
-        # Make sure users can see themselves
-        allowed_ids.add(auth_context.user.id)
-    else:
-        if not auth_context.user.is_admin and not server_config().rbac_enabled:
-            allowed_ids = {auth_context.user.id}
-
-    user_filter_model.configure_rbac(
-        authenticated_user_id=auth_context.user.id, id=allowed_ids
-    )
+    # allowed_ids = get_allowed_resource_ids(resource_type=ResourceType.USER)
+    # if allowed_ids is not None:
+    #     # Make sure users can see themselves
+    #     allowed_ids.add(auth_context.user.id)
+    # else:
+    #     if not auth_context.user.is_admin and not server_config().rbac_enabled:
+    #         allowed_ids = {auth_context.user.id}
+    if not auth_context.user.is_admin and not server_config().rbac_enabled:
+        user_filter_model.configure_rbac(
+            authenticated_user_id=auth_context.user.id,
+            id={auth_context.user.id},
+        )
 
     page = zen_store().list_users(
         user_filter_model=user_filter_model, hydrate=hydrate
@@ -175,11 +172,12 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             auth_context.user.is_admin, "create user"
         )
 
-        new_user = verify_permissions_and_create_entity(
-            request_model=user,
-            resource_type=ResourceType.USER,
-            create_method=zen_store().create_user,
-        )
+        # new_user = verify_permissions_and_create_entity(
+        #     request_model=user,
+        #     resource_type=ResourceType.USER,
+        #     create_method=zen_store().create_user,
+        # )
+        new_user = zen_store().create_user(user)
 
         # add back the original unhashed activation token, if generated, to
         # send it back to the client
@@ -217,10 +215,10 @@ def get_user(
         verify_admin_status_if_no_rbac(
             auth_context.user.is_admin, "get other user"
         )
-        verify_permission_for_model(
-            user,
-            action=Action.READ,
-        )
+        # verify_permission_for_model(
+        #     user,
+        #     action=Action.READ,
+        # )
 
     return dehydrate_response_model(user)
 
@@ -304,10 +302,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             verify_admin_status_if_no_rbac(
                 auth_context.user.is_admin, "update other user account"
             )
-            verify_permission_for_model(
-                user,
-                action=Action.UPDATE,
-            )
+            # verify_permission_for_model(
+            #     user,
+            #     action=Action.UPDATE,
+            # )
 
         # Validate a password change
         if user_update.password is not None:
@@ -497,10 +495,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
         verify_admin_status_if_no_rbac(
             auth_context.user.is_admin, "deactivate user"
         )
-        verify_permission_for_model(
-            user,
-            action=Action.UPDATE,
-        )
+        # verify_permission_for_model(
+        #     user,
+        #     action=Action.UPDATE,
+        # )
 
         user_update = UserUpdate(
             active=False,
@@ -548,10 +546,10 @@ if server_config().auth_scheme != AuthScheme.EXTERNAL:
             verify_admin_status_if_no_rbac(
                 auth_context.user.is_admin, "delete user"
             )
-            verify_permission_for_model(
-                user,
-                action=Action.DELETE,
-            )
+            # verify_permission_for_model(
+            #     user,
+            #     action=Action.DELETE,
+            # )
 
         zen_store().delete_user(user_name_or_id=user_name_or_id)
 
@@ -746,7 +744,7 @@ if server_config().rbac_enabled:
             KeyError: If no resource with the given type and ID exists.
         """
         user = zen_store().get_user(user_name_or_id)
-        verify_permission_for_model(user, action=Action.READ)
+        # verify_permission_for_model(user, action=Action.READ)
 
         if user.id == auth_context.user.id:
             raise ValueError(

zenml/zen_server/routers/workspaces_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for workspaces."""
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -98,13 +98,13 @@ from zenml.zen_server.feature_gate.endpoint_utils import (
 )
 from zenml.zen_server.rbac.endpoint_utils import (
     verify_permissions_and_create_entity,
-    verify_permissions_and_delete_entity,
-    verify_permissions_and_get_entity,
     verify_permissions_and_list_entities,
-    verify_permissions_and_update_entity,
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
+    batch_verify_permissions_for_models,
+    dehydrate_page,
+    dehydrate_response_model,
     get_allowed_resource_ids,
     verify_permission,
     verify_permission_for_model,
@@ -146,12 +146,10 @@ def list_workspaces(
     Returns:
         A list of workspaces.
     """
-    return verify_permissions_and_list_entities(
-        filter_model=workspace_filter_model,
-        resource_type=ResourceType.WORKSPACE,
-        list_method=zen_store().list_workspaces,
-        hydrate=hydrate,
+    workspaces = zen_store().list_workspaces(
+        workspace_filter_model, hydrate=hydrate
     )
+    return dehydrate_page(workspaces)
 
 
 @router.post(
@@ -160,7 +158,7 @@ def list_workspaces(
 )
 @handle_exceptions
 def create_workspace(
-    workspace: WorkspaceRequest,
+    workspace_request: WorkspaceRequest,
     _: AuthContext = Security(authorize),
 ) -> WorkspaceResponse:
     """Creates a workspace based on the requestBody.
@@ -168,16 +166,13 @@ def create_workspace(
     # noqa: DAR401
 
     Args:
-        workspace: Workspace to create.
+        workspace_request: Workspace to create.
 
     Returns:
         The created workspace.
     """
-    return verify_permissions_and_create_entity(
-        request_model=workspace,
-        resource_type=ResourceType.WORKSPACE,
-        create_method=zen_store().create_workspace,
-    )
+    workspace = zen_store().create_workspace(workspace_request)
+    return dehydrate_response_model(workspace)
 
 
 @router.get(
@@ -203,11 +198,10 @@ def get_workspace(
     Returns:
         The requested workspace.
     """
-    return verify_permissions_and_get_entity(
-        id=workspace_name_or_id,
-        get_method=zen_store().get_workspace,
-        hydrate=hydrate,
+    workspace = zen_store().get_workspace(
+        workspace_name_or_id, hydrate=hydrate
     )
+    return dehydrate_response_model(workspace)
 
 
 @router.put(
@@ -231,12 +225,11 @@ def update_workspace(
     Returns:
         The updated workspace.
     """
-    return verify_permissions_and_update_entity(
-        id=workspace_name_or_id,
-        update_model=workspace_update,
-        get_method=zen_store().get_workspace,
-        update_method=zen_store().update_workspace,
+    workspace = zen_store().get_workspace(workspace_name_or_id, hydrate=False)
+    updated_workspace = zen_store().update_workspace(
+        workspace_id=workspace.id, workspace_update=workspace_update
     )
+    return dehydrate_response_model(updated_workspace)
 
 
 @router.delete(
@@ -253,11 +246,7 @@ def delete_workspace(
     Args:
         workspace_name_or_id: Name or ID of the workspace.
     """
-    verify_permissions_and_delete_entity(
-        id=workspace_name_or_id,
-        get_method=zen_store().get_workspace,
-        delete_method=zen_store().delete_workspace,
-    )
+    zen_store().delete_workspace(workspace_name_or_id)
 
 
 @router.get(
@@ -951,20 +940,21 @@ def get_or_create_pipeline_run(
             "is not supported."
         )
 
-    verify_permission(
-        resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
-    )
+    def _pre_creation_hook() -> None:
+        verify_permission(
+            resource_type=ResourceType.PIPELINE_RUN, action=Action.CREATE
+        )
+        check_entitlement(resource_type=ResourceType.PIPELINE_RUN)
 
     run, created = zen_store().get_or_create_run(
-        pipeline_run=pipeline_run,
-        pre_creation_hook=lambda: check_entitlement(
-            resource_type=ResourceType.PIPELINE_RUN
-        ),
+        pipeline_run=pipeline_run, pre_creation_hook=_pre_creation_hook
     )
     if created:
         report_usage(
             resource_type=ResourceType.PIPELINE_RUN, resource_id=run.id
         )
+    else:
+        verify_permission_for_model(run, action=Action.READ)
 
     return run, created
 
@@ -1009,24 +999,23 @@ def create_run_metadata(
             "is not supported."
         )
 
-    if run_metadata.resource_type == MetadataResourceTypes.PIPELINE_RUN:
-        run = zen_store().get_run(run_metadata.resource_id)
-        verify_permission_for_model(run, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.STEP_RUN:
-        step = zen_store().get_run_step(run_metadata.resource_id)
-        verify_permission_for_model(step, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.ARTIFACT_VERSION:
-        artifact_version = zen_store().get_artifact_version(
-            run_metadata.resource_id
-        )
-        verify_permission_for_model(artifact_version, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.MODEL_VERSION:
-        model_version = zen_store().get_model_version(run_metadata.resource_id)
-        verify_permission_for_model(model_version, action=Action.UPDATE)
-    else:
-        raise RuntimeError(
-            f"Unknown resource type: {run_metadata.resource_type}"
-        )
+    verify_models: List[Any] = []
+    for resource in run_metadata.resources:
+        if resource.type == MetadataResourceTypes.PIPELINE_RUN:
+            verify_models.append(zen_store().get_run(resource.id))
+        elif resource.type == MetadataResourceTypes.STEP_RUN:
+            verify_models.append(zen_store().get_run_step(resource.id))
+        elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
+            verify_models.append(zen_store().get_artifact_version(resource.id))
+        elif resource.type == MetadataResourceTypes.MODEL_VERSION:
+            verify_models.append(zen_store().get_model_version(resource.id))
+        else:
+            raise RuntimeError(f"Unknown resource type: {resource.type}")
+
+    batch_verify_permissions_for_models(
+        models=verify_models,
+        action=Action.UPDATE,
+    )
 
     verify_permission(
         resource_type=ResourceType.RUN_METADATA, action=Action.CREATE

zenml/zen_server/template_execution/utils.py

@@ -122,13 +122,15 @@ def run_template(
     placeholder_run = create_placeholder_run(deployment=new_deployment)
     assert placeholder_run
 
-    # We create an API token scoped to the pipeline run
+    # We create an API token scoped to the pipeline run that never expires
     api_token = generate_access_token(
         user_id=auth_context.user.id,
         pipeline_run_id=placeholder_run.id,
         # Keep the original API key or device scopes, if any
         api_key=auth_context.api_key,
         device=auth_context.device,
+        # Never expire the token
+        expires_in=0,
     ).access_token
 
     environment = {
@@ -363,6 +365,7 @@ def deployment_request_from_template(
                     "external_input_artifacts",
                     "model_artifacts_or_metadata",
                     "client_lazy_loaders",
+                    "substitutions",
                     "outputs",
                 }
             ),

zenml/zen_server/utils.py

@@ -421,6 +421,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
     """
     from fastapi import Query
 
+    from zenml.zen_server.exceptions import error_detail
+
     def init_cls_and_handle_errors(*args: Any, **kwargs: Any) -> BaseModel:
         from fastapi import HTTPException
 
@@ -428,9 +430,8 @@ def make_dependable(cls: Type[BaseModel]) -> Callable[..., Any]:
             inspect.signature(init_cls_and_handle_errors).bind(*args, **kwargs)
             return cls(*args, **kwargs)
         except ValidationError as e:
-            for error in e.errors():
-                error["loc"] = tuple(["query"] + list(error["loc"]))
-            raise HTTPException(422, detail=e.errors())
+            detail = error_detail(e, exception_type=ValueError)
+            raise HTTPException(422, detail=detail)
 
     params = {v.name: v for v in inspect.signature(cls).parameters.values()}
     query_params = getattr(cls, "API_MULTI_INPUT_PARAMS", [])

zenml/zen_stores/base_zen_store.py

@@ -36,6 +36,7 @@ from zenml.constants import (
     DEFAULT_STACK_AND_COMPONENT_NAME,
     DEFAULT_WORKSPACE_NAME,
     ENV_ZENML_DEFAULT_WORKSPACE_NAME,
+    ENV_ZENML_SERVER,
     IS_DEBUG_ENV,
 )
 from zenml.enums import (
@@ -155,9 +156,16 @@ class BaseZenStore(
             TypeError: If the store type is unsupported.
         """
         if store_type == StoreType.SQL:
-            from zenml.zen_stores.sql_zen_store import SqlZenStore
+            if os.environ.get(ENV_ZENML_SERVER):
+                from zenml.zen_server.rbac.rbac_sql_zen_store import (
+                    RBACSqlZenStore,
+                )
+
+                return RBACSqlZenStore
+            else:
+                from zenml.zen_stores.sql_zen_store import SqlZenStore
 
-            return SqlZenStore
+                return SqlZenStore
         elif store_type == StoreType.REST:
             from zenml.zen_stores.rest_zen_store import RestZenStore
 

zenml/zen_stores/migrations/versions/0.71.0_release.py

@@ -0,0 +1,23 @@
+"""Release [0.71.0].
+
+Revision ID: 0.71.0
+Revises: cc269488e5a9
+Create Date: 2024-12-04 20:44:19.316139
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = "0.71.0"
+down_revision = "cc269488e5a9"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    pass
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    pass

zenml/zen_stores/migrations/versions/26351d482b9e_add_step_run_unique_constraint.py

@@ -0,0 +1,37 @@
+"""Add step run unique constraint [26351d482b9e].
+
+Revision ID: 26351d482b9e
+Revises: 0.71.0
+Create Date: 2024-12-03 11:46:57.541578
+
+"""
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "26351d482b9e"
+down_revision = "0.71.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("step_run", schema=None) as batch_op:
+        batch_op.create_unique_constraint(
+            "unique_step_name_for_pipeline_run", ["name", "pipeline_run_id"]
+        )
+
+    # ### 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("step_run", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_step_name_for_pipeline_run", type_="unique"
+        )
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/a1237ba94fd8_add_model_version_producer_run_unique_.py

@@ -0,0 +1,68 @@
+"""Add model version producer run unique constraint [a1237ba94fd8].
+
+Revision ID: a1237ba94fd8
+Revises: 26351d482b9e
+Create Date: 2024-12-13 10:28:55.432414
+
+"""
+
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a1237ba94fd8"
+down_revision = "26351d482b9e"
+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("model_version", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "producer_run_id_if_numeric",
+                sqlmodel.sql.sqltypes.GUID(),
+                nullable=True,
+            )
+        )
+
+    # Set the producer_run_id_if_numeric column to the model version ID for
+    # existing rows
+    connection = op.get_bind()
+    metadata = sa.MetaData()
+    metadata.reflect(only=("model_version",), bind=connection)
+    model_version_table = sa.Table("model_version", metadata)
+
+    connection.execute(
+        model_version_table.update().values(
+            producer_run_id_if_numeric=model_version_table.c.id
+        )
+    )
+
+    with op.batch_alter_table("model_version", schema=None) as batch_op:
+        batch_op.alter_column(
+            "producer_run_id_if_numeric",
+            existing_type=sqlmodel.sql.sqltypes.GUID(),
+            nullable=False,
+        )
+        batch_op.create_unique_constraint(
+            "unique_numeric_version_for_pipeline_run",
+            ["model_id", "producer_run_id_if_numeric"],
+        )
+
+    # ### 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("model_version", schema=None) as batch_op:
+        batch_op.drop_constraint(
+            "unique_numeric_version_for_pipeline_run", type_="unique"
+        )
+        batch_op.drop_column("producer_run_id_if_numeric")
+
+    # ### end Alembic commands ###

zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py

@@ -0,0 +1,36 @@
+"""Remove component spec path [b73bc71f1106].
+
+Revision ID: b73bc71f1106
+Revises: ec6307720f92
+Create Date: 2024-11-29 09:36:33.089945
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b73bc71f1106"
+down_revision = "ec6307720f92"
+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("stack_component", schema=None) as batch_op:
+        batch_op.drop_column("component_spec_path")
+
+    # ### 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("stack_component", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("component_spec_path", sa.VARCHAR(), nullable=True)
+        )
+
+    # ### end Alembic commands ###