feat(python): bind python async remote client to rust client (#1700)

Closes [#1638](https://github.com/lancedb/lancedb/issues/1638)

This just binds the Python Async client to the Rust remote client.
This commit is contained in:
Will Jones
2024-10-01 15:46:59 -07:00
committed by GitHub
parent a416925ca1
commit f305f34d9b
6 changed files with 255 additions and 6 deletions

View File

@@ -1,11 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
import http.server
import threading
from unittest.mock import MagicMock
import uuid
import lancedb
import pyarrow as pa
from lancedb.remote.client import VectorQuery, VectorQueryResult
import pytest
class FakeLanceDBClient:
@@ -81,3 +85,57 @@ def test_create_table_with_recordbatches():
table = conn.create_table("test", [batch], schema=batch.schema)
assert table.name == "test"
assert client.post.call_args[0][0] == "/v1/table/test/create/"
def make_mock_http_handler(handler):
class MockLanceDBHandler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
handler(self)
def do_POST(self):
handler(self)
return MockLanceDBHandler
@pytest.mark.asyncio
async def test_async_remote_db():
def handler(request):
# We created a UUID request id
request_id = request.headers["x-request-id"]
assert uuid.UUID(request_id).version == 4
# We set a user agent with the current library version
user_agent = request.headers["User-Agent"]
assert user_agent == f"LanceDB-Python-Client/{lancedb.__version__}"
request.send_response(200)
request.send_header("Content-Type", "application/json")
request.end_headers()
request.wfile.write(b'{"tables": []}')
def run_server():
with http.server.HTTPServer(
("localhost", 8080), make_mock_http_handler(handler)
) as server:
# we will only make one request
server.handle_request()
handle = threading.Thread(target=run_server)
handle.start()
db = await lancedb.connect_async(
"db://dev",
api_key="fake",
host_override="http://localhost:8080",
client_config={
"retry_config": {"retries": 2},
"timeout_config": {
"connect_timeout": 1,
},
},
)
table_names = await db.table_names()
assert table_names == []
handle.join()