add http client connection pooling for efficient connection reuse

Author: purificant

src/bubble_data_api_client/__init__.py

@@ -1,3 +1,4 @@
 from .config import configure, set_config_provider
+from .pool import client_scope, close_clients
 
-__all__ = ["configure", "set_config_provider"]
+__all__ = ["configure", "set_config_provider", "client_scope", "close_clients"]

src/bubble_data_api_client/client/orm.py

@@ -6,19 +6,11 @@
 from pydantic import Field
 
 from bubble_data_api_client.client.raw_client import RawClient
-from bubble_data_api_client.config import get_config
 from bubble_data_api_client.constraints import ConstraintTypes, constraint
 
 
 def _get_client() -> RawClient:
-    config = get_config()
-    api_root = config.get("data_api_root_url")
-    api_key = config.get("api_key")
-    if not api_root:
-        raise RuntimeError("data_api_root_url")
-    if not api_key:
-        raise RuntimeError("api_key")
-    return RawClient(data_api_root_url=api_root, api_key=api_key)
+    return RawClient()
 
 
 class BubbleBaseModel(PydanticBaseModel):

src/bubble_data_api_client/client/raw_client.py

@@ -24,23 +24,13 @@ class RawClient:
     https://www.postman.com/bubbleapi/bubble/request/jigyk5v/
     """
 
-    _data_api_root_url: str
-    _api_key: str
     _transport: Transport
 
-    def __init__(
-        self,
-        data_api_root_url: str,
-        api_key: str,
-    ):
-        self._data_api_root_url = data_api_root_url
-        self._api_key = api_key
+    def __init__(self) -> None:
+        pass
 
     async def __aenter__(self) -> typing.Self:
-        self._transport = Transport(
-            base_url=self._data_api_root_url,
-            api_key=self._api_key,
-        )
+        self._transport = Transport()
         await self._transport.__aenter__()
         return self
 

src/bubble_data_api_client/pool.py

@@ -0,0 +1,99 @@
+"""Client pool for efficient connection reuse."""
+
+import asyncio
+import atexit
+import threading
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import httpx
+
+from .config import get_config
+from .transport import httpx_client_factory
+
+# global client pool keyed by config
+_clients: dict[tuple[str, str], httpx.AsyncClient] = {}
+_lock = threading.Lock()
+
+
+def _make_client_key(config: dict[str, Any]) -> tuple[str, str]:
+    """Generate a unique key for client pooling based on config."""
+    base_url = config.get("data_api_root_url") or ""
+    api_key = config.get("api_key") or ""
+    return (base_url, api_key)
+
+
+def get_client() -> httpx.AsyncClient:
+    """Get or create a client for the current config. Thread-safe."""
+    config = get_config()
+    key = _make_client_key(config)
+
+    # fast path: no lock if client exists
+    if key in _clients:
+        return _clients[key]
+
+    # slow path: acquire lock for creation
+    with _lock:
+        # double-check after acquiring lock
+        if key not in _clients:
+            base_url = config.get("data_api_root_url")
+            if not base_url:
+                raise RuntimeError("data_api_root_url")
+            api_key = config.get("api_key")
+            if not api_key:
+                raise RuntimeError("api_key")
+            _clients[key] = httpx_client_factory(base_url=base_url, api_key=api_key)
+        return _clients[key]
+
+
+async def close_clients() -> None:
+    """Close all clients in the pool. Thread-safe. Safe to call multiple times."""
+    with _lock:
+        clients_to_close = list(_clients.values())
+        _clients.clear()
+
+    for client in clients_to_close:
+        await client.aclose()
+
+
+@asynccontextmanager
+async def client_scope() -> AsyncIterator[None]:
+    """Scope that ensures close_clients() is called on exit."""
+    try:
+        yield
+    finally:
+        await close_clients()
+
+
+def _atexit_cleanup() -> None:
+    """Best-effort cleanup at interpreter exit."""
+    with _lock:
+        clients_to_close = list(_clients.values())
+        _clients.clear()
+
+    if not clients_to_close:
+        return
+
+    # check if there's already a running loop
+    try:
+        running_loop = asyncio.get_running_loop()
+    except RuntimeError:
+        running_loop = None
+
+    try:
+        if running_loop is not None:
+            # loop still running at atexit, schedule cleanup tasks
+            for client in clients_to_close:
+                running_loop.create_task(client.aclose())
+        else:
+            # no running loop, create one for cleanup
+            loop = asyncio.new_event_loop()
+            for client in clients_to_close:
+                loop.run_until_complete(client.aclose())
+            loop.close()
+    except Exception:
+        pass
+
+
+atexit.register(_atexit_cleanup)

src/bubble_data_api_client/transport.py

@@ -24,28 +24,21 @@ def httpx_client_factory(
 class Transport:
     """
     Transport layer focuses on HTTP.
-    - manage connections
-    - authentication
-    - headers
-    - retries, backoff
-    - timeouts
-    - exposes errors to the client
+    - authentication, headers, retries, timeouts: configured via httpx_client_factory
+    - connection lifecycle: managed by pool module
+    - HTTP verb methods: get, post, patch, put, delete
+    - error handling: raise_for_status on responses
     """
 
-    _base_url: str
-    _api_key: str
     _http: httpx.AsyncClient
 
-    def __init__(self, base_url: str, api_key: str):
-        self._base_url = base_url
-        self._api_key = api_key
+    def __init__(self) -> None:
+        pass
 
     async def __aenter__(self) -> typing.Self:
-        self._http = httpx_client_factory(
-            base_url=self._base_url,
-            api_key=self._api_key,
-        )
+        from .pool import get_client
 
+        self._http = get_client()
         return self
 
     async def __aexit__(
@@ -54,7 +47,7 @@ async def __aexit__(
         exc_val: BaseException | None,
         exc_tb: types.TracebackType | None,
     ) -> None:
-        await self._http.aclose()
+        pass
 
     async def request(
         self,