zenml-nightly 0.66.0.dev20240923__py3-none-any.whl → 0.66.0.dev20240927__py3-none-any.whl

zenml/utils/string_utils.py

@@ -15,13 +15,17 @@
 
 import base64
 import datetime
+import functools
 import random
 import string
+from typing import Any, Callable, Dict, TypeVar, cast
 
 from pydantic import BaseModel
 
 from zenml.constants import BANNED_NAME_CHARACTERS
 
+V = TypeVar("V", bound=Any)
+
 
 def get_human_readable_time(seconds: float) -> str:
     """Convert seconds into a human-readable string.
@@ -167,3 +171,49 @@ def format_name_template(
         datetime.datetime.now(datetime.timezone.utc).strftime("%H_%M_%S_%f"),
     )
     return name_template.format(**kwargs)
+
+
+def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:
+    """Recursively substitute strings in objects.
+
+    Args:
+        value: An object in which the strings should be recursively substituted.
+            This can be a pydantic model, dict, set, list, tuple or any
+            primitive type.
+        substitution_func: The function that does the actual string
+            substitution.
+
+    Returns:
+        The object with the substitution function applied to all string values.
+    """
+    substitute_ = functools.partial(
+        substitute_string, substitution_func=substitution_func
+    )
+
+    if isinstance(value, BaseModel):
+        model_values = {}
+
+        for k, v in value.__iter__():
+            new_value = substitute_(v)
+
+            if k not in value.model_fields_set and new_value == getattr(
+                value, k
+            ):
+                # This is a default value on the model and was not set
+                # explicitly. In this case, we don't include it in the model
+                # values to keep the `exclude_unset` behavior the same
+                continue
+
+            model_values[k] = new_value
+
+        return cast(V, type(value).model_validate(model_values))
+    elif isinstance(value, Dict):
+        return cast(
+            V, {substitute_(k): substitute_(v) for k, v in value.items()}
+        )
+    elif isinstance(value, (list, set, tuple)):
+        return cast(V, type(value)(substitute_(v) for v in value))
+    elif isinstance(value, str):
+        return cast(V, substitution_func(value))
+
+    return value

zenml/zen_server/cloud_utils.py

@@ -1,6 +1,7 @@
 """Utils concerning anything concerning the cloud control plane backend."""
 
 import os
+from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Optional
 
 import requests
@@ -19,11 +20,9 @@ class ZenMLCloudConfiguration(BaseModel):
     """ZenML Pro RBAC configuration."""
 
     api_url: str
-
     oauth2_client_id: str
     oauth2_client_secret: str
     oauth2_audience: str
-    auth0_domain: str
 
     @field_validator("api_url")
     @classmethod
@@ -68,6 +67,8 @@ class ZenMLCloudConnection:
         """Initialize the RBAC component."""
         self._config = ZenMLCloudConfiguration.from_environment()
         self._session: Optional[requests.Session] = None
+        self._token: Optional[str] = None
+        self._token_expires_at: Optional[datetime] = None
 
     def get(
         self, endpoint: str, params: Optional[Dict[str, Any]]
@@ -91,7 +92,8 @@ class ZenMLCloudConnection:
 
         response = self.session.get(url=url, params=params, timeout=7)
         if response.status_code == 401:
-            # Refresh the auth token and try again
+            # If we get an Unauthorized error from the API serer, we refresh the
+            # auth token and try again
             self._clear_session()
             response = self.session.get(url=url, params=params, timeout=7)
 
@@ -186,6 +188,8 @@ class ZenMLCloudConnection:
     def _clear_session(self) -> None:
         """Clear the authentication session."""
         self._session = None
+        self._token = None
+        self._token_expires_at = None
 
     def _fetch_auth_token(self) -> str:
         """Fetch an auth token for the Cloud API from auth0.
@@ -196,8 +200,16 @@ class ZenMLCloudConnection:
         Returns:
             Auth token.
         """
+        if (
+            self._token is not None
+            and self._token_expires_at is not None
+            and datetime.now(timezone.utc) + timedelta(minutes=5)
+            < self._token_expires_at
+        ):
+            return self._token
+
         # Get an auth token from auth0
-        auth0_url = f"https://{self._config.auth0_domain}/oauth/token"
+        login_url = f"{self._config.api_url}/auth/login"
         headers = {"content-type": "application/x-www-form-urlencoded"}
         payload = {
             "client_id": self._config.oauth2_client_id,
@@ -207,18 +219,31 @@ class ZenMLCloudConnection:
         }
         try:
             response = requests.post(
-                auth0_url, headers=headers, data=payload, timeout=7
+                login_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", "")
+        json_response = response.json()
+        access_token = json_response.get("access_token", "")
+        expires_in = json_response.get("expires_in", 0)
 
-        if not access_token or not isinstance(access_token, str):
+        if (
+            not access_token
+            or not isinstance(access_token, str)
+            or not expires_in
+            or not isinstance(expires_in, int)
+        ):
             raise RuntimeError("Could not fetch auth token from auth0.")
 
-        return str(access_token)
+        self._token = access_token
+        self._token_expires_at = datetime.now(timezone.utc) + timedelta(
+            seconds=expires_in
+        )
+
+        assert self._token is not None
+        return self._token
 
 
 def cloud_connection() -> ZenMLCloudConnection:

zenml/zen_server/dashboard/assets/{404-iO8vpun1.js → 404-Y50hSt65.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-DnFH_oo1.js";import{f as s,r as t}from"./index-B9wVwe7u.js";import{E as r}from"./EmptyState-Cs3DEmso.js";import{H as a}from"./help-CwN931fX.js";import{L as o}from"./@react-router-APVeuk-U.js";import"./@tanstack-QbMbTrh5.js";import"./@reactflow-B6kq9fJZ.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-DnFH_oo1.js";import{f as s,r as t}from"./index-Ci0nJ8EZ.js";import{E as r}from"./EmptyState-Cs3DEmso.js";import{H as a}from"./help-CwN931fX.js";import{L as o}from"./@react-router-APVeuk-U.js";import"./@tanstack-QbMbTrh5.js";import"./@reactflow-ytavUpwh.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};