diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index a9fc06bdd..49cbad59c 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1033,6 +1033,18 @@ class RemoteTable(Table): """Read the installed LsmWriteSpec, or ``None``.""" return LOOP.run(self._table.get_lsm_write_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()) + def checkpoint_lsm(self) -> None: """Synchronous version of [`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm].""" diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index a7124a4da..1a424339d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1245,6 +1245,27 @@ class Table(ABC): """ raise NotImplementedError + 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. diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 1c887b218..03e038344 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -11,6 +11,7 @@ import re import sys import threading import time +from urllib.parse import parse_qs, urlparse from unittest.mock import MagicMock, patch import uuid from packaging.version import Version @@ -32,6 +33,9 @@ def make_mock_http_handler(handler): def do_POST(self): handler(self) + def do_PUT(self): + handler(self) + return MockLanceDBHandler @@ -1429,6 +1433,118 @@ def test_checkpoint_lsm_sync(): assert called == ["flush_lsm", "get_lsm_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, *, server_version=Version("0.1.0")): def handler(request):