mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-09 23:02:37 +00:00
Jack Ye
e639b1b650
## Summary
Add SQL execution to remote LanceDB connections. On the standard
synchronous connection, `execute_query` waits for the initial result
stream and returns its Arrow reader. `execute_query_async` is called
without Python `await` and immediately returns a query handle for status
inspection, streaming, or cancellation. Local databases report that SQL
is not supported.
The transport and query lifecycle live in Rust. Python exposes
native-backed synchronous and asynchronous connection methods and query
wrappers; it does not use PyArrow's Flight client.
## User experience
The standard synchronous connection supports both direct reads and
background query execution:
```python
db = lancedb.connect(
"db://analytics",
api_key="ldb_...",
sql_host_override="grpc+tls://sql.example.com:10026",
)
# Direct execution waits only until the initial result stream is available.
# Later batches continue streaming as the query progresses.
reader = db.execute_query(
"SELECT * FROM events",
default_namespace_path=["production"],
)
for batch in reader:
print(batch.num_rows)
# Background execution returns a query handle immediately. Despite the
# `_async` suffix, no Python `await` is needed on a synchronous connection.
query = db.execute_query_async("SELECT * FROM events")
print(query.id)
description = db.describe_query(query.id)
print(description.status)
print(description.progress)
print(description.expires_at)
# Start reading as soon as the service advertises partial results. The reader
# continues polling and yields newly available record batches until the query
# and all result endpoints are complete.
reader = query.reader()
for batch in reader:
print(batch.num_rows)
# Or cancel a different still-running query. Its status becomes "cancelling"
# while the server is still working, then "cancelled" once confirmed.
cancelled_query = db.execute_query_async("SELECT * FROM large_events")
cancelled_query.cancel()
```
The less commonly used asynchronous connection exposes the same
operations as coroutines:
```python
async_db = await lancedb.connect_async(
"db://analytics",
api_key="ldb_...",
sql_host_override="grpc+tls://sql.example.com:10026",
)
query = await async_db.execute_query_async("SELECT * FROM events")
async for batch in await query.reader():
print(batch.num_rows)
```
The UUIDv7 query id is scoped to the connection that submitted it. The
connection retains lightweight shared query state used by
`query.describe()` and `db.describe_query(query.id)`; the id does not
encode SQL or a Flight continuation token and is not a cross-connection
resume token. Abandoned state has bounded retention, and terminal state
remains available briefly.
Unqualified table names use the connected database and the `public`
namespace by default. `default_namespace_path` accepts a list such as
`["production", "events"]`. SQL can still use qualified names to
reference other databases and namespaces available to the deployment.
## Design
- Uses Arrow Flight `PollFlightInfo` for submission and long polling,
`DoGet` for results, and `CancelFlightInfo` for cancellation. Each
`PollInfo.info` is treated as the cumulative set of currently available
endpoints, so advertised tickets are consumed once and batches can be
delivered before execution is complete.
- Serializes result completion and cancellation into one lifecycle. A
server-accepted request reports `cancelling` and wakes blocked
status/result work; a later retry can confirm `cancelled`. Result
retrieval is rejected after cancellation is accepted, while cancellation
after a result was already delivered is a no-op.
- Assigns a time-ordered UUIDv7 connection-scoped query id and retains
only shared evolving lifecycle state, keeping SQL, Flight continuation
tokens, and Arrow result data out of public ids and the registry.
- Leaves admission control to the server while honoring server
expiration and a local fallback retention window for abandoned entries.
- Retains terminal ids for five minutes so they remain available for
connection-level description.
- Keeps one lazily initialized SQL client on each remote database
connection and attaches fresh authentication, routing, namespace, and
request metadata to every operation.
- Applies the configured overall timeout to each execution, description,
reader, and cancellation operation. A result reader carries one absolute
deadline from `reader()` through the end of streaming; connect and read
timeouts continue to bound their individual phases.
- Returns a bounded, backpressured, single-consumer Arrow stream rather
than collecting the full result in memory. Dropping the reader stops
downloading but does not implicitly cancel the server query.
- Preserves typed schemas for empty result sets through the stream
schema.
- Accepts Flight result messages up to 1 GiB so a valid row containing a
large blob, string, or vector is not rejected by tonic's 4 MiB default
receive limit.
- Supports the Python client first while keeping the authoritative
implementation in the Rust core.
163 lines
5.5 KiB
Python
163 lines
5.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
from uuid import UUID
|
|
|
|
import pytest
|
|
import pyarrow as pa
|
|
|
|
import lancedb
|
|
from lancedb import _lancedb
|
|
from lancedb.arrow import AsyncRecordBatchReader
|
|
from lancedb.db import AsyncConnection
|
|
from lancedb.remote.db import RemoteDBConnection
|
|
from lancedb.sql import AsyncQuery, Query
|
|
|
|
NIL_QUERY_ID = UUID(int=0)
|
|
|
|
|
|
class FakeNativeQuery:
|
|
id = UUID("0198f1b2-c3d4-7e5f-8123-456789abcdef")
|
|
|
|
async def reader(self):
|
|
return pa.table({"value": [1, 2]})
|
|
|
|
|
|
class FakeNativeConnection:
|
|
async def execute_query_async(self, query, *, default_namespace_path=None):
|
|
return FakeNativeQuery()
|
|
|
|
|
|
class FakeAsyncConnection:
|
|
async def execute_query_async(self, query, *, default_namespace_path=None):
|
|
return AsyncQuery(FakeNativeQuery())
|
|
|
|
|
|
def remote_connection(sql_host_override=None):
|
|
return lancedb.connect(
|
|
"db://analytics",
|
|
api_key="test-key",
|
|
host_override="http://localhost:10024",
|
|
sql_host_override=sql_host_override,
|
|
)
|
|
|
|
|
|
def test_sql_is_connection_scoped():
|
|
assert hasattr(lancedb, "sql")
|
|
assert not callable(lancedb.sql)
|
|
assert not hasattr(_lancedb, "sql")
|
|
assert not hasattr(remote_connection(), "sql")
|
|
assert hasattr(remote_connection(), "execute_query")
|
|
assert hasattr(remote_connection(), "execute_query_async")
|
|
assert hasattr(remote_connection(), "describe_query")
|
|
|
|
|
|
def test_query_id_is_uuid():
|
|
query = AsyncQuery(FakeNativeQuery())
|
|
assert isinstance(query.id, UUID)
|
|
assert Query(query).id == query.id
|
|
|
|
|
|
def test_connection_serializes_sql_host_override():
|
|
endpoint = "grpc+tls://sql.example.com:10026"
|
|
restored = lancedb.deserialize_conn(
|
|
remote_connection(sql_host_override=endpoint).serialize()
|
|
)
|
|
assert restored.sql_host_override == endpoint
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_sql_reader_is_record_batch_stream():
|
|
reader = await AsyncQuery(FakeNativeQuery()).reader()
|
|
assert isinstance(reader, AsyncRecordBatchReader)
|
|
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
|
|
|
|
|
|
def test_sync_sql_reader_is_record_batch_reader():
|
|
reader = Query(AsyncQuery(FakeNativeQuery())).reader()
|
|
assert isinstance(reader, pa.RecordBatchReader)
|
|
assert reader.read_all().column(0).to_pylist() == [1, 2]
|
|
|
|
|
|
def test_execute_query_returns_blocking_reader():
|
|
connection = RemoteDBConnection.__new__(RemoteDBConnection)
|
|
connection._conn = FakeAsyncConnection()
|
|
reader = connection.execute_query("SELECT 1")
|
|
assert isinstance(reader, pa.RecordBatchReader)
|
|
assert reader.read_all().column(0).to_pylist() == [1, 2]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_execute_query_returns_async_reader():
|
|
connection = AsyncConnection(FakeNativeConnection())
|
|
reader = await connection.execute_query("SELECT 1")
|
|
assert isinstance(reader, AsyncRecordBatchReader)
|
|
assert (await reader.read_all())[0].column(0).to_pylist() == [1, 2]
|
|
|
|
|
|
def test_local_connection_rejects_sql(tmp_path):
|
|
connection = lancedb.connect(tmp_path)
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
connection.execute_query("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
connection.execute_query_async("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
connection.describe_query(NIL_QUERY_ID)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_local_async_connection_rejects_sql(tmp_path):
|
|
connection = await lancedb.connect_async(tmp_path)
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.execute_query("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.execute_query_async("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.describe_query(NIL_QUERY_ID)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_namespace_connection_rejects_sql(tmp_path):
|
|
connection = lancedb.connect_namespace_async("dir", {"root": str(tmp_path)})
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.execute_query("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.execute_query_async("SELECT 1")
|
|
with pytest.raises(NotImplementedError, match="SQL"):
|
|
await connection.describe_query(NIL_QUERY_ID)
|
|
|
|
|
|
def test_describe_query_requires_uuid():
|
|
with pytest.raises(TypeError, match="UUID"):
|
|
remote_connection().describe_query(str(NIL_QUERY_ID))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"default_namespace_path",
|
|
["public", ("public",), [1]],
|
|
)
|
|
def test_execute_query_async_requires_namespace_path_list(default_namespace_path):
|
|
with pytest.raises(ValueError, match="default_namespace_path"):
|
|
remote_connection().execute_query_async(
|
|
"SELECT 1", default_namespace_path=default_namespace_path
|
|
)
|
|
|
|
|
|
def test_execute_query_async_rejects_invalid_endpoint():
|
|
connection = remote_connection(sql_host_override="invalid://localhost")
|
|
with pytest.raises(ValueError, match="sql_host_override"):
|
|
connection.execute_query_async("SELECT 1")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"default_namespace_path",
|
|
[[""], ["café"], ["pub\tlic"], ["events$raw"]],
|
|
)
|
|
def test_execute_query_async_rejects_invalid_namespace_components(
|
|
default_namespace_path,
|
|
):
|
|
with pytest.raises(ValueError, match="default_namespace_path"):
|
|
remote_connection().execute_query_async(
|
|
"SELECT 1", default_namespace_path=default_namespace_path
|
|
)
|