Infer base URL from read token in query client (#1088)

Author: Viicos

logfire/experimental/query_client.py

@@ -4,6 +4,10 @@
 from types import TracebackType
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypedDict, TypeVar
 
+from typing_extensions import Self
+
+from logfire._internal.config import get_base_url_from_token
+
 try:
     from httpx import AsyncClient, Client, Response, Timeout
     from httpx._client import BaseClient
@@ -56,8 +60,6 @@ class RowQueryResults(TypedDict):
 
 
 T = TypeVar('T', bound=BaseClient)
-S = TypeVar('S', bound='LogfireQueryClient')
-R = TypeVar('R', bound='AsyncLogfireQueryClient')
 
 
 class _BaseLogfireQueryClient(Generic[T]):
@@ -102,13 +104,14 @@ class LogfireQueryClient(_BaseLogfireQueryClient[Client]):
     def __init__(
         self,
         read_token: str,
-        base_url: str = 'https://logfire-api.pydantic.dev/',
+        base_url: str | None = None,
         timeout: Timeout = DEFAULT_TIMEOUT,
         **client_kwargs: Any,
     ):
+        base_url = base_url or get_base_url_from_token(read_token)
         super().__init__(base_url, read_token, timeout, Client, **client_kwargs)
 
-    def __enter__(self: S) -> S:
+    def __enter__(self) -> Self:
         self.client.__enter__()
         return self
 
@@ -226,13 +229,14 @@ class AsyncLogfireQueryClient(_BaseLogfireQueryClient[AsyncClient]):
     def __init__(
         self,
         read_token: str,
-        base_url: str = 'https://logfire-api.pydantic.dev/',
+        base_url: str | None = None,
         timeout: Timeout = DEFAULT_TIMEOUT,
         **async_client_kwargs: Any,
     ):
+        base_url = base_url or get_base_url_from_token(read_token)
         super().__init__(base_url, read_token, timeout, AsyncClient, **async_client_kwargs)
 
-    async def __aenter__(self: R) -> R:
+    async def __aenter__(self) -> Self:
         await self.client.__aenter__()
         return self
 

tests/aaa_query_client/test_query_client.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import sys
 from datetime import datetime, timezone
 
@@ -21,6 +23,22 @@
 ]
 
 
+@pytest.mark.parametrize('client_class', [AsyncLogfireQueryClient, LogfireQueryClient])
+@pytest.mark.parametrize(
+    ['token', 'expected'],
+    [
+        ('pylf_v1_us_0kYhc414Ys2FNDRdt5vFB05xFx5NjVcbcBMy4Kp6PH0W', 'https://logfire-us.pydantic.dev'),
+        ('pylf_v1_eu_0kYhc414Ys2FNDRdt5vFB05xFx5NjVcbcBMy4Kp6PH0W', 'https://logfire-eu.pydantic.dev'),
+        ('0kYhc414Ys2FNDRdt5vFB05xFx5NjVcbcBMy4Kp6PH0W', 'https://logfire-us.pydantic.dev'),
+    ],
+)
+def test_infers_base_url_from_token(
+    client_class: type[AsyncLogfireQueryClient | LogfireQueryClient], token: str, expected: str
+):
+    client = client_class(read_token=token)
+    assert client.base_url == expected
+
+
 def test_read_sync():
     with LogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
         sql = """