zenml-nightly 0.68.1.dev20241103__py3-none-any.whl → 0.70.0.dev20241115__py3-none-any.whl

zenml/zen_server/auth.py

@@ -14,13 +14,13 @@
 """Authentication module for ZenML server."""
 
 from contextvars import ContextVar
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Callable, Optional, Union
 from urllib.parse import urlencode
 from uuid import UUID
 
 import requests
-from fastapi import Depends
+from fastapi import Depends, Response
 from fastapi.security import (
     HTTPBasic,
     HTTPBasicCredentials,
@@ -37,7 +37,7 @@ from zenml.constants import (
     LOGIN,
     VERSION_1,
 )
-from zenml.enums import AuthScheme, OAuthDeviceStatus
+from zenml.enums import AuthScheme, ExecutionStatus, OAuthDeviceStatus
 from zenml.exceptions import (
     AuthorizationException,
     CredentialsNotValid,
@@ -51,11 +51,13 @@ from zenml.models import (
     ExternalUserModel,
     OAuthDeviceInternalResponse,
     OAuthDeviceInternalUpdate,
+    OAuthTokenResponse,
     UserAuthModel,
     UserRequest,
     UserResponse,
     UserUpdate,
 )
+from zenml.zen_server.cache import cache_result
 from zenml.zen_server.exceptions import http_exception_from_error
 from zenml.zen_server.jwt import JWTToken
 from zenml.zen_server.utils import server_config, zen_store
@@ -329,6 +331,148 @@ def authenticate_credentials(
                 ),
             )
 
+        if decoded_token.schedule_id:
+            # If the token contains a schedule ID, we need to check if the
+            # schedule still exists in the database. We use a cached version
+            # of the schedule active status to avoid unnecessary database
+            # queries.
+
+            @cache_result(expiry=30)
+            def get_schedule_active(schedule_id: UUID) -> Optional[bool]:
+                """Get the active status of a schedule.
+
+                Args:
+                    schedule_id: The schedule ID.
+
+                Returns:
+                    The schedule active status or None if the schedule does not
+                    exist.
+                """
+                try:
+                    schedule = zen_store().get_schedule(
+                        schedule_id, hydrate=False
+                    )
+                except KeyError:
+                    return False
+
+                return schedule.active
+
+            schedule_active = get_schedule_active(decoded_token.schedule_id)
+            if schedule_active is None:
+                error = (
+                    f"Authentication error: error retrieving token schedule "
+                    f"{decoded_token.schedule_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if not schedule_active:
+                error = (
+                    f"Authentication error: schedule {decoded_token.schedule_id} "
+                    "is not active"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.pipeline_run_id:
+            # If the token contains a pipeline run ID, we need to check if the
+            # pipeline run exists in the database and the pipeline run has
+            # not concluded. We use a cached version of the pipeline run status
+            # to avoid unnecessary database queries.
+
+            @cache_result(expiry=30)
+            def get_pipeline_run_status(
+                pipeline_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a pipeline run.
+
+                Args:
+                    pipeline_run_id: The pipeline run ID.
+
+                Returns:
+                    The pipeline run status or None if the pipeline run does not
+                    exist.
+                """
+                try:
+                    pipeline_run = zen_store().get_run(
+                        pipeline_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return pipeline_run.status
+
+            pipeline_run_status = get_pipeline_run_status(
+                decoded_token.pipeline_run_id
+            )
+            if pipeline_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token pipeline run "
+                    f"{decoded_token.pipeline_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if pipeline_run_status in [
+                ExecutionStatus.FAILED,
+                ExecutionStatus.COMPLETED,
+            ]:
+                error = (
+                    f"The execution of pipeline run "
+                    f"{decoded_token.pipeline_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+        if decoded_token.step_run_id:
+            # If the token contains a step run ID, we need to check if the
+            # step run exists in the database and the step run has not concluded.
+            # We use a cached version of the step run status to avoid unnecessary
+            # database queries.
+
+            @cache_result(expiry=30)
+            def get_step_run_status(
+                step_run_id: UUID,
+            ) -> Optional[ExecutionStatus]:
+                """Get the status of a step run.
+
+                Args:
+                    step_run_id: The step run ID.
+
+                Returns:
+                    The step run status or None if the step run does not exist.
+                """
+                try:
+                    step_run = zen_store().get_run_step(
+                        step_run_id, hydrate=False
+                    )
+                except KeyError:
+                    return None
+
+                return step_run.status
+
+            step_run_status = get_step_run_status(decoded_token.step_run_id)
+            if step_run_status is None:
+                error = (
+                    f"Authentication error: error retrieving token step run "
+                    f"{decoded_token.step_run_id}"
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
+            if step_run_status in [
+                ExecutionStatus.FAILED,
+                ExecutionStatus.COMPLETED,
+            ]:
+                error = (
+                    f"The execution of step run "
+                    f"{decoded_token.step_run_id} has already concluded and "
+                    "API tokens scoped to it are no longer valid."
+                )
+                logger.error(error)
+                raise CredentialsNotValid(error)
+
         auth_context = AuthContext(
             user=user_model,
             access_token=decoded_token,
@@ -660,6 +804,80 @@ def authenticate_api_key(
     return AuthContext(user=user_model, api_key=internal_api_key)
 
 
+def generate_access_token(
+    user_id: UUID,
+    response: Optional[Response] = None,
+    device: Optional[OAuthDeviceInternalResponse] = None,
+    api_key: Optional[APIKeyInternalResponse] = None,
+    expires_in: Optional[int] = None,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
+) -> OAuthTokenResponse:
+    """Generates an access token for the given user.
+
+    Args:
+        user_id: The ID of the user.
+        response: The FastAPI response object.
+        device: The device used for authentication.
+        api_key: The service account API key used for authentication.
+        expires_in: The number of seconds until the token expires.
+        schedule_id: The ID of the schedule to scope the token to.
+        pipeline_run_id: The ID of the pipeline run to scope the token to.
+        step_run_id: The ID of the step run to scope the token to.
+
+    Returns:
+        An authentication response with an access token.
+    """
+    config = server_config()
+
+    # If the expiration time is not supplied, the JWT tokens are set to expire
+    # according to the values configured in the server config. Device tokens are
+    # handled separately from regular user tokens.
+    expires: Optional[datetime] = None
+    if expires_in:
+        expires = datetime.utcnow() + timedelta(seconds=expires_in)
+    elif device:
+        # If a device was used for authentication, the token will expire
+        # at the same time as the device.
+        expires = device.expires
+        if expires:
+            expires_in = max(
+                int(expires.timestamp() - datetime.utcnow().timestamp()), 0
+            )
+    elif config.jwt_token_expire_minutes:
+        expires = datetime.utcnow() + timedelta(
+            minutes=config.jwt_token_expire_minutes
+        )
+        expires_in = config.jwt_token_expire_minutes * 60
+
+    access_token = JWTToken(
+        user_id=user_id,
+        device_id=device.id if device else None,
+        api_key_id=api_key.id if api_key else None,
+        schedule_id=schedule_id,
+        pipeline_run_id=pipeline_run_id,
+        step_run_id=step_run_id,
+    ).encode(expires=expires)
+
+    if not device and response:
+        # Also set the access token as an HTTP only cookie in the response
+        response.set_cookie(
+            key=config.get_auth_cookie_name(),
+            value=access_token,
+            httponly=True,
+            samesite="lax",
+            max_age=config.jwt_token_expire_minutes * 60
+            if config.jwt_token_expire_minutes
+            else None,
+            domain=config.auth_cookie_domain,
+        )
+
+    return OAuthTokenResponse(
+        access_token=access_token, expires_in=expires_in, token_type="bearer"
+    )
+
+
 def http_authentication(
     credentials: HTTPBasicCredentials = Depends(HTTPBasic()),
 ) -> AuthContext:

zenml/zen_server/cache.py

@@ -0,0 +1,208 @@
+#  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.
+"""Memory cache module for the ZenML server."""
+
+import time
+from collections import OrderedDict
+from threading import Lock
+from typing import Any, Callable, Optional
+from typing import OrderedDict as OrderedDictType
+from uuid import UUID
+
+from zenml.logger import get_logger
+from zenml.utils.singleton import SingletonMetaClass
+
+logger = get_logger(__name__)
+
+
+class MemoryCacheEntry:
+    """Simple class to hold cache entry data."""
+
+    def __init__(self, value: Any, expiry: int) -> None:
+        """Initialize a cache entry with value and expiry time.
+
+        Args:
+            value: The value to store in the cache.
+            expiry: The expiry time in seconds.
+        """
+        self.value: Any = value
+        self.expiry: int = expiry
+        self.timestamp: float = time.time()
+
+    @property
+    def expired(self) -> bool:
+        """Check if the cache entry has expired.
+
+        Returns:
+            True if the cache entry has expired; otherwise, False.
+        """
+        return time.time() - self.timestamp >= self.expiry
+
+
+class MemoryCache(metaclass=SingletonMetaClass):
+    """Simple in-memory cache with expiry and capacity management.
+
+    This cache is thread-safe and can be used in both synchronous and
+    asynchronous contexts. It uses a simple LRU (Least Recently Used) eviction
+    strategy to manage the cache size.
+
+    Each cache entry has a key, value, timestamp, and expiry. The cache
+    automatically removes expired entries and evicts the oldest entry when
+    the cache reaches its maximum capacity.
+
+
+    Usage Example:
+
+        cache = MemoryCache()
+        uuid_key = UUID("12345678123456781234567812345678")
+
+        if not cache.get(uuid_key):
+            # Get the value from the database or other source
+            value = get_value_from_database()
+            cache.set(uuid_key, value, expiry=60)
+
+    Usage Example with decorator:
+
+        @cache_result(expiry=60)
+        def get_cached_value(key: UUID) -> Any:
+            return get_value_from_database(key)
+
+        uuid_key = UUID("12345678123456781234567812345678")
+
+        value = get_cached_value(uuid_key)
+    """
+
+    def __init__(self, max_capacity: int, default_expiry: int) -> None:
+        """Initialize the cache with a maximum capacity and default expiry time.
+
+        Args:
+            max_capacity: The maximum number of entries the cache can hold.
+            default_expiry: The default expiry time in seconds.
+        """
+        self.cache: OrderedDictType[UUID, MemoryCacheEntry] = OrderedDict()
+        self.max_capacity = max_capacity
+        self.default_expiry = default_expiry
+        self._lock = Lock()
+
+    def set(self, key: UUID, value: Any, expiry: Optional[int] = None) -> None:
+        """Insert value into cache with optional custom expiry time in seconds.
+
+        Args:
+            key: The key to insert the value with.
+            value: The value to insert into the cache.
+            expiry: The expiry time in seconds. If None, uses the default expiry.
+        """
+        with self._lock:
+            self.cache[key] = MemoryCacheEntry(
+                value=value, expiry=expiry or self.default_expiry
+            )
+            self._cleanup()
+
+    def get(self, key: UUID) -> Optional[Any]:
+        """Retrieve value if it's still valid; otherwise, return None.
+
+        Args:
+            key: The key to retrieve the value for.
+
+        Returns:
+            The value if it's still valid; otherwise, None.
+        """
+        with self._lock:
+            return self._get_internal(key)
+
+    def _get_internal(self, key: UUID) -> Optional[Any]:
+        """Helper to retrieve a value without lock (internal use only).
+
+        Args:
+            key: The key to retrieve the value for.
+
+        Returns:
+            The value if it's still valid; otherwise, None.
+        """
+        entry = self.cache.get(key)
+        if entry and not entry.expired:
+            return entry.value
+        elif entry:
+            del self.cache[key]  # Invalidate expired entry
+        return None
+
+    def _cleanup(self) -> None:
+        """Remove expired or excess entries."""
+        # Remove expired entries
+        keys_to_remove = [k for k, v in self.cache.items() if v.expired]
+        for k in keys_to_remove:
+            del self.cache[k]
+
+        # Ensure we don't exceed max capacity
+        while len(self.cache) > self.max_capacity:
+            self.cache.popitem(last=False)
+
+
+F = Callable[[UUID], Any]
+
+
+def cache_result(
+    expiry: Optional[int] = None,
+) -> Callable[[F], F]:
+    """A decorator to cache the result of a function based on a UUID key argument.
+
+    Args:
+        expiry: Custom time in seconds for the cache entry to expire. If None,
+            uses the default expiry time.
+
+    Returns:
+        A decorator that wraps a function, caching its results based on a UUID
+        key.
+    """
+
+    def decorator(func: F) -> F:
+        """The actual decorator that wraps the function with caching logic.
+
+        Args:
+            func: The function to wrap.
+
+        Returns:
+            The wrapped function with caching logic.
+        """
+
+        def wrapper(key: UUID) -> Any:
+            """The wrapped function with caching logic.
+
+            Args:
+                key: The key to use for caching.
+
+            Returns:
+                The result of the original function, either from cache or
+                freshly computed.
+            """
+            from zenml.zen_server.utils import memcache
+
+            cache = memcache()
+
+            # Attempt to retrieve the result from cache
+            cached_value = cache.get(key)
+            if cached_value is not None:
+                logger.debug(
+                    f"Memory cache hit for key: {key} and func: {func.__name__}"
+                )
+                return cached_value
+
+            # Call the original function and cache its result
+            result = func(key)
+            cache.set(key, result, expiry)
+            return result
+
+        return wrapper
+
+    return decorator

zenml/zen_server/dashboard/assets/{404-DT4QRUqN.js → 404-NVXKFp-x.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-DP6vWzyx.js";import{f as s,r as t}from"./index-QQb7wQEC.js";import{E as r}from"./EmptyState-DeK7H4pr.js";import{S as a}from"./help-Dam461dC.js";import{L as o}from"./@react-router-BMhZulnd.js";import"./@tanstack-BUCbhJyH.js";import"./@reactflow-8U9qNlMR.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-DeK6qiuw.js";import{f as s,r as t}from"./index-CCOPpudF.js";import{E as r}from"./EmptyState-BzdlCwp3.js";import{S as a}from"./help-Cc9bBIJH.js";import{L as o}from"./@react-router-B3Z5rLr2.js";import"./@tanstack-DT5WLu9C.js";import"./@reactflow-CK0KJUen.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};