truefoundry 0.2.10__py3-none-any.whl → 0.3.0rc1__py3-none-any.whl

truefoundry/deploy/lib/auth/credential_file_manager.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import os
+import threading
+from functools import lru_cache, wraps
+from pathlib import Path
+from typing import Optional
+
+from filelock import FileLock, Timeout
+
+from truefoundry.deploy.lib.const import CREDENTIAL_FILEPATH
+from truefoundry.deploy.lib.model.entity import CredentialsFileContent
+from truefoundry.logger import logger
+
+
+def _ensure_lock_taken(method):
+    @wraps(method)
+    def lock_guard(self, *method_args, **method_kwargs):
+        if not self.lock_taken():
+            raise Exception(
+                "Trying to write to credential file without using with block"
+            )
+        return method(self, *method_args, **method_kwargs)
+
+    return lock_guard
+
+
+CRED_FILE_THREAD_LOCK = threading.RLock()
+
+
+@lru_cache(maxsize=None)
+def get_file_lock(lock_file_path: str) -> FileLock:
+    return FileLock(lock_file_path)
+
+
+class CredentialsFileManager:
+    def __init__(
+        self,
+        credentials_file_path: Path = CREDENTIAL_FILEPATH,
+        lock_timeout: float = 60.0,
+    ) -> None:
+        credentials_file_path = credentials_file_path.absolute()
+        logger.debug("credential file path %r", credentials_file_path)
+
+        credentials_lock_file_path = f"{credentials_file_path}.lock"
+        logger.debug("credential lock file path %r", credentials_lock_file_path)
+
+        self._credentials_file_path = credentials_file_path
+        cred_file_dir = credentials_file_path.parent
+        cred_file_dir.mkdir(exist_ok=True, parents=True)
+
+        self._file_lock = get_file_lock(credentials_lock_file_path)
+        self._lock_timeout = lock_timeout
+        self._lock_owner: Optional[int] = None
+
+    def __enter__(self) -> CredentialsFileManager:
+        # The lock objects are recursive locks, which means that once acquired, they will not block on successive lock requests:
+        lock_aquired = CRED_FILE_THREAD_LOCK.acquire(timeout=self._lock_timeout)
+        if not lock_aquired:
+            raise Exception(
+                "Could not aquire CRED_FILE_THREAD_LOCK"
+                f" in {self._lock_timeout} seconds"
+            )
+        try:
+            self._file_lock.acquire(timeout=self._lock_timeout)
+        except Timeout as ex:
+            raise Exception(
+                f"Failed to aquire lock on credential file within {self._lock_timeout} seconds.\n"
+                "Is any other process trying to login?"
+            ) from ex
+        logger.debug("Acquired file and thread lock to access credential file")
+        self._lock_owner = threading.get_ident()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self._file_lock.release()
+        CRED_FILE_THREAD_LOCK.release()
+        logger.debug("Released file and thread lock to access credential file")
+        self._lock_owner = None
+
+    def lock_taken(self) -> bool:
+        return self._lock_owner == threading.get_ident()
+
+    @_ensure_lock_taken
+    def read(self) -> CredentialsFileContent:
+        try:
+            return CredentialsFileContent.parse_file(self._credentials_file_path)
+        except Exception as ex:
+            raise Exception(
+                "Error while reading the credentials file "
+                f"{self._credentials_file_path}. Please login again "
+                "using `tfy login --relogin` or `tfy.login(relogin=True)` function"
+            ) from ex
+
+    @_ensure_lock_taken
+    def write(self, credentials_file_content: CredentialsFileContent) -> None:
+        if not isinstance(credentials_file_content, CredentialsFileContent):
+            raise Exception(
+                "Only object of type `CredentialsFileContent` is allowed. "
+                f"Got {type(credentials_file_content)}"
+            )
+        logger.debug("Updating the credential file content")
+        with open(self._credentials_file_path, "w", encoding="utf8") as file:
+            file.write(credentials_file_content.json())
+
+    @_ensure_lock_taken
+    def delete(self) -> bool:
+        if not os.path.exists(self._credentials_file_path):
+            return False
+        os.remove(self._credentials_file_path)
+        return True
+
+    @_ensure_lock_taken
+    def exists(self) -> bool:
+        return self._credentials_file_path.exists()

truefoundry/deploy/lib/auth/credential_provider.py

