Refactor user tokens, introduce Logfire client (#981)

Author: Viicos

logfire/_internal/auth.py

@@ -1,37 +1,198 @@
 from __future__ import annotations
 
 import platform
+import re
 import warnings
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TypedDict
+from typing import TypedDict, cast
 from urllib.parse import urljoin
 
 import requests
+from rich.prompt import IntPrompt
+from typing_extensions import Self
 
 from logfire.exceptions import LogfireConfigError
 
-from .utils import UnexpectedResponse
+from .utils import UnexpectedResponse, read_toml_file
 
 HOME_LOGFIRE = Path.home() / '.logfire'
 """Folder used to store global configuration, and user tokens."""
 DEFAULT_FILE = HOME_LOGFIRE / 'default.toml'
 """File used to store user tokens."""
 
 
+PYDANTIC_LOGFIRE_TOKEN_PATTERN = re.compile(
+    r'^(?P<safe_part>pylf_v(?P<version>[0-9]+)_(?P<region>[a-z]+)_)(?P<token>[a-zA-Z0-9]+)$'
+)
+
+
+class _RegionData(TypedDict):
+    base_url: str
+    gcp_region: str
+
+
+REGIONS: dict[str, _RegionData] = {
+    'us': {
+        'base_url': 'https://logfire-us.pydantic.dev',
+        'gcp_region': 'us-east4',
+    },
+    'eu': {
+        'base_url': 'https://logfire-eu.pydantic.dev',
+        'gcp_region': 'europe-west4',
+    },
+}
+"""The existing Logfire regions."""
+
+
 class UserTokenData(TypedDict):
     """User token data."""
 
     token: str
     expiration: str
 
 
-class DefaultFile(TypedDict):
-    """Content of the default.toml file."""
+class UserTokensFileData(TypedDict, total=False):
+    """Content of the file containing the user tokens."""
 
     tokens: dict[str, UserTokenData]
 
 
+@dataclass
+class UserToken:
+    """A user token."""
+
+    token: str
+    base_url: str
+    expiration: str
+
+    @classmethod
+    def from_user_token_data(cls, base_url: str, token: UserTokenData) -> Self:
+        return cls(
+            token=token['token'],
+            base_url=base_url,
+            expiration=token['expiration'],
+        )
+
+    @property
+    def is_expired(self) -> bool:
+        """Whether the token is expired."""
+        return datetime.now(tz=timezone.utc) >= datetime.fromisoformat(self.expiration.rstrip('Z')).replace(
+            tzinfo=timezone.utc
+        )
+
+    def __str__(self) -> str:
+        region = 'us'
+        if match := PYDANTIC_LOGFIRE_TOKEN_PATTERN.match(self.token):
+            region = match.group('region')
+            if region not in REGIONS:
+                region = 'us'
+
+        token_repr = f'{region.upper()} ({self.base_url}) - '
+        if match:
+            token_repr += match.group('safe_part') + match.group('token')[:5]
+        else:
+            token_repr += self.token[:5]
+        token_repr += '****'
+        return token_repr
+
+
+@dataclass
+class UserTokenCollection:
+    """A collection of user tokens, read from a user tokens file.
+
+    Args:
+        path: The path where the user tokens will be stored. If the path doesn't exist,
+            an empty collection is created. Defaults to `~/.logfire/default.toml`.
+    """
+
+    user_tokens: dict[str, UserToken]
+    """A mapping between base URLs and user tokens."""
+
+    path: Path
+    """The path where the user tokens are stored."""
+
+    def __init__(self, path: Path | None = None) -> None:
+        # FIXME: we can't set the default value of `path` to `DEFAULT_FILE`, otherwise
+        # `mock.patch()` doesn't work:
+        self.path = path if path is not None else DEFAULT_FILE
+        try:
+            data = cast(UserTokensFileData, read_toml_file(self.path))
+        except FileNotFoundError:
+            data: UserTokensFileData = {}
+        self.user_tokens = {url: UserToken(base_url=url, **token) for url, token in data.get('tokens', {}).items()}
+
+    def get_token(self, base_url: str | None = None) -> UserToken:
+        """Get a user token from the collection.
+
+        Args:
+            base_url: Only look for user tokens valid for this base URL. If not provided,
+                all the tokens of the collection will be considered: if only one token is
+                available, it will be used, otherwise the user will be prompted to choose
+                a token.
+
+        Raises:
+            LogfireConfigError: If no user token is found (no token matched the base URL,
+                the collection is empty, or the selected token is expired).
+        """
+        tokens_list = list(self.user_tokens.values())
+
+        if base_url is not None:
+            token = self.user_tokens.get(base_url)
+            if token is None:
+                raise LogfireConfigError(
+                    f'No user token was found matching the {base_url} Logfire URL. '
+                    'Please run `logfire auth` to authenticate.'
+                )
+        elif len(tokens_list) == 1:
+            token = tokens_list[0]
+        elif len(tokens_list) >= 2:
+            choices_str = '\n'.join(
+                f'{i}. {token} ({"expired" if token.is_expired else "valid"})'
+                for i, token in enumerate(tokens_list, start=1)
+            )
+            int_choice = IntPrompt.ask(
+                f'Multiple user tokens found. Please select one:\n{choices_str}\n',
+                choices=[str(i) for i in range(1, len(tokens_list) + 1)],
+            )
+            token = tokens_list[int_choice - 1]
+        else:  # tokens_list == []
+            raise LogfireConfigError('No user tokens are available. Please run `logfire auth` to authenticate.')
+
+        if token.is_expired:
+            raise LogfireConfigError(f'User token {token} is expired. Please run `logfire auth` to authenticate.')
+        return token
+
+    def is_logged_in(self, base_url: str | None = None) -> bool:
+        """Check whether the user token collection contains at least one valid user token.
+
+        Args:
+            base_url: Only check for user tokens valid for this base URL. If not provided,
+                all the tokens of the collection will be considered.
+        """
+        if base_url is not None:
+            tokens = (t for t in self.user_tokens.values() if t.base_url == base_url)
+        else:
+            tokens = self.user_tokens.values()
+        return any(not t.is_expired for t in tokens)
+
+    def add_token(self, base_url: str, token: UserTokenData) -> UserToken:
+        """Add a user token to the collection."""
+        self.user_tokens[base_url] = user_token = UserToken.from_user_token_data(base_url, token)
+        self._dump()
+        return user_token
+
+    def _dump(self) -> None:
+        """Dump the user token collection as TOML to the provided path."""
+        # There's no standard library package to write TOML files, so we'll write it manually.
+        with self.path.open('w') as f:
+            for base_url, user_token in self.user_tokens.items():
+                f.write(f'[tokens."{base_url}"]\n')
+                f.write(f'token = "{user_token.token}"\n')
+                f.write(f'expiration = "{user_token.expiration}"\n')
+
+
 class NewDeviceFlow(TypedDict):
     """Matches model of the same name in the backend."""
 
@@ -91,17 +252,3 @@ def poll_for_token(session: requests.Session, device_code: str, base_api_url: st
             opt_user_token: UserTokenData | None = res.json()
             if opt_user_token:
                 return opt_user_token
-
-
-def is_logged_in(data: DefaultFile, logfire_url: str) -> bool:
-    """Check if the user is logged in.
-
-    Returns:
-        True if the user is logged in, False otherwise.
-    """
-    for url, info in data['tokens'].items():  # pragma: no branch
-        # token expirations are in UTC
-        expiry_date = datetime.fromisoformat(info['expiration'].rstrip('Z')).replace(tzinfo=timezone.utc)
-        if url == logfire_url and datetime.now(tz=timezone.utc) < expiry_date:  # pragma: no branch
-            return True
-    return False  # pragma: no cover

logfire/_internal/cli.py

@@ -14,7 +14,7 @@
 from collections.abc import Sequence
 from operator import itemgetter
 from pathlib import Path
-from typing import Any, cast
+from typing import Any
 from urllib.parse import urlparse
 
 import requests
@@ -24,11 +24,17 @@
 from logfire.propagate import ContextCarrier, get_context
 
 from ..version import VERSION
-from .auth import DEFAULT_FILE, HOME_LOGFIRE, DefaultFile, is_logged_in, poll_for_token, request_device_code
+from .auth import (
+    DEFAULT_FILE,
+    HOME_LOGFIRE,
+    UserTokenCollection,
+    poll_for_token,
+    request_device_code,
+)
+from .client import LogfireClient
 from .config import REGIONS, LogfireCredentials, get_base_url_from_token
 from .config_params import ParamManager
 from .tracer import SDKTracerProvider
-from .utils import read_toml_file
 
 BASE_OTEL_INTEGRATION_URL = 'https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/'
 BASE_DOCS_URL = 'https://logfire.pydantic.dev/docs'
@@ -51,7 +57,7 @@ def parse_whoami(args: argparse.Namespace) -> None:
     """Show user authenticated username and the URL to your Logfire project."""
     data_dir = Path(args.data_dir)
     param_manager = ParamManager.create(data_dir)
-    base_url = param_manager.load_param('base_url', args.logfire_url)
+    base_url: str | None = param_manager.load_param('base_url', args.logfire_url)
     token = param_manager.load_param('token')
 
     if token:
@@ -61,12 +67,15 @@ def parse_whoami(args: argparse.Namespace) -> None:
             credentials.print_token_summary()
             return
 
-    current_user = LogfireCredentials.get_current_user(session=args._session, logfire_api_url=base_url)
-    if current_user is None:
+    try:
+        client = LogfireClient.from_url(base_url)
+    except LogfireConfigError:
         sys.stderr.write('Not logged in. Run `logfire auth` to log in.\n')
     else:
+        current_user = client.get_user_information()
         username = current_user['name']
         sys.stderr.write(f'Logged in as: {username}\n')
+
     credentials = LogfireCredentials.load_creds_file(data_dir)
     if credentials is None:
         sys.stderr.write(f'No Logfire credentials found in {data_dir.resolve()}\n')
@@ -198,16 +207,10 @@ def parse_auth(args: argparse.Namespace) -> None:
 
     This will authenticate your machine with Logfire and store the credentials.
     """
-    logfire_url = args.logfire_url
-    if DEFAULT_FILE.is_file():
-        data = cast(DefaultFile, read_toml_file(DEFAULT_FILE))
-    else:
-        data: DefaultFile = {'tokens': {}}
+    logfire_url: str | None = args.logfire_url
 
-    if logfire_url:
-        logged_in = is_logged_in(data, logfire_url)
-    else:
-        logged_in = any(is_logged_in(data, url) for url in data['tokens'])
+    tokens_collection = UserTokenCollection()
+    logged_in = tokens_collection.is_logged_in(logfire_url)
 
     if logged_in:
         sys.stderr.writelines(
@@ -256,22 +259,16 @@ def parse_auth(args: argparse.Namespace) -> None:
         )
     )
 
-    data['tokens'][logfire_url] = poll_for_token(args._session, device_code, logfire_url)
+    tokens_collection.add_token(logfire_url, poll_for_token(args._session, device_code, logfire_url))
     sys.stderr.write('Successfully authenticated!\n')
-
-    # There's no standard library package to write TOML files, so we'll write it manually.
-    with DEFAULT_FILE.open('w') as f:
-        for url, info in data['tokens'].items():
-            f.write(f'[tokens."{url}"]\n')
-            f.write(f'token = "{info["token"]}"\n')
-            f.write(f'expiration = "{info["expiration"]}"\n')
-
     sys.stderr.write(f'\nYour Logfire credentials are stored in {DEFAULT_FILE}\n')
 
 
 def parse_list_projects(args: argparse.Namespace) -> None:
     """List user projects."""
-    projects = LogfireCredentials.get_user_projects(session=args._session, logfire_api_url=args.logfire_url)
+    client = LogfireClient.from_url(args.logfire_url)
+
+    projects = client.get_user_projects()
     if projects:
         sys.stderr.write(
             _pretty_table(
@@ -300,42 +297,37 @@ def _write_credentials(project_info: dict[str, Any], data_dir: Path, logfire_api
 def parse_create_new_project(args: argparse.Namespace) -> None:
     """Create a new project."""
     data_dir = Path(args.data_dir)
-    logfire_url = args.logfire_url
-    if logfire_url is None:  # pragma: no cover
-        _, logfire_url = LogfireCredentials._get_user_token_data()  # type: ignore
+    client = LogfireClient.from_url(args.logfire_url)
+
     project_name = args.project_name
     organization = args.org
     default_organization = args.default_org
     project_info = LogfireCredentials.create_new_project(
-        session=args._session,
-        logfire_api_url=logfire_url,
+        client=client,
         organization=organization,
         default_organization=default_organization,
         project_name=project_name,
     )
-    credentials = _write_credentials(project_info, data_dir, logfire_url)
+    credentials = _write_credentials(project_info, data_dir, client.base_url)
     sys.stderr.write(f'Project created successfully. You will be able to view it at: {credentials.project_url}\n')
 
 
 def parse_use_project(args: argparse.Namespace) -> None:
     """Use an existing project."""
     data_dir = Path(args.data_dir)
-    logfire_url = args.logfire_url
-    if logfire_url is None:  # pragma: no cover
-        _, logfire_url = LogfireCredentials._get_user_token_data()  # type: ignore
+    client = LogfireClient.from_url(args.logfire_url)
+
     project_name = args.project_name
     organization = args.org
-
-    projects = LogfireCredentials.get_user_projects(session=args._session, logfire_api_url=logfire_url)
+    projects = client.get_user_projects()
     project_info = LogfireCredentials.use_existing_project(
-        session=args._session,
-        logfire_api_url=logfire_url,
+        client=client,
         projects=projects,
         organization=organization,
         project_name=project_name,
     )
     if project_info:
-        credentials = _write_credentials(project_info, data_dir, logfire_url)
+        credentials = _write_credentials(project_info, data_dir, client.base_url)
         sys.stderr.write(
             f'Project configured successfully. You will be able to view it at: {credentials.project_url}\n'
         )