mirror of
https://github.com/lancedb/lancedb.git
synced 2026-01-11 06:12:58 +00:00
feat: extra headers parameter in client options (#2091)
Closes #1106 Unfortunately, these need to be set at the connection level. I investigated whether if we let users provide a callback they could use `AsyncLocalStorage` to access their context. However, it doesn't seem like NAPI supports this right now. I filed an issue: https://github.com/napi-rs/napi-rs/issues/2456
This commit is contained in:
@@ -109,6 +109,7 @@ class ClientConfig:
|
||||
user_agent: str = f"LanceDB-Python-Client/{__version__}"
|
||||
retry_config: RetryConfig = field(default_factory=RetryConfig)
|
||||
timeout_config: Optional[TimeoutConfig] = field(default_factory=TimeoutConfig)
|
||||
extra_headers: Optional[dict] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if isinstance(self.retry_config, dict):
|
||||
|
||||
@@ -57,7 +57,7 @@ def mock_lancedb_connection(handler):
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def mock_lancedb_connection_async(handler):
|
||||
async def mock_lancedb_connection_async(handler, **client_config):
|
||||
with http.server.HTTPServer(
|
||||
("localhost", 8080), make_mock_http_handler(handler)
|
||||
) as server:
|
||||
@@ -73,6 +73,7 @@ async def mock_lancedb_connection_async(handler):
|
||||
"timeout_config": {
|
||||
"connect_timeout": 1,
|
||||
},
|
||||
**client_config,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -522,3 +523,19 @@ def test_create_client():
|
||||
|
||||
with pytest.warns(DeprecationWarning):
|
||||
lancedb.connect(**mandatory_args, request_thread_pool=10)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pass_through_headers():
|
||||
def handler(request):
|
||||
assert request.headers["foo"] == "bar"
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"tables": []}')
|
||||
|
||||
async with mock_lancedb_connection_async(
|
||||
handler, extra_headers={"foo": "bar"}
|
||||
) as db:
|
||||
table_names = await db.table_names()
|
||||
assert table_names == []
|
||||
|
||||
Reference in New Issue
Block a user