@@ -0,0 +1,131 @@
+import os
+import threading
+from abc import ABC, abstractmethod
+
+from truefoundry.deploy.lib.auth.auth_service_client import AuthServiceClient
+from truefoundry.deploy.lib.auth.credential_file_manager import (
+    CredentialsFileContent,
+    CredentialsFileManager,
+)
+from truefoundry.deploy.lib.clients.utils import resolve_base_url
+from truefoundry.deploy.lib.const import API_KEY_ENV_NAME
+from truefoundry.deploy.lib.model.entity import Token
+from truefoundry.logger import logger
+
+TOKEN_REFRESH_LOCK = threading.RLock()
+
+
+class CredentialProvider(ABC):
+    @property
+    @abstractmethod
+    def token(self) -> Token: ...
+
+    @property
+    @abstractmethod
+    def base_url(self) -> str: ...
+
+    @staticmethod
+    @abstractmethod
+    def can_provide() -> bool: ...
+
+
+class EnvCredentialProvider(CredentialProvider):
+    def __init__(self) -> None:
+        from truefoundry.deploy.lib.clients.servicefoundry_client import (
+            ServiceFoundryServiceClient,
+        )
+
+        logger.debug("Using env var credential provider")
+        api_key = os.getenv(API_KEY_ENV_NAME)
+        if not api_key:
+            raise Exception(
+                f"Value of {API_KEY_ENV_NAME} env var should be non-empty string"
+            )
+        # TODO: Read host from cred file as well.
+        base_url = resolve_base_url().strip("/")
+        self._host = base_url
+        self._auth_service = AuthServiceClient(base_url=base_url)
+
+        servicefoundry_client = ServiceFoundryServiceClient(
+            init_session=False, base_url=base_url
+        )
+        self._token: Token = servicefoundry_client.get_token_from_api_key(
+            api_key=api_key
+        )
+
+    @staticmethod
+    def can_provide() -> bool:
+        return API_KEY_ENV_NAME in os.environ
+
+    @property
+    def token(self) -> Token:
+        with TOKEN_REFRESH_LOCK:
+            if self._token.is_going_to_be_expired():
+                logger.info("Refreshing access token")
+                self._token = self._auth_service.refresh_token(
+                    self._token, self.base_url
+                )
+            return self._token
+
+    @property
+    def base_url(self) -> str:
+        return self._host
+
+
+class FileCredentialProvider(CredentialProvider):
+    def __init__(self) -> None:
+        logger.debug("Using file credential provider")
+        self._cred_file = CredentialsFileManager()
+
+        with self._cred_file:
+            self._last_cred_file_content = self._cred_file.read()
+            self._token = self._last_cred_file_content.to_token()
+            self._host = self._last_cred_file_content.host
+
+        self._auth_service = AuthServiceClient(base_url=self._host)
+
+    @staticmethod
+    def can_provide() -> bool:
+        with CredentialsFileManager() as cred_file:
+            return cred_file.exists()
+
+    @property
+    def token(self) -> Token:
+        with TOKEN_REFRESH_LOCK:
+            if not self._token.is_going_to_be_expired():
+                return self._token
+
+            logger.info("Refreshing access token")
+            with self._cred_file:
+                new_cred_file_content = self._cred_file.read()
+                new_token = new_cred_file_content.to_token()
+                new_host = new_cred_file_content.host
+
+                if new_cred_file_content == self._last_cred_file_content:
+                    self._token = self._auth_service.refresh_token(
+                        self._token, self.base_url
+                    )
+                    self._last_cred_file_content = CredentialsFileContent(
+                        host=self._host,
+                        access_token=self._token.access_token,
+                        refresh_token=self._token.refresh_token,
+                    )
+                    self._cred_file.write(self._last_cred_file_content)
+                    return self._token
+
+                if (
+                    new_host == self._host
+                    and new_token.to_user_info() == self._token.to_user_info()
+                ):
+                    self._last_cred_file_content = new_cred_file_content
+                    self._token = new_token
+                    # recursive
+                    return self.token
+
+                raise Exception(
+                    "Credentials on disk changed while mlfoundry was running."
+                )
+
+    @property
+    def base_url(self) -> str:
+        return self._host

truefoundry/deploy/lib/auth/servicefoundry_session.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from truefoundry.deploy.lib.auth.credential_provider import (
+    CredentialProvider,
+    EnvCredentialProvider,
+    FileCredentialProvider,
+)
+from truefoundry.deploy.lib.model.entity import UserInfo
+from truefoundry.logger import logger
+
+ACTIVE_SESSION: Optional[ServiceFoundrySession] = None
+
+
+class ServiceFoundrySession:
+    def __init__(self) -> None:
+        self._cred_provider = self._get_cred_provider()
+        self._user_info: UserInfo = self._cred_provider.token.to_user_info()
+
+        global ACTIVE_SESSION
+        if (ACTIVE_SESSION is None) or (
+            ACTIVE_SESSION
+            and ACTIVE_SESSION.base_url != self.base_url
+            and ACTIVE_SESSION.user_info != self.user_info
+        ):
+            logger.info(
+                "Logged in to %r as %r (%s)",
+                self.base_url,
+                self.user_info.user_id,
+                self.user_info.email or self.user_info.user_type.value,
+            )
+        ACTIVE_SESSION = self
+
+    @staticmethod
+    def _get_cred_provider() -> CredentialProvider:
+        final_cred_provider = None
+        for cred_provider in [EnvCredentialProvider, FileCredentialProvider]:
+            if cred_provider.can_provide():
+                final_cred_provider = cred_provider()
+                break
+        if final_cred_provider is None:
+            raise Exception(
+                "Please login again using `tfy login --relogin`"
+                "or `tfy.login(relogin=True)` function"
+            )
+        return final_cred_provider
+
+    @property
+    def access_token(self) -> str:
+        return self._cred_provider.token.access_token
+
+    @property
+    def base_url(self) -> str:
+        return self._cred_provider.base_url
+
+    @property
+    def user_info(self) -> UserInfo:
+        return self._user_info