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:
Will Jones
2025-02-04 17:26:45 -08:00
committed by GitHub
parent c269524b2f
commit 16851389ea
7 changed files with 73 additions and 2 deletions

View File

@@ -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):

View File

@@ -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 == []

View File

@@ -223,6 +223,7 @@ pub struct PyClientConfig {
user_agent: String,
retry_config: Option<PyClientRetryConfig>,
timeout_config: Option<PyClientTimeoutConfig>,
extra_headers: Option<HashMap<String, String>>,
}
#[derive(FromPyObject)]
@@ -274,6 +275,7 @@ impl From<PyClientConfig> for lancedb::remote::ClientConfig {
user_agent: value.user_agent,
retry_config: value.retry_config.map(Into::into).unwrap_or_default(),
timeout_config: value.timeout_config.map(Into::into).unwrap_or_default(),
extra_headers: value.extra_headers.unwrap_or_default(),
}
}
}