Partial feedback

Author: Viicos

logfire/_internal/auth.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import platform
+import re
 import sys
 import warnings
 from dataclasses import dataclass
@@ -30,14 +31,37 @@
 """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 TokensFileData(TypedDict):
+class UserTokensFileData(TypedDict):
     """Content of the file containing the user tokens."""
 
     tokens: dict[str, UserTokenData]
@@ -87,43 +111,46 @@ def __str__(self) -> str:
 @dataclass
 class UserTokenCollection:
     """A collection of user tokens."""
-    user_tokens: list[UserToken]
+
+    user_tokens: dict[str, UserToken]
+    """A mapping between base URLs and user tokens."""
 
     @classmethod
-    def from_tokens(cls, tokens: TokensFileData) -> Self:
-        return cls(
-            user_tokens=[
-                UserToken.from_user_token_data(url, data)  # pyright: ignore[reportArgumentType], waiting for PEP 728
-                for url, data in tokens.items()
-            ]
-        )
+    def empty(cls) -> Self:
+        """Create an empty user token collection."""
+        return cls(user_tokens={})
+
+    @classmethod
+    def from_file_data(cls, file_data: UserTokensFileData) -> Self:
+        return cls(user_tokens={url: UserToken(base_url=url, **data) for url, data in file_data['tokens'].items()})
 
     @classmethod
     def from_tokens_file(cls, file: Path) -> Self:
-        return cls.from_tokens(cast(TokensFileData, read_toml_file(file)))
+        return cls.from_file_data(cast(UserTokensFileData, read_toml_file(file)))
 
     def get_token(self, base_url: str | None = None) -> UserToken:
+        tokens_list = list(self.user_tokens.values())
         if base_url is not None:
-            token = next((t for t in self.user_tokens if t.base_url == base_url), None)
+            token = next((t for t in tokens_list if t.base_url == base_url), None)
             if token is None:
                 raise LogfireConfigError(
                     f'No user token was found matching the {base_url} Logfire URL. '
                     'Please run `logfire auth` to authenticate.'
                 )
         else:
-            if len(self.user_tokens) == 1:
-                token = self.user_tokens[0]
-            elif len(self.user_tokens) >= 2:
+            if 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(self.user_tokens, start=1)
+                    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(self.user_tokens) + 1)],
+                    choices=[str(i) for i in range(1, len(tokens_list) + 1)],
                 )
-                token = self.user_tokens[int_choice - 1]
-            else:  # self.user_tokens == []
+                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:
@@ -132,28 +159,21 @@ def get_token(self, base_url: str | None = None) -> UserToken:
 
     def is_logged_in(self, base_url: str | None = None) -> bool:
         if base_url is not None:
-            tokens = (t for t in self.user_tokens if t.base_url == base_url)
+            tokens = (t for t in self.user_tokens.values() if t.base_url == base_url)
         else:
-            tokens = self.user_tokens
+            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:
-        existing_token = next((t for t in self.user_tokens if t.base_url == base_url), None)
-        if existing_token:
-            token_index = self.user_tokens.index(existing_token)
-            self.user_tokens.remove(existing_token)
-        else:
-            token_index = len(self.user_tokens)
-
         user_token = UserToken.from_user_token_data(base_url, token)
-        self.user_tokens.insert(token_index, user_token)
+        self.user_tokens[base_url] = UserToken.from_user_token_data(base_url, token)
         return user_token
 
     def dump(self, path: Path) -> None:
         # There's no standard library package to write TOML files, so we'll write it manually.
         with path.open('w') as f:
-            for user_token in self.user_tokens:
-                f.write(f'[tokens."{user_token.base_url}"]\n')
+            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')
 

logfire/_internal/cli.py

@@ -209,7 +209,7 @@ def parse_auth(args: argparse.Namespace) -> None:
     if DEFAULT_FILE.is_file():
         tokens_collection = default_token_collection()
     else:
-        tokens_collection = UserTokenCollection.from_tokens({'tokens': {}})
+        tokens_collection = UserTokenCollection.empty()
 
     logged_in = tokens_collection.is_logged_in(logfire_url)
 
@@ -270,8 +270,7 @@ def parse_auth(args: argparse.Namespace) -> None:
 def parse_list_projects(args: argparse.Namespace) -> None:
     """List user projects."""
     logfire_url: str | None = args.logfire_url
-    user_token = default_token_collection().get_token(logfire_url)
-    client = LogfireClient(user_token)
+    client = LogfireClient.from_url(logfire_url)
 
     projects = client.get_user_projects()
     if projects:
@@ -303,8 +302,7 @@ def parse_create_new_project(args: argparse.Namespace) -> None:
     """Create a new project."""
     data_dir = Path(args.data_dir)
     logfire_url: str | None = args.logfire_url
-    user_token = default_token_collection().get_token(logfire_url)
-    client = LogfireClient(user_token)
+    client = LogfireClient.from_url(logfire_url)
 
     project_name = args.project_name
     organization = args.org
@@ -315,16 +313,15 @@ def parse_create_new_project(args: argparse.Namespace) -> None:
         default_organization=default_organization,
         project_name=project_name,
     )
-    credentials = _write_credentials(project_info, data_dir, user_token.base_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: str | None = args.logfire_url
-    user_token = default_token_collection().get_token(logfire_url)
-    client = LogfireClient(user_token)
+    client = LogfireClient.from_url(logfire_url)
 
     project_name = args.project_name
     organization = args.org

logfire/_internal/client.py

@@ -1,19 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from urllib.parse import urljoin
 
 from requests import Response, Session
+from typing_extensions import Self
 
 from logfire.exceptions import LogfireConfigError
 from logfire.version import VERSION
 
+from .auth import UserToken, UserTokenCollection, default_token_collection
 from .utils import UnexpectedResponse
 
-if TYPE_CHECKING:
-    from .auth import UserToken
-
-
 UA_HEADER = f'logfire/{VERSION}'
 
 
@@ -27,6 +25,12 @@ def __init__(self, reason: str, /) -> None:
 
 
 class LogfireClient:
+    """A Logfire HTTP client to interact with the API.
+
+    Args:
+        user_token: The user token to use when authenticating against the API.
+    """
+
     def __init__(self, user_token: UserToken) -> None:
         if user_token.is_expired:
             raise RuntimeError
@@ -35,6 +39,20 @@ def __init__(self, user_token: UserToken) -> None:
         self._session = Session()
         self._session.headers.update({'Authorization': self._token, 'User-Agent': UA_HEADER})
 
+    @classmethod
+    def from_url(cls, base_url: str | None, token_collection: UserTokenCollection | None = None) -> Self:
+        """Create a client from the provided base URL.
+
+        Args:
+            base_url: The base URL to use when looking for a user token. If `None`, will prompt
+                the user into selecting a token from the token collection (or, if only one available,
+                use it directly).
+            token_collection: The token collection to use when looking for the user token. Defaults
+                to the default token collection from `~/.logfire/default.toml`.
+        """
+        token_collection = token_collection or default_token_collection()
+        return cls(user_token=token_collection.get_token(base_url))
+
     def _get(self, endpoint: str) -> Response:
         response = self._session.get(urljoin(self.base_url, endpoint))
         UnexpectedResponse.raise_for_status(response)

logfire/_internal/config.py

@@ -65,7 +65,6 @@
 from logfire.version import VERSION
 
 from ..propagate import NoExtractTraceContextPropagator, WarnOnExtractTraceContextPropagator
-from .auth import default_token_collection
 from .client import InvalidProjectName, LogfireClient, ProjectAlreadyExists
 from .config_params import ParamManager, PydanticPluginRecordValues
 from .constants import (
@@ -897,8 +896,7 @@ def add_span_processor(span_processor: SpanProcessor) -> None:
                     # if we still don't have a token, try initializing a new project and writing a new creds file
                     # note, we only do this if `send_to_logfire` is explicitly `True`, not 'if-token-present'
                     if self.send_to_logfire is True and credentials is None:
-                        user_token = default_token_collection().get_token(self.advanced.base_url)
-                        client = LogfireClient(user_token)
+                        client = LogfireClient.from_url(self.advanced.base_url)
                         credentials = LogfireCredentials.initialize_project(client=client)
                         credentials.write_creds_file(self.data_dir)
 
@@ -1596,23 +1594,6 @@ def _get_creds_file(creds_dir: Path) -> Path:
     return creds_dir / CREDENTIALS_FILENAME
 
 
-def _get_token_repr(url: str, token: str) -> str:
-    region = 'us'
-    if match := PYDANTIC_LOGFIRE_TOKEN_PATTERN.match(token):
-        region = match.group('region')
-        if region not in REGIONS:
-            region = 'us'
-
-    token_repr = f'{region.upper()} ({url}) - '
-    if match:
-        # new_token, include prefix and 5 chars
-        token_repr += match.group('safe_part') + match.group('token')[:5]
-    else:
-        token_repr += token[:5]
-    token_repr += '****'
-    return token_repr
-
-
 def get_base_url_from_token(token: str) -> str:
     """Get the base API URL from the token's region."""
     # default to US for tokens that were created before regions were added: