diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 7c6611608..5f83060d0 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index d9c27ba81..a7124a4da 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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. diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 0b3dc01c3..ceeb67441 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -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 diff --git a/python/src/table.rs b/python/src/table.rs index 72e0095c6..f1348de5b 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1603,6 +1603,45 @@ impl Table { }) } + pub fn get_table_overrides(self_: PyRef<'_, Self>) -> PyResult> { + 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> { + 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> { + 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 diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index b7a492598..17c739687 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -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 RestfulLanceDbClient { 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 { if let Some(ref provider) = self.header_provider { diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 7061ac979..63c8a3eb0 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -216,6 +216,7 @@ impl RemoteDatabaseOptionsBuilder { pub struct RemoteDatabase { client: RestfulLanceDbClient, table_cache: Cache>>, + db_name: String, uri: String, /// Headers to pass to the namespace client for authentication namespace_headers: HashMap, @@ -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 Database for RemoteDatabase { 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 Database for RemoteDatabase { 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 Database for RemoteDatabase { 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 Database for RemoteDatabase { 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 Database for RemoteDatabase { 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, diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 79bf2e354..fabc51fb6 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -489,6 +489,7 @@ impl Tags for RemoteTags<'_, S> { pub struct RemoteTable { client: RestfulLanceDbClient, + db_name: String, name: String, namespace: Vec, identifier: String, @@ -514,6 +515,48 @@ impl std::fmt::Debug for RemoteTable { } impl RemoteTable { + 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, + patch: serde_json::Map, + ) { + 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> { self.check_mutable().await?; let request = self @@ -651,6 +694,7 @@ impl RemoteTable { pub fn new( client: RestfulLanceDbClient, + db_name: String, name: String, namespace: Vec, identifier: String, @@ -658,6 +702,7 @@ impl RemoteTable { ) -> Self { Self { client, + db_name, name, namespace, identifier, @@ -691,6 +736,7 @@ impl RemoteTable { fn with_branch(&self, branch: Option) -> 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 BaseTable for RemoteTable { Ok(Some(spec)) } + async fn get_table_overrides(&self) -> Result { + 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 { + 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 { + 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> { 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!({ diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index fbec73dfb..f128d391b 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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 { + 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 { + 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 { + 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 { + 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 { + self.inner.update_table_overrides(overrides).await + } + + /// Clear Enterprise table overrides for this table. + pub async fn reset_table_overrides(&self) -> Result { + 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