zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250115__py3-none-any.whl

zenml/zen_server/rbac/rbac_sql_zen_store.py

@@ -0,0 +1,175 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""RBAC SQL Zen Store implementation."""
+
+from typing import (
+    Optional,
+    Tuple,
+)
+from uuid import UUID
+
+from zenml.logger import get_logger
+from zenml.models import (
+    ModelRequest,
+    ModelResponse,
+    ModelVersionRequest,
+    ModelVersionResponse,
+)
+from zenml.zen_server.feature_gate.endpoint_utils import (
+    check_entitlement,
+    report_usage,
+)
+from zenml.zen_server.rbac.models import Action, ResourceType
+from zenml.zen_server.rbac.utils import (
+    verify_permission,
+    verify_permission_for_model,
+)
+from zenml.zen_stores.sql_zen_store import SqlZenStore
+
+logger = get_logger(__name__)
+
+
+class RBACSqlZenStore(SqlZenStore):
+    """Wrapper around the SQLZenStore that implements RBAC functionality."""
+
+    def _get_or_create_model(
+        self, model_request: ModelRequest
+    ) -> Tuple[bool, ModelResponse]:
+        """Get or create a model.
+
+        Args:
+            model_request: The model request.
+
+        # noqa: DAR401
+        Raises:
+            Exception: If the user is not allowed to create a model.
+
+        Returns:
+            A boolean whether the model was created or not, and the model.
+        """
+        allow_model_creation = True
+        error = None
+
+        try:
+            verify_permission(
+                resource_type=ResourceType.MODEL, action=Action.CREATE
+            )
+            check_entitlement(resource_type=ResourceType.MODEL)
+        except Exception as e:
+            allow_model_creation = False
+            error = e
+
+        if allow_model_creation:
+            created, model_response = super()._get_or_create_model(
+                model_request
+            )
+        else:
+            try:
+                model_response = self.get_model(model_request.name)
+                created = False
+            except KeyError:
+                # The model does not exist. We now raise the error that
+                # explains why the model could not be created, instead of just
+                # the KeyError that it doesn't exist
+                assert error
+                raise error from None
+
+        if created:
+            report_usage(
+                resource_type=ResourceType.MODEL, resource_id=model_response.id
+            )
+        else:
+            verify_permission_for_model(model_response, action=Action.READ)
+
+        return created, model_response
+
+    def _get_model_version(
+        self,
+        model_id: UUID,
+        version_name: Optional[str] = None,
+        producer_run_id: Optional[UUID] = None,
+    ) -> ModelVersionResponse:
+        """Get a model version.
+
+        Args:
+            model_id: The ID of the model.
+            version_name: The name of the model version.
+            producer_run_id: The ID of the producer pipeline run. If this is
+                set, only numeric versions created as part of the pipeline run
+                will be returned.
+
+        Returns:
+            The model version.
+        """
+        model_version = super()._get_model_version(
+            model_id=model_id,
+            version_name=version_name,
+            producer_run_id=producer_run_id,
+        )
+        verify_permission_for_model(model_version, action=Action.READ)
+        return model_version
+
+    def _get_or_create_model_version(
+        self,
+        model_version_request: ModelVersionRequest,
+        producer_run_id: Optional[UUID] = None,
+    ) -> Tuple[bool, ModelVersionResponse]:
+        """Get or create a model version.
+
+        Args:
+            model_version_request: The model version request.
+            producer_run_id: ID of the producer pipeline run.
+
+        # noqa: DAR401
+        Raises:
+            Exception: If the authenticated user is not allowed to
+                create a model version.
+
+        Returns:
+            A boolean whether the model version was created or not, and the
+            model version.
+        """
+        allow_creation = True
+        error = None
+
+        try:
+            verify_permission(
+                resource_type=ResourceType.MODEL_VERSION, action=Action.CREATE
+            )
+        except Exception as e:
+            allow_creation = False
+            error = e
+
+        if allow_creation:
+            created, model_version_response = (
+                super()._get_or_create_model_version(
+                    model_version_request, producer_run_id=producer_run_id
+                )
+            )
+        else:
+            try:
+                model_version_response = self._get_model_version(
+                    model_id=model_version_request.model,
+                    version_name=model_version_request.name,
+                    producer_run_id=producer_run_id,
+                )
+                created = False
+            except KeyError:
+                # The model version does not exist. We now raise the error that
+                # explains why the version could not be created, instead of just
+                # the KeyError that it doesn't exist
+                assert error
+                raise error from None
+
+        return created, model_version_response

zenml/zen_server/rbac/utils.py

@@ -29,7 +29,12 @@ from uuid import UUID
 from pydantic import BaseModel
 
 from zenml.exceptions import IllegalOperationError
-from zenml.models import BaseResponse, Page, UserResponse, UserScopedResponse
+from zenml.models import (
+    BaseIdentifiedResponse,
+    Page,
+    UserResponse,
+    UserScopedResponse,
+)
 from zenml.zen_server.auth import get_auth_context
 from zenml.zen_server.rbac.models import Action, Resource, ResourceType
 from zenml.zen_server.utils import rbac, server_config
@@ -37,7 +42,7 @@ from zenml.zen_server.utils import rbac, server_config
 if TYPE_CHECKING:
     from zenml.zen_stores.schemas import BaseSchema
 
-AnyResponse = TypeVar("AnyResponse", bound=BaseResponse)  # type: ignore[type-arg]
+AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
 AnyModel = TypeVar("AnyModel", bound=BaseModel)
 
 
@@ -67,7 +72,7 @@ def dehydrate_page(page: Page[AnyResponse]) -> Page[AnyResponse]:
         for item in page.items
     ]
 
-    return page.copy(update={"items": new_items})
+    return page.model_copy(update={"items": new_items})
 
 
 def dehydrate_response_model(
@@ -98,15 +103,14 @@ def dehydrate_response_model(
         )
 
     dehydrated_values = {}
-    for key, value in model.__dict__.items():
-        if key in model.__private_attributes__:
-            dehydrated_values[key] = value
-        else:
-            dehydrated_values[key] = _dehydrate_value(
-                value, permissions=permissions
-            )
+    # See `get_subresources_for_model(...)` for a detailed explanation why we
+    # need to use `model.__iter__()` here
+    for key, value in model.__iter__():
+        dehydrated_values[key] = _dehydrate_value(
+            value, permissions=permissions
+        )
 
-    return type(model).parse_obj(dehydrated_values)
+    return type(model).model_validate(dehydrated_values)
 
 
 def _dehydrate_value(
@@ -124,7 +128,7 @@ def _dehydrate_value(
     Returns:
         The recursively dehydrated value.
     """
