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

zenml/zen_server/jwt.py

@@ -25,7 +25,7 @@ from uuid import UUID
 import jwt
 from pydantic import BaseModel
 
-from zenml.exceptions import AuthorizationException
+from zenml.exceptions import CredentialsNotValid
 from zenml.logger import get_logger
 from zenml.zen_server.utils import server_config
 
@@ -40,16 +40,20 @@ class JWTToken(BaseModel):
         device_id: The id of the authenticated device.
         api_key_id: The id of the authenticated API key for which this token
             was issued.
-        pipeline_id: The id of the pipeline for which the token was issued.
         schedule_id: The id of the schedule for which the token was issued.
+        pipeline_run_id: The id of the pipeline run for which the token was
+            issued.
+        step_run_id: The id of the step run for which the token was
+            issued.
         claims: The original token claims.
     """
 
     user_id: UUID
     device_id: Optional[UUID] = None
     api_key_id: Optional[UUID] = None
-    pipeline_id: Optional[UUID] = None
     schedule_id: Optional[UUID] = None
+    pipeline_run_id: Optional[UUID] = None
+    step_run_id: Optional[UUID] = None
     claims: Dict[str, Any] = {}
 
     @classmethod
@@ -71,7 +75,7 @@ class JWTToken(BaseModel):
             The decoded JWT access token.
 
         Raises:
-            AuthorizationException: If the token is invalid.
+            CredentialsNotValid: If the token is invalid.
         """
         config = server_config()
 
@@ -87,27 +91,27 @@ class JWTToken(BaseModel):
             )
             claims = cast(Dict[str, Any], claims_data)
         except jwt.PyJWTError as e:
-            raise AuthorizationException(f"Invalid JWT token: {e}") from e
+            raise CredentialsNotValid(f"Invalid JWT token: {e}") from e
 
-        subject: str = claims.get("sub", "")
+        subject: str = claims.pop("sub", "")
         if not subject:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is missing"
             )
 
         try:
             user_id = UUID(subject)
         except ValueError:
