zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250115__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, HTTPException, status
+from fastapi import Depends, Response
 from fastapi.security import (
     HTTPBasic,
     HTTPBasicCredentials,
@@ -37,8 +37,17 @@ from zenml.constants import (
     LOGIN,
     VERSION_1,
 )
-from zenml.enums import AuthScheme, OAuthDeviceStatus
-from zenml.exceptions import AuthorizationException, OAuthError
+from zenml.enums import (
+    AuthScheme,
+    ExecutionStatus,
+    OAuthDeviceStatus,
+    OnboardingStep,
+)
+from zenml.exceptions import (
+    AuthorizationException,
+    CredentialsNotValid,
+    OAuthError,
+)
 from zenml.logger import get_logger
 from zenml.models import (
     APIKey,
@@ -47,11 +56,14 @@ 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
 
@@ -109,7 +121,7 @@ def _fetch_and_verify_api_key(
         The fetched API key.
 
     Raises:
-        AuthorizationException: If the API key could not be found, is not
+        CredentialsNotValid: If the API key could not be found, is not
             active, if it could not be verified against the supplied key value
             or if the associated service account is not active.
     """
@@ -118,11 +130,9 @@ def _fetch_and_verify_api_key(
     try:
         api_key = zen_store().get_internal_api_key(api_key_id)
     except KeyError:
-        error = (
-            f"Authentication error: error retrieving API key " f"{api_key_id}"
-        )
+        error = f"Authentication error: error retrieving API key {api_key_id}"
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if not api_key.service_account.active:
         error = (
@@ -131,7 +141,7 @@ def _fetch_and_verify_api_key(
             f"associated with API key {api_key.name} is not active"
         )
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if not api_key.active:
         error = (
@@ -141,7 +151,7 @@ def _fetch_and_verify_api_key(
             f"{api_key.service_account.name} is not active"
         )
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     if key_to_verify and not api_key.verify_key(key_to_verify):
         error = (
@@ -149,12 +159,12 @@ def _fetch_and_verify_api_key(
             f"{api_key.name}"
         )
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     # Update the "last used" timestamp of the API key
     store.update_internal_api_key(
         api_key.id,
-        APIKeyInternalUpdate(update_last_login=True),  # type: ignore[call-arg]
+        APIKeyInternalUpdate(update_last_login=True),
     )
 
     return api_key
@@ -188,7 +198,7 @@ def authenticate_credentials(
         The authenticated account details.
 
     Raises:
-        AuthorizationException: If the credentials are invalid.
+        CredentialsNotValid: If the credentials are invalid.
     """
     user: Optional[UserAuthModel] = None
     auth_context: Optional[AuthContext] = None
@@ -218,11 +228,11 @@ def authenticate_credentials(
         if not UserAuthModel.verify_password(password, user):
             error = "Authentication error: invalid username or password"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
         if user and not user.active:
             error = f"Authentication error: user {user.name} is not active"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     elif activation_token is not None:
         if not UserAuthModel.verify_activation_token(activation_token, user):
@@ -231,17 +241,17 @@ def authenticate_credentials(
                 f"{user_name_or_id}"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     elif access_token is not None:
         try:
             decoded_token = JWTToken.decode_token(
                 token=access_token,
             )
-        except AuthorizationException:
-            error = "Authentication error: error decoding access token. You may need to rerun zenml connect."
+        except CredentialsNotValid as e:
+            error = f"Authentication error: error decoding access token: {e}."
             logger.exception(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
         try:
             user_model = zen_store().get_user(
@@ -253,7 +263,7 @@ def authenticate_credentials(
                 f"{decoded_token.user_id}"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
         if not user_model.active:
             error = (
@@ -261,7 +271,7 @@ def authenticate_credentials(
                 f"active"
             )
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
         api_key_model: Optional[APIKeyInternalResponse] = None
         if decoded_token.api_key_id:
@@ -285,7 +295,7 @@ def authenticate_credentials(
                     f"{decoded_token.device_id}"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if (
                 device_model.user is None
@@ -296,7 +306,7 @@ def authenticate_credentials(
                     f"does not belong to user {user_model.name}"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if device_model.status != OAuthDeviceStatus.ACTIVE:
                 error = (
@@ -304,7 +314,7 @@ def authenticate_credentials(
                     f"is not active"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             if (
                 device_model.expires
@@ -315,7 +325,7 @@ def authenticate_credentials(
                     "has expired"
                 )
                 logger.error(error)
-                raise AuthorizationException(error)
+                raise CredentialsNotValid(error)
 
             zen_store().update_internal_authorized_device(
                 device_id=device_model.id,
@@ -324,6 +334,142 @@ 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.is_finished:
+                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.is_finished:
+                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,
@@ -340,12 +486,12 @@ def authenticate_credentials(
         if server_config().auth_scheme != AuthScheme.NO_AUTH:
             error = "Authentication error: no credentials provided"
             logger.error(error)
-            raise AuthorizationException(error)
+            raise CredentialsNotValid(error)
 
     if not auth_context:
         error = "Authentication error: invalid credentials"
         logger.error(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     return auth_context
 
@@ -487,12 +633,15 @@ def authenticate_device(client_id: UUID, device_code: str) -> AuthContext:
     return AuthContext(user=device_model.user, device=device_model)
 
 
-def authenticate_external_user(external_access_token: str) -> AuthContext:
+def authenticate_external_user(
+    external_access_token: str, request: Request
+) -> AuthContext:
     """Implement external authentication.
 
     Args:
         external_access_token: The access token used to authenticate the user
             to the external authenticator.
+        request: The request object.
 
     Returns:
         The authentication context reflecting the authenticated user.
@@ -522,8 +671,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
         )
     except Exception as e:
         logger.exception(
-            f"Error fetching user information from external authenticator: "
-            f"{e}"
+            f"Error fetching user information from external authenticator: {e}"
         )
         raise AuthorizationException(
             "Error fetching user information from external authenticator."
@@ -544,7 +692,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
 
         if isinstance(payload, dict):
             try:
-                external_user = ExternalUserModel.parse_obj(payload)
+                external_user = ExternalUserModel.model_validate(payload)
             except Exception as e:
                 logger.exception(
                     f"Error parsing user information from external "
@@ -588,6 +736,7 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
                 email_opted_in=True,
                 active=True,
                 email=external_user.email,
+                is_admin=external_user.is_admin,
             ),
         )
     except KeyError:
@@ -603,15 +752,30 @@ def authenticate_external_user(external_access_token: str) -> AuthContext:
                 email_opted_in=True,
                 active=True,
                 email=external_user.email,
+                is_admin=external_user.is_admin,
             )
         )
 
         with AnalyticsContext() as context:
             context.user_id = user.id
             context.identify(
-                traits={"email": user.email, "source": "external_auth"}
+                traits={
+                    "email": external_user.email,
+                    "source": "external_auth",
+                }
             )
-            context.alias(user_id=user.id, previous_id=external_user.id)
+            context.alias(user_id=external_user.id, previous_id=user.id)
+
+    # This is the best spot to update the onboarding state to mark the
+    # "zenml login" step as completed for ZenML Pro servers, because the
+    # user has just successfully logged in. However, we need to differentiate
+    # between web clients (i.e. the dashboard) and CLI clients (i.e. the
+    # zenml CLI).
+    user_agent = request.headers.get("User-Agent", "").lower()
+    if "zenml/" in user_agent:
+        store.update_onboarding_state(
+            completed_steps={OnboardingStep.DEVICE_VERIFIED}
+        )
 
     return AuthContext(user=user)
 
@@ -629,14 +793,14 @@ def authenticate_api_key(
         The authentication context reflecting the authenticated service account.
 
     Raises:
-        AuthorizationException: If the service account could not be authorized.
+        CredentialsNotValid: If the service account could not be authorized.
     """
     try:
         decoded_api_key = APIKey.decode_api_key(api_key)
     except ValueError:
         error = "Authentication error: error decoding API key"
         logger.exception(error)
-        raise AuthorizationException(error)
+        raise CredentialsNotValid(error)
 
     internal_api_key = _fetch_and_verify_api_key(
         api_key_id=decoded_api_key.id, key_to_verify=decoded_api_key.key
@@ -650,6 +814,85 @@ 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. If not set,
+            the default value is determined automatically based on the server
+            configuration and type of token. If set to 0, the token will not
+            expire.
+        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 == 0:
+        expires_in = None
+    elif expires_in is not None:
+        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:
@@ -661,19 +904,19 @@ def http_authentication(
     Returns:
         The authentication context reflecting the authenticated user.
 
-    Raises:
-        HTTPException: If the credentials are invalid.
+    # noqa: DAR401
     """
     try:
         return authenticate_credentials(
             user_name_or_id=credentials.username, password=credentials.password
         )
-    except AuthorizationException as e:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=str(e),
-            headers={"WWW-Authenticate": "Basic"},
-        )
+    except CredentialsNotValid as e:
+        # We want to be very explicit here and return a CredentialsNotValid
+        # exception encoded as a 401 Unauthorized error encoded, so that the
+        # client can distinguish between a 401 error due to invalid credentials
+        # and other 401 errors and handle them accordingly by throwing away the
+        # current access token and re-authenticating.
+        raise http_exception_from_error(e)
 
 
 class CookieOAuth2TokenBearer(OAuth2PasswordBearer):
@@ -716,17 +959,17 @@ def oauth2_authentication(
     Returns:
         The authentication context reflecting the authenticated user.
 
-    Raises:
-        HTTPException: If the JWT token could not be authorized.
+    # noqa: DAR401
     """
     try:
         auth_context = authenticate_credentials(access_token=token)
-    except AuthorizationException as e:
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail=str(e),
-            headers={"WWW-Authenticate": "Bearer"},
-        )
+    except CredentialsNotValid as e:
+        # We want to be very explicit here and return a CredentialsNotValid
+        # exception encoded as a 401 Unauthorized error encoded, so that the
+        # client can distinguish between a 401 error due to invalid credentials
+        # and other 401 errors and handle them accordingly by throwing away the
+        # current access token and re-authenticating.
+        raise http_exception_from_error(e)
 
     return auth_context