Identify Python SDK requests

Amp-Thread-ID: https://ampcode.com/threads/T-01a03960-16bc-735f-afac-14505af47e7b

Author: snamber

tilebox-grpc/_tilebox/grpc/aio/channel.py

@@ -5,6 +5,7 @@
     CHANNEL_OPTIONS,
     ChannelInfo,
     ChannelProtocol,
+    _client_metadata,
     add_metadata,
     parse_channel_info,
     update_method,
@@ -34,7 +35,7 @@ def open_channel(url: str, auth_token: str | None = None, rpc_method_prefix: str
         A gRPC channel.
     """
     channel_info = parse_channel_info(url)
-    interceptors: list[ClientInterceptor] = []
+    interceptors: list[ClientInterceptor] = [_ClientMetadataInterceptor()]
     if auth_token is not None:
         interceptors = [_AuthMetadataInterceptor(auth_token), *interceptors]  # add auth interceptor as the first one
     if rpc_method_prefix is not None:
@@ -90,6 +91,20 @@ async def intercept_unary_unary(
         return await continuation(add_metadata(client_call_details, [self._auth]), request)
 
 
+class _ClientMetadataInterceptor(UnaryUnaryClientInterceptor):
+    def __init__(self) -> None:
+        super().__init__()
+        self._metadata = list(_client_metadata().items())
+
+    async def intercept_unary_unary(
+        self,
+        continuation: Callable[[ClientCallDetails, RequestType], UnaryUnaryCall],
+        client_call_details: ClientCallDetails,
+        request: RequestType,
+    ) -> UnaryUnaryCall:
+        return await continuation(add_metadata(client_call_details, self._metadata), request)
+
+
 class _RpcMethodPrefixInterceptor(UnaryUnaryClientInterceptor):
     def __init__(self, prefix: str) -> None:
         """A gRPC channel interceptor which prefixes every outgoing RPC method path."""

tilebox-grpc/_tilebox/grpc/channel.py

@@ -3,6 +3,8 @@
 from collections.abc import Awaitable, Callable, Mapping
 from dataclasses import dataclass
 from enum import Enum
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as package_version
 from typing import Literal, TypeVar
 
 from _tilebox.grpc.error import async_wrap_connect_rpc, wrap_connect_rpc
@@ -46,6 +48,9 @@
     ("grpc.service_config", json.dumps(_SERVICE_CONFIG)),
 ]
 
+CLIENT_SOURCE_HEADER = "tilebox-client-source"
+CLIENT_VERSION_HEADER = "tilebox-client-version"
+
 
 class ChannelProtocol(Enum):
     HTTPS = 1
@@ -77,7 +82,7 @@ def open_channel(url: str, auth_token: str | None = None, rpc_method_prefix: str
         A sync gRPC channel.
     """
     channel_info = parse_channel_info(url)
-    interceptors: list[UnaryUnaryClientInterceptor] = []
+    interceptors: list[UnaryUnaryClientInterceptor] = [_ClientMetadataInterceptor()]
     if auth_token is not None:
         interceptors = [_AuthMetadataInterceptor(auth_token), *interceptors]  # add auth interceptor as the first one
     if rpc_method_prefix is not None:
@@ -185,7 +190,7 @@ def __init__(
         method_path_prefix = _rpc_method_prefix_path(rpc_method_prefix)
         service_name = _connect_service_name(client)
         self._client = client
-        self._headers = headers
+        self._headers = {**(headers or {}), **_client_metadata()}
 
         for connect_name in _connect_client_methods(client):
             grpc_name = _snake_to_pascal_case(connect_name)
@@ -211,7 +216,7 @@ def __init__(
         method_path_prefix = _rpc_method_prefix_path(rpc_method_prefix)
         service_name = _connect_service_name(client)
         self._client = client
-        self._headers = headers
+        self._headers = {**(headers or {}), **_client_metadata()}
 
         for connect_name in _connect_client_methods(client):
             grpc_name = _snake_to_pascal_case(connect_name)
@@ -278,6 +283,20 @@ def intercept_unary_unary(
         return continuation(add_metadata(client_call_details, [self._auth]), request)
 
 
+class _ClientMetadataInterceptor(UnaryUnaryClientInterceptor):
+    def __init__(self) -> None:
+        super().__init__()
+        self._metadata = list(_client_metadata().items())
+
+    def intercept_unary_unary(
+        self,
+        continuation: Callable[[ClientCallDetails, RequestType], ResponseType],
+        client_call_details: ClientCallDetails,
+        request: RequestType,
+    ) -> ResponseType:
+        return continuation(add_metadata(client_call_details, self._metadata), request)
+
+
 class _RpcMethodPrefixInterceptor(UnaryUnaryClientInterceptor):
     def __init__(self, prefix: str) -> None:
         """A sync gRPC channel interceptor which prefixes every outgoing RPC method path."""
@@ -301,6 +320,17 @@ def add_metadata(
     return _replace_call_details(client_call_details, metadata=metadata)
 
 
+def _client_metadata() -> dict[str, str]:
+    try:
+        client_version = package_version("tilebox-grpc")
+    except PackageNotFoundError:
+        client_version = "dev"
+    return {
+        CLIENT_SOURCE_HEADER: "python_sdk",
+        CLIENT_VERSION_HEADER: client_version,
+    }
+
+
 def update_method(client_call_details: ClientCallDetails, prefix: str) -> ClientCallDetails:
     return _replace_call_details(client_call_details, method=prefix_rpc_method(client_call_details.method, prefix))
 

tilebox-grpc/_tilebox/grpc/replay.py

@@ -8,6 +8,7 @@
 
 from _tilebox.grpc.channel import (
     _AuthMetadataInterceptor,
+    _ClientMetadataInterceptor,
     _open_channel,
     parse_channel_info,
 )
@@ -32,7 +33,7 @@
 def open_recording_channel(url: str, auth_token: str | None, recording: str | Path) -> Channel:
     """Open a gRPC channel to the given URL and record all requests and responses to a file."""
     channel_info = parse_channel_info(url)
-    interceptors: list[UnaryUnaryClientInterceptor] = [_RecordRPCsInterceptor(recording)]
+    interceptors: list[UnaryUnaryClientInterceptor] = [_ClientMetadataInterceptor(), _RecordRPCsInterceptor(recording)]
     if auth_token is not None:
         interceptors = [_AuthMetadataInterceptor(auth_token), *interceptors]  # add auth interceptor as the first one
 

tilebox-grpc/tests/aio/test_channel.py

@@ -2,7 +2,13 @@
 
 import pytest
 
-from _tilebox.grpc.aio.channel import ClientCallDetails, _AuthMetadataInterceptor, _RpcMethodPrefixInterceptor
+from _tilebox.grpc.aio.channel import (
+    ClientCallDetails,
+    _AuthMetadataInterceptor,
+    _ClientMetadataInterceptor,
+    _RpcMethodPrefixInterceptor,
+)
+from _tilebox.grpc.channel import CLIENT_SOURCE_HEADER, CLIENT_VERSION_HEADER
 
 
 @pytest.mark.asyncio
@@ -22,6 +28,23 @@ async def test_auth_interceptor(req_metadata: list[tuple[str, str]] | None) -> N
     assert ("authorization", "Bearer very-secret") in updated_call_details.metadata
 
 
+@pytest.mark.asyncio
+async def test_client_metadata_interceptor() -> None:
+    interceptor = _ClientMetadataInterceptor()
+    mock_method = AsyncMock()
+
+    await interceptor.intercept_unary_unary(
+        mock_method,
+        ClientCallDetails("/some-rpc-method", 10, [("authorization", "Bearer token")], None, True),
+        AsyncMock(),
+    )
+
+    metadata = mock_method.call_args[0][0].metadata
+    assert ("authorization", "Bearer token") in metadata
+    assert (CLIENT_SOURCE_HEADER, "python_sdk") in metadata
+    assert any(key == CLIENT_VERSION_HEADER and value for key, value in metadata)
+
+
 @pytest.mark.asyncio
 async def test_rpc_method_prefix_interceptor() -> None:
     interceptor = _RpcMethodPrefixInterceptor("/public")

tilebox-grpc/tests/test_channel.py

@@ -1,11 +1,17 @@
+from importlib.metadata import PackageNotFoundError
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from _tilebox.grpc.channel import (
     CHANNEL_OPTIONS,
+    CLIENT_SOURCE_HEADER,
+    CLIENT_VERSION_HEADER,
+    AsyncConnectStubAdapter,
     ChannelProtocol,
     ClientCallDetails,
+    ConnectStubAdapter,
+    _ClientMetadataInterceptor,
     _RpcMethodPrefixInterceptor,
     connect_address,
     open_channel,
@@ -39,14 +45,80 @@ def test_open_authenticated_channel(open_func: MagicMock, intercept_func: MagicM
     assert intercept_func.call_args[0][1]._auth == ("authorization", "Bearer very-secret")
 
 
+@patch("_tilebox.grpc.channel.package_version", return_value="1.2.3")
+def test_client_metadata_interceptor(package_version: MagicMock) -> None:
+    interceptor = _ClientMetadataInterceptor()
+    continuation = MagicMock()
+
+    interceptor.intercept_unary_unary(
+        continuation,
+        ClientCallDetails("/some-rpc-method", 10, [("authorization", "Bearer token")], None, True),
+        MagicMock(),
+    )
+
+    package_version.assert_called_once_with("tilebox-grpc")
+    metadata = continuation.call_args[0][0].metadata
+    assert ("authorization", "Bearer token") in metadata
+    assert (CLIENT_SOURCE_HEADER, "python_sdk") in metadata
+    assert (CLIENT_VERSION_HEADER, "1.2.3") in metadata
+
+
+def test_client_metadata_uses_dev_version_when_package_is_not_installed() -> None:
+    with patch("_tilebox.grpc.channel.package_version", side_effect=PackageNotFoundError):
+        interceptor = _ClientMetadataInterceptor()
+
+    assert (CLIENT_VERSION_HEADER, "dev") in interceptor._metadata
+
+
+class _ConnectClient:
+    def get_value(self, request: str, *, headers: dict[str, str]) -> tuple[str, dict[str, str]]:
+        return request, headers
+
+
+class _AsyncConnectClient:
+    async def get_value(self, request: str, *, headers: dict[str, str]) -> tuple[str, dict[str, str]]:
+        return request, headers
+
+
+@patch("_tilebox.grpc.channel.package_version", return_value="1.2.3")
+def test_connect_stub_adapter_adds_client_metadata(package_version: MagicMock) -> None:
+    adapter = ConnectStubAdapter(_ConnectClient(), {"authorization": "Bearer token"})
+
+    request, headers = adapter.GetValue("request")  # ty: ignore[unresolved-attribute]  # added dynamically
+
+    assert request == "request"
+    assert headers == {
+        "authorization": "Bearer token",
+        CLIENT_SOURCE_HEADER: "python_sdk",
+        CLIENT_VERSION_HEADER: "1.2.3",
+    }
+    package_version.assert_called_once_with("tilebox-grpc")
+
+
+@pytest.mark.asyncio
+@patch("_tilebox.grpc.channel.package_version", return_value="1.2.3")
+async def test_async_connect_stub_adapter_adds_client_metadata(package_version: MagicMock) -> None:
+    adapter = AsyncConnectStubAdapter(_AsyncConnectClient(), {"authorization": "Bearer token"})
+
+    request, headers = await adapter.GetValue("request")  # ty: ignore[unresolved-attribute]  # added dynamically
+
+    assert request == "request"
+    assert headers == {
+        "authorization": "Bearer token",
+        CLIENT_SOURCE_HEADER: "python_sdk",
+        CLIENT_VERSION_HEADER: "1.2.3",
+    }
+    package_version.assert_called_once_with("tilebox-grpc")
+
+
 @patch("_tilebox.grpc.channel.intercept_channel")
 @patch("_tilebox.grpc.channel.secure_channel")
 def test_open_channel_with_rpc_method_prefix(open_func: MagicMock, intercept_func: MagicMock) -> None:
     open_channel("api.tilebox.com", rpc_method_prefix="/public")
     open_func.assert_called_once()
     intercept_func.assert_called_once()
 
-    assert intercept_func.call_args[0][1]._prefix == "/public"
+    assert intercept_func.call_args[0][2]._prefix == "/public"
 
 
 def test_rpc_method_prefix_interceptor() -> None: