feat(sdk): expose table overrides for remote tables

This commit is contained in:
ChilePiquin
2026-09-18 12:06:48 -06:00
parent 7955c50929
commit 2ef1173724
8 changed files with 444 additions and 1 deletions
+3
View File
@@ -494,6 +494,9 @@ class Table:
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
async def get_table_overrides(self) -> str: ...
async def update_table_overrides(self, overrides_json: str) -> str: ...
async def reset_table_overrides(self) -> str: ...
async def checkpoint_lsm(self) -> None: ...
async def flush_lsm(self) -> None: ...
async def compact_lsm(self) -> None: ...
+62
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import asyncio
import inspect
import json
import deprecation
import warnings
from abc import ABC, abstractmethod
@@ -4449,6 +4450,36 @@ class LanceTable(Table):
[`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec]."""
return LOOP.run(self._table.get_lsm_write_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 existing
overrides are preserved by the remote service update flow.
Examples
--------
>>> table.update_table_overrides(
... 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_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
@@ -5188,6 +5219,37 @@ class AsyncTable:
"""
return await self._inner.get_lsm_write_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 existing
overrides are preserved by the remote service update flow.
Examples
--------
>>> await table.update_table_overrides(
... 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())
async def checkpoint_lsm(self) -> None:
"""Converge this table's LSM write path into its base table.
+61
View File
@@ -4,6 +4,7 @@
import ctypes
import gc
import json
import os
import sys
import threading
@@ -1640,6 +1641,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}}}
@patch("lancedb.table.AsyncTable.create_index")
def test_create_index_name_and_train_parameters(
mock_create_index, mem_db: DBConnection
+39
View File
@@ -1603,6 +1603,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}"))
})
})
}
/// Converge the table's LSM write path into its base table.
///
/// Best-effort: with writes flowing, new rows may land after the last
+6 -1
View File
@@ -448,7 +448,7 @@ enum BodyLogging {
/// always splits back into the parts that made it. The configuration field
/// exists because the identifier grammar comes from the Lance REST catalog
/// standard, which carries a delimiter setting for other catalogs to adopt.
pub(crate) const ID_DELIMITER: &str = "$";
pub const ID_DELIMITER: &str = "$";
fn validate_id_delimiter(delimiter: &str) -> Result<()> {
if delimiter != ID_DELIMITER {
@@ -767,6 +767,11 @@ impl<S: HttpSend> RestfulLanceDbClient<S> {
self.client.post(full_uri)
}
pub fn put(&self, uri: &str) -> RequestBuilder {
let full_uri = format!("{}{}", self.host, uri);
self.client.put(full_uri)
}
/// Apply dynamic headers from the header provider if configured
pub(crate) async fn apply_dynamic_headers(&self, mut request: Request) -> Result<Request> {
if let Some(ref provider) = self.header_provider {
+9
View File
@@ -216,6 +216,7 @@ impl RemoteDatabaseOptionsBuilder {
pub struct RemoteDatabase<S: HttpSend = Sender> {
client: RestfulLanceDbClient<S>,
table_cache: Cache<String, Arc<RemoteTable<S>>>,
db_name: String,
uri: String,
/// Headers to pass to the namespace client for authentication
namespace_headers: HashMap<String, String>,
@@ -410,6 +411,7 @@ impl RemoteDatabase {
Ok(Self {
client,
table_cache,
db_name: parsed.db_name,
uri: uri.to_owned(),
namespace_headers,
namespace_context_provider,
@@ -525,6 +527,7 @@ mod test_utils {
Self {
client,
table_cache: Cache::new(0),
db_name: "default".to_string(),
uri: "http://localhost".to_string(),
namespace_headers: HashMap::new(),
namespace_context_provider: None,
@@ -548,6 +551,7 @@ mod test_utils {
Self {
client,
table_cache: Cache::new(0),
db_name: "default".to_string(),
uri: "http://localhost".to_string(),
namespace_headers: config.extra_headers.clone(),
namespace_context_provider,
@@ -1218,6 +1222,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let cache_key = build_cache_key(table, &request.namespace_path);
let remote_table = Arc::new(RemoteTable::new(
self.client.clone(),
self.db_name.clone(),
table.clone(),
request.namespace_path.clone(),
table_identifier.clone(),
@@ -1254,6 +1259,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let cache_key = build_cache_key(table, &namespace_vec);
let remote_table = Arc::new(RemoteTable::new(
self.client.clone(),
self.db_name.clone(),
table.clone(),
namespace_vec.clone(),
table_identifier.clone(),
@@ -1324,6 +1330,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let cache_key = build_cache_key(&request.name, &request.namespace_path);
let table = Arc::new(RemoteTable::new(
self.client.clone(),
self.db_name.clone(),
request.name.clone(),
request.namespace_path.clone(),
table_identifier,
@@ -1366,6 +1373,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let cache_key = build_cache_key(&request.target_table_name, &request.target_namespace_path);
let table = Arc::new(RemoteTable::new(
self.client.clone(),
self.db_name.clone(),
request.target_table_name.clone(),
request.target_namespace_path.clone(),
table_identifier,
@@ -1396,6 +1404,7 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
let table_identifier = build_table_identifier(&request.name, &request.namespace_path)?;
let table = Arc::new(RemoteTable::new(
self.client.clone(),
self.db_name.clone(),
request.name.clone(),
request.namespace_path.clone(),
table_identifier,
+209
View File
@@ -489,6 +489,7 @@ impl<S: HttpSend + 'static> Tags for RemoteTags<'_, S> {
pub struct RemoteTable<S: HttpSend = Sender> {
client: RestfulLanceDbClient<S>,
db_name: String,
name: String,
namespace: Vec<String>,
identifier: String,
@@ -514,6 +515,48 @@ impl<S: HttpSend> std::fmt::Debug for RemoteTable<S> {
}
impl<S: HttpSend> RemoteTable<S> {
fn table_overrides_query(&self) -> Vec<(&str, &str)> {
let mut query = Vec::with_capacity(3 + self.namespace.len());
query.push(("db", self.db_name.as_str()));
query.push(("table", self.name.as_str()));
for namespace in &self.namespace {
query.push(("namespace", namespace.as_str()));
}
if let Some(branch) = self.branch.as_deref() {
query.push(("branch", branch));
}
query
}
fn table_overrides_set_body(&self, overrides: serde_json::Value) -> serde_json::Value {
let mut body = serde_json::json!({
"db": &self.db_name,
"namespace": &self.namespace,
"table": &self.name,
"overrides": overrides,
});
if let Some(branch) = &self.branch {
body["branch"] = serde_json::Value::String(branch.clone());
}
body
}
fn merge_table_overrides(
target: &mut serde_json::Map<String, serde_json::Value>,
patch: serde_json::Map<String, serde_json::Value>,
) {
for (key, value) in patch {
match (target.get_mut(&key), value) {
(Some(serde_json::Value::Object(existing)), serde_json::Value::Object(patch)) => {
Self::merge_table_overrides(existing, patch);
}
(_, value) => {
target.insert(key, value);
}
}
}
}
async fn submit_create_index(&self, mut index: IndexBuilder) -> Result<Option<String>> {
self.check_mutable().await?;
let request = self
@@ -651,6 +694,7 @@ impl<S: HttpSend> RemoteTable<S> {
pub fn new(
client: RestfulLanceDbClient<S>,
db_name: String,
name: String,
namespace: Vec<String>,
identifier: String,
@@ -658,6 +702,7 @@ impl<S: HttpSend> RemoteTable<S> {
) -> Self {
Self {
client,
db_name,
name,
namespace,
identifier,
@@ -691,6 +736,7 @@ impl<S: HttpSend> RemoteTable<S> {
fn with_branch(&self, branch: Option<String>) -> Self {
Self {
client: self.client.clone(),
db_name: self.db_name.clone(),
name: self.name.clone(),
namespace: self.namespace.clone(),
identifier: self.identifier.clone(),
@@ -1594,6 +1640,7 @@ mod test_utils {
let client = client_with_handler(handler);
Self {
client,
db_name: "default".to_string(),
name: name.clone(),
namespace: vec![],
identifier: name,
@@ -1618,6 +1665,7 @@ mod test_utils {
let client = client_with_handler_and_interval(handler, read_consistency_interval);
Self {
client,
db_name: "default".to_string(),
name: name.clone(),
namespace: vec![],
identifier: name,
@@ -1651,6 +1699,7 @@ mod test_utils {
let client = client_with_handler_and_config(handler, config);
Self {
client,
db_name: "default".to_string(),
name: name.clone(),
namespace: vec![],
identifier: name,
@@ -3278,6 +3327,56 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
Ok(Some(spec))
}
async fn get_table_overrides(&self) -> Result<serde_json::Value> {
let request = self
.client
.get("/admin/table/overrides")
.query(&self.table_overrides_query());
let (request_id, response) = self.send(request, true).await?;
let response = self.client.check_response(&request_id, response).await?;
response.json().await.err_to_http(request_id)
}
async fn update_table_overrides(
&self,
overrides: serde_json::Value,
) -> Result<serde_json::Value> {
self.check_mutable().await?;
let serde_json::Value::Object(patch) = overrides else {
return Err(Error::InvalidInput {
message: "table overrides update must be a JSON object".into(),
});
};
let current = self.get_table_overrides().await?;
let mut current = match current {
serde_json::Value::Object(current) => current,
other => {
return Err(Error::InvalidInput {
message: format!("table overrides response must be an object, got {other}"),
});
}
};
Self::merge_table_overrides(&mut current, patch);
let merged = serde_json::Value::Object(current);
let body = self.table_overrides_set_body(merged);
let request = self.client.put("/admin/table/overrides").json(&body);
let (request_id, response) = self.send(request, true).await?;
let response = self.client.check_response(&request_id, response).await?;
response.json().await.err_to_http(request_id)
}
async fn reset_table_overrides(&self) -> Result<serde_json::Value> {
self.check_mutable().await?;
let request = self
.client
.post("/admin/table/overrides/reset")
.query(&self.table_overrides_query());
let (request_id, response) = self.send(request, true).await?;
let response = self.client.check_response(&request_id, response).await?;
response.json().await.err_to_http(request_id)
}
async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
Ok(Box::new(RemoteTags { inner: self }))
}
@@ -9045,6 +9144,116 @@ mod tests {
assert!(table.get_lsm_write_spec().await.unwrap().is_none());
}
#[tokio::test]
async fn test_get_table_overrides() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "GET");
assert_eq!(request.url().path(), "/admin/table/overrides");
assert_eq!(request.url().query(), Some("db=default&table=my_table"));
http::Response::builder()
.status(200)
.body(
serde_json::json!({
"job_types": {
"cleanup": { "enabled": false }
}
})
.to_string(),
)
.unwrap()
});
let overrides = table.get_table_overrides().await.unwrap();
assert_eq!(overrides["job_types"]["cleanup"]["enabled"], false);
}
#[tokio::test]
async fn test_update_table_overrides_merges_existing_fields() {
let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let seen = calls.clone();
let table = Table::new_with_handler("my_table", move |request| {
match seen.fetch_add(1, std::sync::atomic::Ordering::SeqCst) {
0 => {
assert_eq!(request.method(), "GET");
assert_eq!(request.url().path(), "/admin/table/overrides");
http::Response::builder()
.status(200)
.body(
serde_json::json!({
"job_types": {
"compaction": { "enabled": true },
"reindex_vector": {
"table": { "num_workers": 3 }
}
},
"unknown": { "keep": true }
})
.to_string(),
)
.unwrap()
}
1 => {
assert_eq!(request.method(), "PUT");
assert_eq!(request.url().path(), "/admin/table/overrides");
let body = request_body_json(&request);
assert_eq!(body["db"], "default");
assert_eq!(body["table"], "my_table");
assert_eq!(body["namespace"], serde_json::json!([]));
assert_eq!(
body["overrides"],
serde_json::json!({
"job_types": {
"compaction": { "enabled": true },
"cleanup": { "enabled": false },
"reindex_vector": {
"table": { "num_workers": 3 }
}
},
"unknown": { "keep": true }
})
);
http::Response::builder()
.status(200)
.body(body["overrides"].to_string())
.unwrap()
}
_ => panic!("unexpected request"),
}
});
let updated = table
.update_table_overrides(serde_json::json!({
"job_types": {
"cleanup": { "enabled": false }
}
}))
.await
.unwrap();
assert_eq!(updated["job_types"]["cleanup"]["enabled"], false);
assert_eq!(updated["job_types"]["compaction"]["enabled"], true);
assert_eq!(
updated["job_types"]["reindex_vector"]["table"]["num_workers"],
3
);
assert_eq!(updated["unknown"]["keep"], true);
assert_eq!(calls.load(std::sync::atomic::Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_reset_table_overrides() {
let table = Table::new_with_handler("my_table", |request| {
assert_eq!(request.method(), "POST");
assert_eq!(request.url().path(), "/admin/table/overrides/reset");
assert_eq!(request.url().query(), Some("db=default&table=my_table"));
http::Response::builder().status(200).body("{}").unwrap()
});
assert_eq!(
table.reset_table_overrides().await.unwrap(),
serde_json::json!({})
);
}
/// Build a `get_lsm_stats` body for one bucket holding `generations`.
fn stats_body(generations: &[u64], compacting: bool) -> String {
serde_json::json!({
+55
View File
@@ -39,6 +39,7 @@ use lance_table::io::commit::CommitHandler;
use lance_table::io::commit::ManifestNamingScheme;
use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
use serde::{Deserialize, Serialize};
pub type TableOverrides = serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::format;
use std::path::Path;
@@ -706,6 +707,33 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
message: "get_lsm_write_spec is not supported on this table type".into(),
})
}
/// Read Enterprise table overrides.
///
/// The default implementation returns `NotSupported`; Enterprise remote
/// tables override it.
async fn get_table_overrides(&self) -> Result<TableOverrides> {
Err(Error::NotSupported {
message: "get_table_overrides is supported only on remote Enterprise tables".into(),
})
}
/// Merge Enterprise table overrides into the currently stored object.
///
/// The default implementation returns `NotSupported`; Enterprise remote
/// tables override it.
async fn update_table_overrides(&self, _overrides: TableOverrides) -> Result<TableOverrides> {
Err(Error::NotSupported {
message: "update_table_overrides is supported only on remote Enterprise tables".into(),
})
}
/// Clear Enterprise table overrides.
///
/// The default implementation returns `NotSupported`; Enterprise remote
/// tables override it.
async fn reset_table_overrides(&self) -> Result<TableOverrides> {
Err(Error::NotSupported {
message: "reset_table_overrides is supported only on remote Enterprise tables".into(),
})
}
/// Seal every bucket's active memtable into L0.
///
/// The default implementation returns `NotSupported`.
@@ -2009,6 +2037,33 @@ impl Table {
self.inner.get_lsm_write_spec().await
}
/// Read Enterprise table overrides for this table.
///
/// Table overrides configure Enterprise-managed background behavior such as
/// automatic compaction, cleanup, and reindex execution. Local tables return
/// `NotSupported` because they do not have an Enterprise background planner
/// consuming these overrides.
pub async fn get_table_overrides(&self) -> Result<TableOverrides> {
self.inner.get_table_overrides().await
}
/// Merge Enterprise table overrides into the current stored overrides.
///
/// This method uses patch semantics: fields present in `overrides` are
/// recursively merged into the existing table overrides and unrelated fields
/// are preserved.
pub async fn update_table_overrides(
&self,
overrides: TableOverrides,
) -> Result<TableOverrides> {
self.inner.update_table_overrides(overrides).await
}
/// Clear Enterprise table overrides for this table.
pub async fn reset_table_overrides(&self) -> Result<TableOverrides> {
self.inner.reset_table_overrides().await
}
/// Converge this table's LSM write path into its base table.
///
/// One `flush` to seal every memtable into L0, then compaction triggers