Add info method to logfire query clients (#1204)

Author: hramezani

docs/how-to-guides/query-api.md

@@ -88,6 +88,10 @@ Here's an example of how to use these clients:
             df_from_csv = pl.read_csv(StringIO(await client.query_csv(sql=query)))
             print(df_from_csv)
 
+            # Get read token info
+            read_token_info = await client.info()
+            print(read_token_info)
+
 
     if __name__ == '__main__':
         import asyncio
@@ -132,6 +136,10 @@ Here's an example of how to use these clients:
             df_from_csv = pl.read_csv(StringIO(client.query_csv(sql=query)))
             print(df_from_csv)
 
+            # Get read token info
+            read_token_info = client.info()
+            print(read_token_info)
+
 
     if __name__ == '__main__':
         main()

logfire-api/logfire_api/experimental/query_client.pyi

@@ -1,5 +1,6 @@
 from _typeshed import Incomplete
 from datetime import datetime
+from uuid import UUID
 from httpx import AsyncClient, Client, Response, Timeout
 from httpx._client import BaseClient
 from logfire._internal.config import get_base_url_from_token as get_base_url_from_token
@@ -15,6 +16,11 @@ class QueryExecutionError(RuntimeError):
 class QueryRequestError(RuntimeError):
     """Raised when the query request is invalid."""
 
+class ReadTokenInfo(TypedDict, total=True):
+    """Information about the read token."""
+    organization_name: str
+    project_name: str
+
 class ColumnDetails(TypedDict):
     """The details of a column in the row-oriented JSON-format query results."""
     name: str
@@ -49,6 +55,8 @@ class LogfireQueryClient(_BaseLogfireQueryClient[Client]):
     def __init__(self, read_token: str, base_url: str | None = None, timeout: Timeout = ..., **client_kwargs: Any) -> None: ...
     def __enter__(self) -> Self: ...
     def __exit__(self, exc_type: type[BaseException] | None = None, exc_value: BaseException | None = None, traceback: TracebackType | None = None) -> None: ...
+    def info(self) -> ReadTokenInfo:
+       """Get information about the read token."""
     def query_json(self, sql: str, min_timestamp: datetime | None = None, max_timestamp: datetime | None = None, limit: int | None = None) -> QueryResults:
         """Query Logfire data and return the results as a column-oriented dictionary."""
     def query_json_rows(self, sql: str, min_timestamp: datetime | None = None, max_timestamp: datetime | None = None, limit: int | None = None) -> RowQueryResults:
@@ -71,6 +79,8 @@ class AsyncLogfireQueryClient(_BaseLogfireQueryClient[AsyncClient]):
     def __init__(self, read_token: str, base_url: str | None = None, timeout: Timeout = ..., **async_client_kwargs: Any) -> None: ...
     async def __aenter__(self) -> Self: ...
     async def __aexit__(self, exc_type: type[BaseException] | None = None, exc_value: BaseException | None = None, traceback: TracebackType | None = None) -> None: ...
+    async def info(self) -> ReadTokenInfo:
+       """Get information about the read token."""
     async def query_json(self, sql: str, min_timestamp: datetime | None = None, max_timestamp: datetime | None = None, limit: int | None = None) -> QueryResults:
         """Query Logfire data and return the results as a column-oriented dictionary."""
     async def query_json_rows(self, sql: str, min_timestamp: datetime | None = None, max_timestamp: datetime | None = None, limit: int | None = None) -> RowQueryResults:

logfire/experimental/query_client.py

@@ -32,6 +32,19 @@ class QueryRequestError(RuntimeError):
     pass
 
 
+class InfoRequestError(RuntimeError):
+    """Raised when the request for read token info fails because of unavailable information."""
+
+    pass
+
+
+class ReadTokenInfo(TypedDict, total=False):
+    """Information about the read token."""
+
+    organization_name: str
+    project_name: str
+
+
 class ColumnDetails(TypedDict):
     """The details of a column in the row-oriented JSON-format query results."""
 
@@ -123,6 +136,22 @@ def __exit__(
     ) -> None:
         self.client.__exit__(exc_type, exc_value, traceback)
 
+    def info(self) -> ReadTokenInfo:
+        """Get information about the read token."""
+        response = self.client.get('/api/read-token-info')
+        self.handle_response_errors(response)
+        token_info = response.json()
+        try:
+            # Keep keys defined in ReadTokenInfo
+            return {
+                'organization_name': token_info['organization_name'],
+                'project_name': token_info['project_name'],
+            }
+        except KeyError:
+            raise InfoRequestError(
+                'The read token info response is missing required fields: organization_name or project_name'
+            )
+
     def query_json(
         self,
         sql: str,
@@ -248,6 +277,22 @@ async def __aexit__(
     ) -> None:
         await self.client.__aexit__(exc_type, exc_value, traceback)
 
+    async def info(self) -> ReadTokenInfo:
+        """Get information about the read token."""
+        response = await self.client.get('/api/read-token-info')
+        self.handle_response_errors(response)
+        token_info = response.json()
+        # Keep keys defined in ReadTokenInfo
+        try:
+            return {
+                'organization_name': token_info['organization_name'],
+                'project_name': token_info['project_name'],
+            }
+        except KeyError:
+            raise InfoRequestError(
+                'The read token info response is missing required fields: organization_name or project_name'
+            )
+
     async def query_json(
         self,
         sql: str,

tests/aaa_query_client/cassettes/test_query_client/test_info_async.yaml

@@ -0,0 +1,40 @@
+interactions:
+- request:
+    body: ''
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      host:
+      - localhost:8000
+      user-agent:
+      - python-httpx/0.28.1
+    method: GET
+    uri: http://localhost:8000/api/read-token-info
+  response:
+    body:
+      string: '{"token_id":"b630bd4e-6a20-486b-95cb-c02a175767b2","organization_id":"710c0de0-dda1-40fb-9601-b72d120acb71","project_id":"11d5d03b-7fc6-4ad6-bcd9-d1cc4bc87986","organization_name":"test-org","project_name":"starter-project"}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '221'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 06 Jul 2025 12:57:40 GMT
+      Server:
+      - nginx/1.29.0
+      access-control-expose-headers:
+      - traceresponse
+      traceresponse:
+      - 00-0197dfd050205c9d29c0a1a4a934222c-03c3392466156875-01
+      x-api-version:
+      - UClDZTsUC0iVNKhEvlUL2JNm0ur0dBUm1u2/5He9ZtE=
+    status:
+      code: 200
+      message: OK
+version: 1

tests/aaa_query_client/cassettes/test_query_client/test_info_invalid_schema_async.yaml

@@ -0,0 +1,40 @@
+interactions:
+- request:
+    body: ''
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      host:
+      - localhost:8000
+      user-agent:
+      - python-httpx/0.28.1
+    method: GET
+    uri: http://localhost:8000/api/read-token-info
+  response:
+    body:
+      string: '{"token_id":"b630bd4e-6a20-486b-95cb-c02a175767b2","organization_id":"710c0de0-dda1-40fb-9601-b72d120acb71","project_id":"11d5d03b-7fc6-4ad6-bcd9-d1cc4bc87986","project_name":"starter-project"}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '221'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 06 Jul 2025 13:40:40 GMT
+      Server:
+      - nginx/1.29.0
+      access-control-expose-headers:
+      - traceresponse
+      traceresponse:
+      - 00-0197dff7ac3bae870d3350ef4f961b65-8c6f724ec401043c-01
+      x-api-version:
+      - UClDZTsUC0iVNKhEvlUL2JNm0ur0dBUm1u2/5He9ZtE=
+    status:
+      code: 200
+      message: OK
+version: 1

tests/aaa_query_client/cassettes/test_query_client/test_info_invalid_schema_sync.yaml

@@ -0,0 +1,40 @@
+interactions:
+- request:
+    body: ''
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      host:
+      - localhost:8000
+      user-agent:
+      - python-httpx/0.28.1
+    method: GET
+    uri: http://localhost:8000/api/read-token-info
+  response:
+    body:
+      string: '{"token_id":"b630bd4e-6a20-486b-95cb-c02a175767b2","organization_id":"710c0de0-dda1-40fb-9601-b72d120acb71","project_id":"11d5d03b-7fc6-4ad6-bcd9-d1cc4bc87986","project_name":"starter-project"}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '221'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 06 Jul 2025 13:44:21 GMT
+      Server:
+      - nginx/1.29.0
+      access-control-expose-headers:
+      - traceresponse
+      traceresponse:
+      - 00-0197dffb0d73874ace548d5d84ed156e-d4a83dfee7bcd915-01
+      x-api-version:
+      - UClDZTsUC0iVNKhEvlUL2JNm0ur0dBUm1u2/5He9ZtE=
+    status:
+      code: 200
+      message: OK
+version: 1

tests/aaa_query_client/cassettes/test_query_client/test_info_sync.yaml

@@ -0,0 +1,40 @@
+interactions:
+- request:
+    body: ''
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      host:
+      - localhost:8000
+      user-agent:
+      - python-httpx/0.28.1
+    method: GET
+    uri: http://localhost:8000/api/read-token-info
+  response:
+    body:
+      string: '{"token_id":"b630bd4e-6a20-486b-95cb-c02a175767b2","organization_id":"710c0de0-dda1-40fb-9601-b72d120acb71","project_id":"11d5d03b-7fc6-4ad6-bcd9-d1cc4bc87986","organization_name":"test-org","project_name":"starter-project"}'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '221'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 06 Jul 2025 12:58:04 GMT
+      Server:
+      - nginx/1.29.0
+      access-control-expose-headers:
+      - traceresponse
+      traceresponse:
+      - 00-0197dfd0ab965a5d8713381b433ff674-77253b890c5a0f74-01
+      x-api-version:
+      - UClDZTsUC0iVNKhEvlUL2JNm0ur0dBUm1u2/5He9ZtE=
+    status:
+      code: 200
+      message: OK
+version: 1

tests/aaa_query_client/test_query_client.py

@@ -12,7 +12,7 @@
 # To update, set the `CLIENT_BASE_URL` and `CLIENT_READ_TOKEN` values to match the local development environment,
 # and run the tests with `--record-mode=rewrite --inline-snapshot=fix` to update the cassettes and snapshots.
 CLIENT_BASE_URL = 'http://localhost:8000/'
-CLIENT_READ_TOKEN = '06KJCLLch8TCYx1FX4N1VGbr2mHrR760Z87zWjpb0TPm'
+CLIENT_READ_TOKEN = 'pylf_v1_local_wk3Vg7NQP1BLtK62PTB0sRqFmn3ThjqvbQn5R27MDZpd'
 pytestmark = [
     pytest.mark.vcr(),
     pytest.mark.skipif(
@@ -39,6 +39,45 @@ def test_infers_base_url_from_token(
     assert client.base_url == expected
 
 
+def test_info_sync():
+    with LogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
+        info = client.info()
+        assert info == snapshot(
+            {
+                'organization_name': 'test-org',
+                'project_name': 'starter-project',
+            }
+        )
+
+
+def test_info_invalid_schema_sync():
+    with pytest.raises(
+        RuntimeError, match='The read token info response is missing required fields: organization_name or project_name'
+    ):
+        with LogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
+            client.info()
+
+
+@pytest.mark.anyio
+async def test_info_async():
+    async with AsyncLogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
+        info = await client.info()
+        assert info == snapshot(
+            {
+                'organization_name': 'test-org',
+                'project_name': 'starter-project',
+            }
+        )
+
+
+async def test_info_invalid_schema_async():
+    with pytest.raises(
+        RuntimeError, match='The read token info response is missing required fields: organization_name or project_name'
+    ):
+        async with AsyncLogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
+            await client.info()
+
+
 def test_read_sync():
     with LogfireQueryClient(read_token=CLIENT_READ_TOKEN, base_url=CLIENT_BASE_URL) as client:
         sql = """