mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(sdk): expose table overrides for remote tables
Expose table-level maintenance override read, update, and reset APIs for remote
Rust and Python tables. Local tables continue to return `NotSupported`.
Updates recursively merge the caller-provided JSON object into the current
override document so unrelated and unknown service fields are preserved. Reset
uses the explicit reset endpoint.
The integrated change keeps current main's remote database identity, namespace,
branch, SQL, WAL, and threaded mock-server behavior.
Validated with focused Rust and Python override tests and the complete applicable
CI run for 43bccbadc5.
This commit is contained in:
@@ -593,6 +593,9 @@ class Table:
|
||||
async def enable_rte(self, spec: RTEWriteSpec) -> None: ...
|
||||
async def disable_rte(self) -> None: ...
|
||||
async def get_rte_spec(self) -> Optional[RTEWriteSpec]: ...
|
||||
async def get_table_overrides(self) -> str: ...
|
||||
async def update_table_overrides(self, overrides_json: str) -> str: ...
|
||||
async def reset_table_overrides(self) -> str: ...
|
||||
# WAL-PK-FUSION: delete both.
|
||||
def hybrid_pk_fusion_learned(self) -> bool: ...
|
||||
def note_hybrid_pk_fusion(self) -> None: ...
|
||||
|
||||
@@ -1153,6 +1153,18 @@ class RemoteTable(Table):
|
||||
"""Read the installed RTEWriteSpec, or ``None``."""
|
||||
return LOOP.run(self._table.get_rte_spec())
|
||||
|
||||
def get_table_overrides(self) -> dict[str, Any]:
|
||||
"""Read Enterprise table overrides for this table."""
|
||||
return LOOP.run(self._table.get_table_overrides())
|
||||
|
||||
def update_table_overrides(self, **overrides: Any) -> dict[str, Any]:
|
||||
"""Merge Enterprise table overrides into the current table overrides."""
|
||||
return LOOP.run(self._table.update_table_overrides(**overrides))
|
||||
|
||||
def reset_table_overrides(self) -> dict[str, Any]:
|
||||
"""Clear Enterprise table overrides for this table."""
|
||||
return LOOP.run(self._table.reset_table_overrides())
|
||||
|
||||
# WAL-PK-FUSION: delete both hooks.
|
||||
def _hybrid_pk_fusion_learned(self) -> bool:
|
||||
# `self._table` is the async table, which owns the Rust handle.
|
||||
|
||||
@@ -1931,6 +1931,27 @@ class Table(ABC):
|
||||
"create_index_async is not supported for this table type"
|
||||
)
|
||||
|
||||
def get_table_overrides(self) -> dict[str, Any]:
|
||||
"""Read Enterprise table overrides for this table.
|
||||
|
||||
Table overrides configure Enterprise-managed background behavior such as
|
||||
automatic compaction, cleanup, and reindex execution. This method is
|
||||
supported only on remote Enterprise tables.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def update_table_overrides(self, **overrides: Any) -> dict[str, Any]:
|
||||
"""Merge Enterprise table overrides into the current table overrides.
|
||||
|
||||
Only fields supplied in ``overrides`` are changed; unrelated existing
|
||||
overrides are preserved by the remote service update flow.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def reset_table_overrides(self) -> dict[str, Any]:
|
||||
"""Clear Enterprise table overrides for this table."""
|
||||
raise NotImplementedError
|
||||
|
||||
def drop_index(self, name: str) -> None:
|
||||
"""
|
||||
Drop an index from the table.
|
||||
@@ -5193,6 +5214,36 @@ class LanceTable(Table):
|
||||
[`AsyncTable.get_rte_spec`][lancedb.AsyncTable.get_rte_spec]."""
|
||||
return LOOP.run(self._table.get_rte_spec())
|
||||
|
||||
def get_table_overrides(self) -> dict[str, Any]:
|
||||
"""Read Enterprise table overrides for this table.
|
||||
|
||||
Table overrides configure Enterprise-managed background behavior such as
|
||||
automatic compaction, cleanup, and reindex execution. This method is
|
||||
supported only on remote Enterprise tables.
|
||||
"""
|
||||
return LOOP.run(self._table.get_table_overrides())
|
||||
|
||||
def update_table_overrides(self, **overrides: Any) -> dict[str, Any]:
|
||||
"""Merge Enterprise table overrides into the current table overrides.
|
||||
|
||||
Only fields supplied in ``overrides`` are changed; unrelated overrides
|
||||
from the current read snapshot are preserved.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> table.update_table_overrides( # doctest: +SKIP
|
||||
... job_types={
|
||||
... "compaction": {"enabled": True},
|
||||
... "cleanup": {"enabled": False},
|
||||
... }
|
||||
... )
|
||||
"""
|
||||
return LOOP.run(self._table.update_table_overrides(**overrides))
|
||||
|
||||
def reset_table_overrides(self) -> dict[str, Any]:
|
||||
"""Clear Enterprise table overrides for this table."""
|
||||
return LOOP.run(self._table.reset_table_overrides())
|
||||
|
||||
def checkpoint_rte(self) -> None:
|
||||
"""Synchronous version of
|
||||
[`AsyncTable.checkpoint_rte`][lancedb.AsyncTable.checkpoint_rte]."""
|
||||
@@ -5942,6 +5993,37 @@ class AsyncTable:
|
||||
"""
|
||||
return await self._inner.get_rte_spec()
|
||||
|
||||
async def get_table_overrides(self) -> dict[str, Any]:
|
||||
"""Read Enterprise table overrides for this table.
|
||||
|
||||
Table overrides configure Enterprise-managed background behavior such as
|
||||
automatic compaction, cleanup, and reindex execution. This method is
|
||||
supported only on remote Enterprise tables.
|
||||
"""
|
||||
return json.loads(await self._inner.get_table_overrides())
|
||||
|
||||
async def update_table_overrides(self, **overrides: Any) -> dict[str, Any]:
|
||||
"""Merge Enterprise table overrides into the current table overrides.
|
||||
|
||||
Only fields supplied in ``overrides`` are changed; unrelated overrides
|
||||
from the current read snapshot are preserved.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> await table.update_table_overrides( # doctest: +SKIP
|
||||
... job_types={
|
||||
... "compaction": {"enabled": True},
|
||||
... "cleanup": {"enabled": False},
|
||||
... }
|
||||
... )
|
||||
"""
|
||||
encoded = json.dumps(overrides)
|
||||
return json.loads(await self._inner.update_table_overrides(encoded))
|
||||
|
||||
async def reset_table_overrides(self) -> dict[str, Any]:
|
||||
"""Clear Enterprise table overrides for this table."""
|
||||
return json.loads(await self._inner.reset_table_overrides())
|
||||
|
||||
# WAL-PK-FUSION: delete both hooks.
|
||||
def _hybrid_pk_fusion_learned(self) -> bool:
|
||||
"""See [`Table._hybrid_pk_fusion_learned`][lancedb.table.Table]."""
|
||||
|
||||
@@ -17,9 +17,9 @@ import sys
|
||||
import textwrap
|
||||
import threading
|
||||
import time
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from unittest.mock import MagicMock, patch
|
||||
import uuid
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
from packaging.version import Version
|
||||
|
||||
import lancedb
|
||||
@@ -55,6 +55,9 @@ def make_mock_http_handler(handler):
|
||||
def do_POST(self):
|
||||
handler(self)
|
||||
|
||||
def do_PUT(self):
|
||||
handler(self)
|
||||
|
||||
return MockLanceDBHandler
|
||||
|
||||
|
||||
@@ -2346,6 +2349,118 @@ def test_checkpoint_rte_sync():
|
||||
assert called == ["flush_rte", "get_rte_stats"]
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def table_overrides_test_table(overrides_handler):
|
||||
def handler(request):
|
||||
parsed = urlparse(request.path)
|
||||
if parsed.path in (
|
||||
"/admin/table/overrides",
|
||||
"/admin/table/overrides/reset",
|
||||
):
|
||||
overrides_handler(request, parsed)
|
||||
elif parsed.path == "/v1/table/test/describe/":
|
||||
request.send_response(200)
|
||||
request.send_header("Content-Type", "application/json")
|
||||
request.end_headers()
|
||||
request.wfile.write(b'{"version": 1, "schema": {"fields": []}}')
|
||||
else:
|
||||
request.send_response(404)
|
||||
request.end_headers()
|
||||
|
||||
with mock_lancedb_connection(handler) as db:
|
||||
yield db.open_table("test")
|
||||
|
||||
|
||||
def assert_table_overrides_query(parsed):
|
||||
query = parse_qs(parsed.query)
|
||||
assert query == {"db": ["dev"], "table": ["test"]}
|
||||
|
||||
|
||||
def test_remote_table_exposes_table_overrides_methods():
|
||||
from lancedb.remote.table import RemoteTable
|
||||
|
||||
for name in (
|
||||
"get_table_overrides",
|
||||
"update_table_overrides",
|
||||
"reset_table_overrides",
|
||||
):
|
||||
assert hasattr(RemoteTable, name)
|
||||
|
||||
|
||||
def test_get_table_overrides_sync_remote_table():
|
||||
def overrides_handler(request, parsed):
|
||||
assert request.command == "GET"
|
||||
assert parsed.path == "/admin/table/overrides"
|
||||
assert_table_overrides_query(parsed)
|
||||
send_json(request, {"job_types": {"cleanup": {"enabled": False}}})
|
||||
|
||||
with table_overrides_test_table(overrides_handler) as table:
|
||||
assert table.get_table_overrides() == {
|
||||
"job_types": {"cleanup": {"enabled": False}}
|
||||
}
|
||||
|
||||
|
||||
def test_update_table_overrides_sync_remote_table_merges_existing_fields():
|
||||
calls = []
|
||||
|
||||
def overrides_handler(request, parsed):
|
||||
calls.append(request.command)
|
||||
assert parsed.path == "/admin/table/overrides"
|
||||
if request.command == "GET":
|
||||
assert_table_overrides_query(parsed)
|
||||
send_json(
|
||||
request,
|
||||
{
|
||||
"job_types": {
|
||||
"compaction": {"enabled": True},
|
||||
"reindex_vector": {"table": {"num_workers": 3}},
|
||||
},
|
||||
"unknown": {"keep": True},
|
||||
},
|
||||
)
|
||||
elif request.command == "PUT":
|
||||
body = read_json_body(request)
|
||||
assert body["db"] == "dev"
|
||||
assert body["namespace"] == []
|
||||
assert body["table"] == "test"
|
||||
assert body["overrides"] == {
|
||||
"job_types": {
|
||||
"compaction": {"enabled": True},
|
||||
"cleanup": {"enabled": False},
|
||||
"reindex_vector": {"table": {"num_workers": 3}},
|
||||
},
|
||||
"unknown": {"keep": True},
|
||||
}
|
||||
send_json(request, body["overrides"])
|
||||
else:
|
||||
raise AssertionError(f"unexpected method: {request.command}")
|
||||
|
||||
with table_overrides_test_table(overrides_handler) as table:
|
||||
assert table.update_table_overrides(
|
||||
job_types={"cleanup": {"enabled": False}}
|
||||
) == {
|
||||
"job_types": {
|
||||
"compaction": {"enabled": True},
|
||||
"cleanup": {"enabled": False},
|
||||
"reindex_vector": {"table": {"num_workers": 3}},
|
||||
},
|
||||
"unknown": {"keep": True},
|
||||
}
|
||||
|
||||
assert calls == ["GET", "PUT"]
|
||||
|
||||
|
||||
def test_reset_table_overrides_sync_remote_table():
|
||||
def overrides_handler(request, parsed):
|
||||
assert request.command == "POST"
|
||||
assert parsed.path == "/admin/table/overrides/reset"
|
||||
assert_table_overrides_query(parsed)
|
||||
send_json(request, {})
|
||||
|
||||
with table_overrides_test_table(overrides_handler) as table:
|
||||
assert table.reset_table_overrides() == {}
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def query_test_table(
|
||||
query_handler,
|
||||
|
||||
@@ -1933,6 +1933,66 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection):
|
||||
)
|
||||
|
||||
|
||||
@patch("lancedb.table.AsyncTable.get_table_overrides")
|
||||
def test_get_table_overrides_method(mock_get_table_overrides, mem_db: DBConnection):
|
||||
expected = {"job_types": {"cleanup": {"enabled": False}}}
|
||||
mock_get_table_overrides.return_value = expected
|
||||
table = mem_db.create_table("test", data=[{"id": 1}])
|
||||
|
||||
assert table.get_table_overrides() == expected
|
||||
mock_get_table_overrides.assert_called_once_with()
|
||||
|
||||
|
||||
@patch("lancedb.table.AsyncTable.update_table_overrides")
|
||||
def test_update_table_overrides_method(
|
||||
mock_update_table_overrides, mem_db: DBConnection
|
||||
):
|
||||
overrides = {
|
||||
"job_types": {
|
||||
"cleanup": {"enabled": False},
|
||||
"index": {
|
||||
"scalar": {
|
||||
"reindex": {
|
||||
"row_threshold": 1000,
|
||||
"ttl_millis": 300000,
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
mock_update_table_overrides.return_value = overrides
|
||||
table = mem_db.create_table("test", data=[{"id": 1}])
|
||||
|
||||
assert table.update_table_overrides(**overrides) == overrides
|
||||
mock_update_table_overrides.assert_called_once_with(**overrides)
|
||||
|
||||
|
||||
@patch("lancedb.table.AsyncTable.reset_table_overrides")
|
||||
def test_reset_table_overrides_method(mock_reset_table_overrides, mem_db: DBConnection):
|
||||
expected = {}
|
||||
mock_reset_table_overrides.return_value = expected
|
||||
table = mem_db.create_table("test", data=[{"id": 1}])
|
||||
|
||||
assert table.reset_table_overrides() == expected
|
||||
mock_reset_table_overrides.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_update_table_overrides_encodes_json():
|
||||
class MockInner:
|
||||
async def update_table_overrides(self, encoded_overrides):
|
||||
assert json.loads(encoded_overrides) == {
|
||||
"job_types": {"cleanup": {"enabled": False}}
|
||||
}
|
||||
return json.dumps({"job_types": {"cleanup": {"enabled": False}}})
|
||||
|
||||
table = lancedb.table.AsyncTable(MockInner())
|
||||
|
||||
assert await table.update_table_overrides(
|
||||
job_types={"cleanup": {"enabled": False}}
|
||||
) == {"job_types": {"cleanup": {"enabled": False}}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api", ["legacy", "unified"])
|
||||
@pytest.mark.parametrize("cuda_failure", ["unavailable", "init_error"])
|
||||
def test_create_index_cuda_failure_suggests_cpu(
|
||||
|
||||
@@ -1849,6 +1849,45 @@ impl Table {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_table_overrides(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let overrides = inner.get_table_overrides().await.infer_error()?;
|
||||
serde_json::to_string(&overrides).map_err(|err| {
|
||||
PyRuntimeError::new_err(format!("Failed to serialize table overrides: {err}"))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn update_table_overrides<'a>(
|
||||
self_: PyRef<'a, Self>,
|
||||
overrides_json: String,
|
||||
) -> PyResult<Bound<'a, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let overrides = serde_json::from_str(&overrides_json).map_err(|err| {
|
||||
PyValueError::new_err(format!("Invalid table overrides JSON: {err}"))
|
||||
})?;
|
||||
let overrides = inner
|
||||
.update_table_overrides(overrides)
|
||||
.await
|
||||
.infer_error()?;
|
||||
serde_json::to_string(&overrides).map_err(|err| {
|
||||
PyRuntimeError::new_err(format!("Failed to serialize table overrides: {err}"))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub fn reset_table_overrides(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let overrides = inner.reset_table_overrides().await.infer_error()?;
|
||||
serde_json::to_string(&overrides).map_err(|err| {
|
||||
PyRuntimeError::new_err(format!("Failed to serialize table overrides: {err}"))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Whether a hybrid query on this table has already been refused `_rowid`.
|
||||
/// Learned from a previous refusal, never probed, so this is free.
|
||||
///
|
||||
|
||||
Reference in New Issue
Block a user