-    if isinstance(value, BaseResponse):
+    if isinstance(value, BaseIdentifiedResponse):
         permission_model = get_surrogate_permission_model_for_model(
             value, action=Action.READ
         )
@@ -139,6 +143,8 @@ def _dehydrate_value(
             return dehydrate_response_model(value, permissions=permissions)
         else:
             return get_permission_denied_model(value)
+    elif isinstance(value, Page):
+        return dehydrate_page(page=value)
     elif isinstance(value, BaseModel):
         return dehydrate_response_model(value, permissions=permissions)
     elif isinstance(value, Dict):
@@ -184,8 +190,13 @@ def get_permission_denied_model(model: AnyResponse) -> AnyResponse:
     Returns:
         The permission denied model.
     """
-    return model.copy(
-        exclude={"body", "metadata"}, update={"permission_denied": True}
+    return model.model_copy(
+        update={
+            "body": None,
+            "metadata": None,
+            "resources": None,
+            "permission_denied": True,
+        }
     )
 
 
@@ -228,7 +239,10 @@ def verify_permission_for_model(model: AnyResponse, action: Action) -> None:
     batch_verify_permissions_for_models(models=[model], action=action)
 
 
-def batch_verify_permissions(resources: Set[Resource], action: Action) -> None:
+def batch_verify_permissions(
+    resources: Set[Resource],
+    action: Action,
+) -> None:
     """Batch permission verification.
 
     Args:
@@ -337,7 +351,7 @@ def get_resource_for_model(model: AnyResponse) -> Optional[Resource]:
 
 def get_surrogate_permission_model_for_model(
     model: AnyResponse, action: str
-) -> BaseResponse[Any, Any]:
+) -> BaseIdentifiedResponse[Any, Any, Any]:
     """Get a surrogate permission model for a model.
 
     In some cases a different model instead of the original model is used to
@@ -377,10 +391,12 @@ def get_resource_type_for_model(
         is not associated with any resource type.
     """
     from zenml.models import (
+        ActionResponse,
         ArtifactResponse,
         ArtifactVersionResponse,
         CodeRepositoryResponse,
         ComponentResponse,
+        EventSourceResponse,
         FlavorResponse,
         ModelResponse,
         ModelVersionResponse,
@@ -388,20 +404,23 @@ def get_resource_type_for_model(
         PipelineDeploymentResponse,
         PipelineResponse,
         PipelineRunResponse,
-        RunMetadataResponse,
+        RunTemplateResponse,
         SecretResponse,
         ServiceAccountResponse,
         ServiceConnectorResponse,
+        ServiceResponse,
         StackResponse,
         TagResponse,
-        UserResponse,
-        WorkspaceResponse,
+        TriggerExecutionResponse,
+        TriggerResponse,
     )
 
     mapping: Dict[
         Any,
         ResourceType,
     ] = {
+        ActionResponse: ResourceType.ACTION,
+        EventSourceResponse: ResourceType.EVENT_SOURCE,
         FlavorResponse: ResourceType.FLAVOR,
         ServiceConnectorResponse: ResourceType.SERVICE_CONNECTOR,
         ComponentResponse: ResourceType.STACK_COMPONENT,
@@ -413,14 +432,17 @@ def get_resource_type_for_model(
         ModelVersionResponse: ResourceType.MODEL_VERSION,
         ArtifactResponse: ResourceType.ARTIFACT,
         ArtifactVersionResponse: ResourceType.ARTIFACT_VERSION,
-        WorkspaceResponse: ResourceType.WORKSPACE,
-        UserResponse: ResourceType.USER,
-        RunMetadataResponse: ResourceType.RUN_METADATA,
+        # WorkspaceResponse: ResourceType.WORKSPACE,
+        # UserResponse: ResourceType.USER,
         PipelineDeploymentResponse: ResourceType.PIPELINE_DEPLOYMENT,
         PipelineBuildResponse: ResourceType.PIPELINE_BUILD,
         PipelineRunResponse: ResourceType.PIPELINE_RUN,
+        RunTemplateResponse: ResourceType.RUN_TEMPLATE,
         TagResponse: ResourceType.TAG,
+        TriggerResponse: ResourceType.TRIGGER,
+        TriggerExecutionResponse: ResourceType.TRIGGER_EXECUTION,
         ServiceAccountResponse: ResourceType.SERVICE_ACCOUNT,
+        ServiceResponse: ResourceType.SERVICE,
     }
 
     return mapping.get(type(model))
@@ -462,10 +484,18 @@ def get_subresources_for_model(
     """
     resources = set()
 
-    for key, value in model.__dict__.items():
-        if key in model.__private_attributes__:
-            continue
-        resources.update(_get_subresources_for_value(value))
+    # We don't want to use `model.model_dump()` here as that recursively
+    # converts models to dicts, but we want to preserve those classes for
+    # the recursive `_get_subresources_for_value` calls.
+    # We previously used `dict(model)` here, but that lead to issues with
+    # models overwriting `__getattr__`, this `model.__iter__()` has the same
+    # results though.
+    if isinstance(model, Page):
+        for item in model:
+            resources.update(_get_subresources_for_value(item))
+    else:
+        for _, value in model.__iter__():
+            resources.update(_get_subresources_for_value(value))
 
     return resources
 
@@ -479,7 +509,7 @@ def _get_subresources_for_value(value: Any) -> Set[Resource]:
     Returns:
         All resources of the value which need permission verification.
     """
-    if isinstance(value, BaseResponse):
+    if isinstance(value, BaseIdentifiedResponse):
         resources = set()
         if not is_owned_by_authenticated_user(value):
             value = get_surrogate_permission_model_for_model(
@@ -515,9 +545,11 @@ def get_schema_for_resource_type(
         The database schema.
     """
     from zenml.zen_stores.schemas import (
+        ActionSchema,
         ArtifactSchema,
         ArtifactVersionSchema,
         CodeRepositorySchema,
+        EventSourceSchema,
         FlavorSchema,
         ModelSchema,
         ModelVersionSchema,
@@ -526,13 +558,16 @@ def get_schema_for_resource_type(
         PipelineRunSchema,
         PipelineSchema,
         RunMetadataSchema,
+        RunTemplateSchema,
         SecretSchema,
         ServiceConnectorSchema,
+        ServiceSchema,
         StackComponentSchema,
         StackSchema,
         TagSchema,
+        TriggerExecutionSchema,
+        TriggerSchema,
         UserSchema,
-        WorkspaceSchema,
     )
 
     mapping: Dict[ResourceType, Type["BaseSchema"]] = {
@@ -547,14 +582,20 @@ def get_schema_for_resource_type(
         ResourceType.ARTIFACT: ArtifactSchema,
         ResourceType.ARTIFACT_VERSION: ArtifactVersionSchema,
         ResourceType.SECRET: SecretSchema,
+        ResourceType.SERVICE: ServiceSchema,
         ResourceType.TAG: TagSchema,
         ResourceType.SERVICE_ACCOUNT: UserSchema,
-        ResourceType.WORKSPACE: WorkspaceSchema,
+        # ResourceType.WORKSPACE: WorkspaceSchema,
         ResourceType.PIPELINE_RUN: PipelineRunSchema,
         ResourceType.PIPELINE_DEPLOYMENT: PipelineDeploymentSchema,
         ResourceType.PIPELINE_BUILD: PipelineBuildSchema,
+        ResourceType.RUN_TEMPLATE: RunTemplateSchema,
         ResourceType.RUN_METADATA: RunMetadataSchema,
-        ResourceType.USER: UserSchema,
+        # ResourceType.USER: UserSchema,
+        ResourceType.ACTION: ActionSchema,
+        ResourceType.EVENT_SOURCE: EventSourceSchema,
+        ResourceType.TRIGGER: TriggerSchema,
+        ResourceType.TRIGGER_EXECUTION: TriggerExecutionSchema,
     }
 
     return mapping[resource_type]

zenml/zen_server/rbac/zenml_cloud_rbac.py

@@ -12,13 +12,10 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Cloud RBAC implementation."""
-import os
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple
 
-import requests
-from pydantic import BaseModel, validator
-from requests.adapters import HTTPAdapter, Retry
+from typing import TYPE_CHECKING, Dict, List, Set, Tuple
 
+from zenml.zen_server.cloud_utils import cloud_connection
 from zenml.zen_server.rbac.models import Action, Resource
 from zenml.zen_server.rbac.rbac_interface import RBACInterface
 from zenml.zen_server.utils import server_config
@@ -27,7 +24,6 @@ if TYPE_CHECKING:
     from zenml.models import UserResponse
 
 
-ZENML_CLOUD_RBAC_ENV_PREFIX = "ZENML_CLOUD_"
 PERMISSIONS_ENDPOINT = "/rbac/check_permissions"
 ALLOWED_RESOURCE_IDS_ENDPOINT = "/rbac/allowed_resource_ids"
 RESOURCE_MEMBERSHIP_ENDPOINT = "/rbac/resource_members"
@@ -38,7 +34,7 @@ SERVER_ID = server_config().external_server_id
 
 
 def _convert_to_cloud_resource(resource: Resource) -> str:
-    """Convert a resource to a ZenML Cloud API resource.
+    """Convert a resource to a ZenML Pro Management Plane resource.
 
     Args:
         resource: The resource to convert.
@@ -78,59 +74,12 @@ def _convert_from_cloud_resource(cloud_resource: str) -> Resource:
         return Resource(type=resource_type_and_id)
 
 
-class ZenMLCloudRBACConfiguration(BaseModel):
-    """ZenML Cloud RBAC configuration."""
-
-    api_url: str
-
-    oauth2_client_id: str
-    oauth2_client_secret: str
-    oauth2_audience: str
-    auth0_domain: str
-
-    @validator("api_url")
-    def _strip_trailing_slashes_url(cls, url: str) -> str:
-        """Strip any trailing slashes on the API URL.
-
-        Args:
-            url: The API URL.
-
-        Returns:
-            The API URL with potential trailing slashes removed.
-        """
-        return url.rstrip("/")
-
-    @classmethod
-    def from_environment(cls) -> "ZenMLCloudRBACConfiguration":
-        """Get the RBAC configuration from environment variables.
-
-        Returns:
-            The RBAC configuration.
-        """
-        env_config: Dict[str, Any] = {}
-        for k, v in os.environ.items():
-            if v == "":
-                continue
-            if k.startswith(ZENML_CLOUD_RBAC_ENV_PREFIX):
-                env_config[k[len(ZENML_CLOUD_RBAC_ENV_PREFIX) :].lower()] = v
-
-        return ZenMLCloudRBACConfiguration(**env_config)
-
-    class Config:
-        """Pydantic configuration class."""
-
-        # Allow extra attributes from configs of previous ZenML versions to
-        # permit downgrading
-        extra = "allow"
-
-
 class ZenMLCloudRBAC(RBACInterface):
-    """RBAC implementation that uses the ZenML Cloud API as a backend."""
+    """RBAC implementation that uses the ZenML Pro Management Plane as a backend."""
 
     def __init__(self) -> None:
-        """Initialize the RBAC component."""
-        self._config = ZenMLCloudRBACConfiguration.from_environment()
-        self._session: Optional[requests.Session] = None
+        """Initialize the object."""
+        self._connection = cloud_connection()
 
     def check_permissions(
         self, user: "UserResponse", resources: Set[Resource], action: Action
@@ -154,7 +103,7 @@ class ZenMLCloudRBAC(RBACInterface):
             # Service accounts have full permissions for now
             return {resource: True for resource in resources}
 
-        # At this point it's a regular user, which in the ZenML cloud with RBAC
+        # At this point it's a regular user, which in a ZenML Pro with RBAC
         # enabled is always authenticated using external authentication
         assert user.external_user_id
 
@@ -165,7 +114,9 @@ class ZenMLCloudRBAC(RBACInterface):
             ],
             "action": str(action),
         }
-        response = self._get(endpoint=PERMISSIONS_ENDPOINT, params=params)
+        response = self._connection.get(
+            endpoint=PERMISSIONS_ENDPOINT, params=params
+        )
         value = response.json()
 
         assert isinstance(value, dict)
@@ -194,7 +145,7 @@ class ZenMLCloudRBAC(RBACInterface):
             # Service accounts have full permissions for now
             return True, []
 
-        # At this point it's a regular user, which in the ZenML cloud with RBAC
+        # At this point it's a regular user, which in the ZenML Pro with RBAC
         # enabled is always authenticated using external authentication
         assert user.external_user_id
         params = {
@@ -202,7 +153,7 @@ class ZenMLCloudRBAC(RBACInterface):
             "resource": _convert_to_cloud_resource(resource),
             "action": str(action),
         }
-        response = self._get(
+        response = self._connection.get(
             endpoint=ALLOWED_RESOURCE_IDS_ENDPOINT, params=params
         )
         response_json = response.json()
@@ -232,130 +183,4 @@ class ZenMLCloudRBAC(RBACInterface):
             "resource": _convert_to_cloud_resource(resource),
             "actions": [str(action) for action in actions],
         }
-        self._post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)
-
-    def _get(self, endpoint: str, params: Dict[str, Any]) -> requests.Response:
-        """Send a GET request using the active session.
-
-        Args:
-            endpoint: The endpoint to send the request to. This will be appended
-                to the base URL.
-            params: Parameters to include in the request.
-
-        Raises:
-            RuntimeError: If the request failed.
-
-        Returns:
-            The response.
-        """
-        url = self._config.api_url + endpoint
-
-        response = self.session.get(url=url, params=params, timeout=7)
-        if response.status_code == 401:
-            # Refresh the auth token and try again
-            self._clear_session()
-            response = self.session.get(url=url, params=params, timeout=7)
-
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            raise RuntimeError(
-                f"Failed while trying to contact RBAC service: {e}"
-            )
-
-        return response
-
-    def _post(
-        self,
-        endpoint: str,
-        params: Optional[Dict[str, Any]] = None,
-        data: Optional[Dict[str, Any]] = None,
-    ) -> requests.Response:
-        """Send a POST request using the active session.
-
-        Args:
-            endpoint: The endpoint to send the request to. This will be appended
-                to the base URL.
-            params: Parameters to include in the request.
-            data: Data to include in the request.
-
-        Raises:
-            RuntimeError: If the request failed.
-
-        Returns:
-            The response.
-        """
-        url = self._config.api_url + endpoint
-
-        response = self.session.post(
-            url=url, params=params, json=data, timeout=7
-        )
-        if response.status_code == 401:
-            # Refresh the auth token and try again
-            self._clear_session()
-            response = self.session.post(
-                url=url, params=params, json=data, timeout=7
-            )
-
-        try:
-            response.raise_for_status()
-        except requests.HTTPError as e:
-            raise RuntimeError(
-                f"Failed while trying to contact RBAC service: {e}"
-            )
-
-        return response
-
-    @property
-    def session(self) -> requests.Session:
-        """Authenticate to the ZenML Cloud API.
-
-        Returns:
-            A requests session with the authentication token.
-        """
-        if self._session is None:
-            self._session = requests.Session()
-            token = self._fetch_auth_token()
-            self._session.headers.update({"Authorization": "Bearer " + token})
-
-            retries = Retry(total=5, backoff_factor=0.1)
-            self._session.mount("https://", HTTPAdapter(max_retries=retries))
-
-        return self._session
-
-    def _clear_session(self) -> None:
-        """Clear the authentication session."""
-        self._session = None
-
-    def _fetch_auth_token(self) -> str:
-        """Fetch an auth token for the Cloud API from auth0.
-
-        Raises:
-            RuntimeError: If the auth token can't be fetched.
-
-        Returns:
-            Auth token.
-        """
-        # Get an auth token from auth0
-        auth0_url = f"https://{self._config.auth0_domain}/oauth/token"
-        headers = {"content-type": "application/x-www-form-urlencoded"}
-        payload = {
-            "client_id": self._config.oauth2_client_id,
-            "client_secret": self._config.oauth2_client_secret,
-            "audience": self._config.oauth2_audience,
-            "grant_type": "client_credentials",
-        }
-        try:
-            response = requests.post(
-                auth0_url, headers=headers, data=payload, timeout=7
-            )
-            response.raise_for_status()
-        except Exception as e:
-            raise RuntimeError(f"Error fetching auth token from auth0: {e}")
-
-        access_token = response.json().get("access_token", "")
-
-        if not access_token or not isinstance(access_token, str):
-            raise RuntimeError("Could not fetch auth token from auth0.")
-
-        return str(access_token)
+        self._connection.post(endpoint=RESOURCE_MEMBERSHIP_ENDPOINT, data=data)