-            raise AuthorizationException(
+            raise CredentialsNotValid(
                 "Invalid JWT token: the subject claim is not a valid UUID"
             )
 
         device_id: Optional[UUID] = None
         if "device_id" in claims:
             try:
-                device_id = UUID(claims["device_id"])
+                device_id = UUID(claims.pop("device_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the device_id claim is not a valid "
                     "UUID"
                 )
@@ -115,30 +119,40 @@ class JWTToken(BaseModel):
         api_key_id: Optional[UUID] = None
         if "api_key_id" in claims:
             try:
-                api_key_id = UUID(claims["api_key_id"])
+                api_key_id = UUID(claims.pop("api_key_id"))
             except ValueError:
-                raise AuthorizationException(
+                raise CredentialsNotValid(
                     "Invalid JWT token: the api_key_id claim is not a valid "
                     "UUID"
                 )
 
-        pipeline_id: Optional[UUID] = None
-        if "pipeline_id" in claims:
+        schedule_id: Optional[UUID] = None
+        if "schedule_id" in claims:
             try:
-                pipeline_id = UUID(claims["pipeline_id"])
+                schedule_id = UUID(claims.pop("schedule_id"))
             except ValueError:
-                raise AuthorizationException(
-                    "Invalid JWT token: the pipeline_id claim is not a valid "
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the schedule_id claim is not a valid "
                     "UUID"
                 )
 
-        schedule_id: Optional[UUID] = None
-        if "schedule_id" in claims:
+        pipeline_run_id: Optional[UUID] = None
+        if "pipeline_run_id" in claims:
             try:
-                schedule_id = UUID(claims["schedule_id"])
+                pipeline_run_id = UUID(claims.pop("pipeline_run_id"))
             except ValueError:
-                raise AuthorizationException(
-                    "Invalid JWT token: the schedule_id claim is not a valid "
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the pipeline_run_id claim is not a valid "
+                    "UUID"
+                )
+
+        step_run_id: Optional[UUID] = None
+        if "step_run_id" in claims:
+            try:
+                step_run_id = UUID(claims.pop("step_run_id"))
+            except ValueError:
+                raise CredentialsNotValid(
+                    "Invalid JWT token: the step_run_id claim is not a valid "
                     "UUID"
                 )
 
@@ -146,8 +160,9 @@ class JWTToken(BaseModel):
             user_id=user_id,
             device_id=device_id,
             api_key_id=api_key_id,
-            pipeline_id=pipeline_id,
             schedule_id=schedule_id,
+            pipeline_run_id=pipeline_run_id,
+            step_run_id=step_run_id,
             claims=claims,
         )
 
@@ -165,25 +180,27 @@ class JWTToken(BaseModel):
         """
         config = server_config()
 
-        claims: Dict[str, Any] = dict(
-            sub=str(self.user_id),
-        )
+        claims: Dict[str, Any] = self.claims.copy()
+
+        claims["sub"] = str(self.user_id)
         claims["iss"] = config.get_jwt_token_issuer()
         claims["aud"] = config.get_jwt_token_audience()
 
         if expires:
             claims["exp"] = expires
+        else:
+            claims.pop("exp", None)
+
         if self.device_id:
             claims["device_id"] = str(self.device_id)
         if self.api_key_id:
             claims["api_key_id"] = str(self.api_key_id)
-        if self.pipeline_id:
-            claims["pipeline_id"] = str(self.pipeline_id)
         if self.schedule_id:
             claims["schedule_id"] = str(self.schedule_id)
-
-        # Apply custom claims
-        claims.update(self.claims)
+        if self.pipeline_run_id:
+            claims["pipeline_run_id"] = str(self.pipeline_run_id)
+        if self.step_run_id:
+            claims["step_run_id"] = str(self.step_run_id)
 
         return jwt.encode(
             claims,

zenml/zen_server/rate_limit.py

@@ -0,0 +1,200 @@
+#  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.
+"""Rate limiting for the ZenML Server."""
+
+import inspect
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+from functools import wraps
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    TypeVar,
+    cast,
+)
+
+from starlette.requests import Request
+
+from zenml.logger import get_logger
+from zenml.zen_server.utils import server_config
+
+logger = get_logger(__name__)
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class RequestLimiter:
+    """Simple in-memory rate limiter."""
+
+    def __init__(
+        self,
+        day_limit: Optional[int] = None,
+        minute_limit: Optional[int] = None,
+    ):
+        """Initializes the limiter.
+
+        Args:
+            day_limit: The number of requests allowed per day.
+            minute_limit: The number of requests allowed per minute.
+
+        Raises:
+            ValueError: If both day_limit and minute_limit are None.
+        """
+        self.limiting_enabled = server_config().rate_limit_enabled
+        if not self.limiting_enabled:
+            return
+        if day_limit is None and minute_limit is None:
+            raise ValueError("Pass either day or minuter limits, or both.")
+        self.day_limit = day_limit
+        self.minute_limit = minute_limit
+        self.limiter: Dict[str, List[float]] = defaultdict(list)
+
+    def hit_limiter(self, request: Request) -> None:
+        """Increase the number of hits in the limiter.
+
+        Args:
+            request: Request object.
+
+        Raises:
+            HTTPException: If the request limit is exceeded.
+        """
+        if not self.limiting_enabled:
+            return
+        from fastapi import HTTPException
+
+        requester = self._get_ipaddr(request)
+        now = time.time()
+        minute_ago = now - 60
+        day_ago = now - 60 * 60 * 24
+        self.limiter[requester].append(now)
+
+        from bisect import bisect_left
+
+        # remove failures older than a day
+        older_index = bisect_left(self.limiter[requester], day_ago)
+        self.limiter[requester] = self.limiter[requester][older_index:]
+
+        if self.day_limit and len(self.limiter[requester]) > self.day_limit:
+            raise HTTPException(
+                status_code=429, detail="Daily request limit exceeded."
+            )
+        minute_requests = len(
+            [
+                limiter_hit
+                for limiter_hit in self.limiter[requester][::-1]
+                if limiter_hit >= minute_ago
+            ]
+        )
+        if self.minute_limit and minute_requests > self.minute_limit:
+            raise HTTPException(
+                status_code=429, detail="Minute request limit exceeded."
+            )
+
+    def reset_limiter(self, request: Request) -> None:
+        """Resets the limiter on successful request.
+
+        Args:
+            request: Request object.
+        """
+        if self.limiting_enabled:
+            requester = self._get_ipaddr(request)
+            if requester in self.limiter:
+                del self.limiter[requester]
+
+    def _get_ipaddr(self, request: Request) -> str:
+        """Returns the IP address for the current request.
+
+        Based on the X-Forwarded-For headers or client information.
+
+        Args:
+            request: The request object.
+
+        Returns:
+            The ip address for the current request (or 127.0.0.1 if none found).
+        """
+        if "X_FORWARDED_FOR" in request.headers:
+            return request.headers["X_FORWARDED_FOR"]
+        else:
+            if not request.client or not request.client.host:
+                return "127.0.0.1"
+
+            return request.client.host
+
+    @contextmanager
+    def limit_failed_requests(
+        self, request: Request
+    ) -> Generator[None, Any, Any]:
+        """Limits the number of failed requests.
+
+        Args:
+            request: Request object.
+
+        Yields:
+            None
+        """
+        self.hit_limiter(request)
+
+        yield
+
+        # if request was successful - reset limiter
+        self.reset_limiter(request)
+
+
+def rate_limit_requests(
+    day_limit: Optional[int] = None,
+    minute_limit: Optional[int] = None,
+) -> Callable[..., Any]:
+    """Decorator to handle exceptions in the API.
+
+    Args:
+        day_limit: Number of requests allowed per day.
+        minute_limit: Number of requests allowed per minute.
+
+    Returns:
+        Decorated function.
+    """
+    limiter = RequestLimiter(day_limit=day_limit, minute_limit=minute_limit)
+
+    def decorator(func: F) -> F:
+        request_arg, request_kwarg = None, None
+        parameters = inspect.signature(func).parameters
+        for arg_num, arg_name in enumerate(parameters):
+            if parameters[arg_name].annotation == Request:
+                request_arg = arg_num
+                request_kwarg = arg_name
+                break
+        if request_arg is None or request_kwarg is None:
+            raise ValueError(
+                "Rate limiting APIs must have argument of `Request` type."
+            )
+
+        @wraps(func)
+        def decorated(
+            *args: Any,
+            **kwargs: Any,
+        ) -> Any:
+            if request_kwarg in kwargs:
+                request = kwargs[request_kwarg]
+            else:
+                request = args[request_arg]
+            with limiter.limit_failed_requests(request):
+                return func(*args, **kwargs)
+
+        return cast(F, decorated)
+
+    return decorator

zenml/zen_server/rbac/endpoint_utils.py

@@ -1,18 +1,40 @@
+#  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.
 """High-level helper functions to write endpoints with RBAC."""
-from typing import Any, Callable, TypeVar, Union
+
+from typing import Any, Callable, List, TypeVar, Union
 from uuid import UUID
 
 from pydantic import BaseModel
 
+from zenml.constants import (
+    REPORTABLE_RESOURCES,
+    REQUIRES_CUSTOM_RESOURCE_REPORTING,
+)
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
     BaseFilter,
+    BaseIdentifiedResponse,
     BaseRequest,
-    BaseResponse,
     Page,
     UserScopedRequest,
 )
 from zenml.zen_server.auth import get_auth_context
+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 (
     dehydrate_page,
@@ -23,7 +45,7 @@ from zenml.zen_server.rbac.utils import (
 )
 
 AnyRequest = TypeVar("AnyRequest", bound=BaseRequest)
-AnyResponse = TypeVar("AnyResponse", bound=BaseResponse)  # type: ignore[type-arg]
+AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: ignore[type-arg]
 AnyFilter = TypeVar("AnyFilter", bound=BaseFilter)
 AnyUpdate = TypeVar("AnyUpdate", bound=BaseModel)
 UUIDOrStr = TypeVar("UUIDOrStr", UUID, Union[UUID, str])
@@ -57,9 +79,63 @@ def verify_permissions_and_create_entity(
                 f"Not allowed to create resource '{resource_type}' for a "
                 "different user."
             )
+    verify_permission(resource_type=resource_type, action=Action.CREATE)
+
+    needs_usage_increment = (
+        resource_type in REPORTABLE_RESOURCES
+        and resource_type not in REQUIRES_CUSTOM_RESOURCE_REPORTING
+    )
+    if needs_usage_increment:
+        check_entitlement(resource_type)
+
+    created = create_method(request_model)
+
+    if needs_usage_increment:
+        report_usage(resource_type, resource_id=created.id)
+
+    return created
+
+
+def verify_permissions_and_batch_create_entity(
+    batch: List[AnyRequest],
+    resource_type: ResourceType,
+    create_method: Callable[[List[AnyRequest]], List[AnyResponse]],
+) -> List[AnyResponse]:
+    """Verify permissions and create a batch of entities if authorized.
+
+    Args:
+        batch: The batch to create.
+        resource_type: The resource type of the entities to create.
+        create_method: The method to create the entities.
+
+    Raises:
+        IllegalOperationError: If the request model has a different owner then
+            the currently authenticated user.
+        RuntimeError: If the resource type is usage-tracked.
+
+    Returns:
+        The created entities.
+    """
+    auth_context = get_auth_context()
+    assert auth_context
+
+    for request_model in batch:
+        if isinstance(request_model, UserScopedRequest):
+            if request_model.user != auth_context.user.id:
+                raise IllegalOperationError(
+                    f"Not allowed to create resource '{resource_type}' for a "
+                    "different user."
+                )
 
     verify_permission(resource_type=resource_type, action=Action.CREATE)
-    return create_method(request_model)
+
+    if resource_type in REPORTABLE_RESOURCES:
+        raise RuntimeError(
+            "Batch requests are currently not possible with usage-tracked features."
+        )
+
+    created = create_method(batch)
+    return created
 
 
 def verify_permissions_and_get_entity(
@@ -113,7 +189,7 @@ def verify_permissions_and_list_entities(
 def verify_permissions_and_update_entity(
     id: UUIDOrStr,
     update_model: AnyUpdate,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     update_method: Callable[[UUIDOrStr, AnyUpdate], AnyResponse],
 ) -> AnyResponse:
     """Verify permissions and update an entity.
@@ -127,7 +203,8 @@ def verify_permissions_and_update_entity(
     Returns:
         A model of the updated entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.UPDATE)
     updated_model = update_method(model.id, update_model)
     return dehydrate_response_model(updated_model)
@@ -135,20 +212,26 @@ def verify_permissions_and_update_entity(
 
 def verify_permissions_and_delete_entity(
     id: UUIDOrStr,
-    get_method: Callable[[UUIDOrStr], AnyResponse],
+    get_method: Callable[[UUIDOrStr, bool], AnyResponse],
     delete_method: Callable[[UUIDOrStr], None],
-) -> None:
+) -> AnyResponse:
     """Verify permissions and delete an entity.
 
     Args:
         id: The ID of the entity to delete.
         get_method: The method to fetch the entity.
         delete_method: The method to delete the entity.
+
+    Returns:
+        The deleted entity.
     """
-    model = get_method(id)
+    # We don't need the hydrated version here
+    model = get_method(id, False)
     verify_permission_for_model(model, action=Action.DELETE)
     delete_method(model.id)
 
+    return model
+
 
 def verify_permissions_and_prune_entities(
     resource_type: ResourceType,

zenml/zen_server/rbac/models.py

@@ -16,7 +16,7 @@
 from typing import Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.utils.enum_utils import StrEnum
 
@@ -46,25 +46,32 @@ class Action(StrEnum):
 class ResourceType(StrEnum):
     """Resource types of the server API."""
 
-    STACK = "stack"
-    FLAVOR = "flavor"
-    STACK_COMPONENT = "stack_component"
-    PIPELINE = "pipeline"
+    ACTION = "action"
+    ARTIFACT = "artifact"
+    ARTIFACT_VERSION = "artifact_version"
     CODE_REPOSITORY = "code_repository"
+    EVENT_SOURCE = "event_source"
+    FLAVOR = "flavor"
     MODEL = "model"
     MODEL_VERSION = "model_version"
-    SERVICE_CONNECTOR = "service_connector"
-    ARTIFACT = "artifact"
-    ARTIFACT_VERSION = "artifact_version"
-    SECRET = "secret"
-    TAG = "tag"
-    SERVICE_ACCOUNT = "service_account"
-    WORKSPACE = "workspace"
+    PIPELINE = "pipeline"
     PIPELINE_RUN = "pipeline_run"
     PIPELINE_DEPLOYMENT = "pipeline_deployment"
     PIPELINE_BUILD = "pipeline_build"
+    RUN_TEMPLATE = "run_template"
+    SERVICE = "service"
     RUN_METADATA = "run_metadata"
-    USER = "user"
+    SECRET = "secret"
+    SERVICE_ACCOUNT = "service_account"
+    SERVICE_CONNECTOR = "service_connector"
+    STACK = "stack"
+    STACK_COMPONENT = "stack_component"
+    TAG = "tag"
+    TRIGGER = "trigger"
+    TRIGGER_EXECUTION = "trigger_execution"
+    # Deactivated for now
+    # USER = "user"
+    # WORKSPACE = "workspace"
 
 
 class Resource(BaseModel):
@@ -85,7 +92,4 @@ class Resource(BaseModel):
 
         return representation
 
-    class Config:
-        """Pydantic configuration class."""
-
-        frozen = True
+    model_config = ConfigDict(frozen=True)