diff --git a/docs/src/js/classes/Branches.md b/docs/src/js/classes/Branches.md index 4629620ac..5296d0d08 100644 --- a/docs/src/js/classes/Branches.md +++ b/docs/src/js/classes/Branches.md @@ -83,6 +83,24 @@ Delete a branch. *** +### diff() + +```ts +diff(fromBranch): Promise +``` + +Compare a branch against main without modifying either branch. + +#### Parameters + +* **fromBranch**: `string` + +#### Returns + +`Promise`<[`BranchDiff`](../interfaces/BranchDiff.md)> + +*** + ### list() ```ts @@ -94,3 +112,28 @@ List all branches, mapping name to branch metadata. #### Returns `Promise`<`Record`<`string`, [`BranchContents`](BranchContents.md)>> + +*** + +### merge() + +```ts +merge(fromBranch, dryRun): Promise +``` + +Merge a branch into main. + +Set `dryRun` to `true` to preview the merge. A rejected merge resolves +with `status: "rejected"` instead of throwing. + +#### Parameters + +* **fromBranch**: `string` + Branch to merge from. + +* **dryRun**: `boolean` = `false` + When true, only preview the merge. Defaults to false. + +#### Returns + +`Promise`<[`MergeBranchResult`](../interfaces/MergeBranchResult.md)> diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index e664bda29..c214a5f13 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -52,6 +52,11 @@ - [AddDataOptions](interfaces/AddDataOptions.md) - [AddResult](interfaces/AddResult.md) - [AlterColumnsResult](interfaces/AlterColumnsResult.md) +- [BranchColumnChange](interfaces/BranchColumnChange.md) +- [BranchColumnSummary](interfaces/BranchColumnSummary.md) +- [BranchDiff](interfaces/BranchDiff.md) +- [BranchIndexSummary](interfaces/BranchIndexSummary.md) +- [BranchRowCountSummary](interfaces/BranchRowCountSummary.md) - [ClientConfig](interfaces/ClientConfig.md) - [ColumnAlteration](interfaces/ColumnAlteration.md) - [ColumnOrdering](interfaces/ColumnOrdering.md) @@ -86,6 +91,9 @@ - [ListNamespacesOptions](interfaces/ListNamespacesOptions.md) - [ListNamespacesResponse](interfaces/ListNamespacesResponse.md) - [LsmWriteSpec](interfaces/LsmWriteSpec.md) +- [MergeBlocker](interfaces/MergeBlocker.md) +- [MergeBranchResult](interfaces/MergeBranchResult.md) +- [MergePreview](interfaces/MergePreview.md) - [MergeResult](interfaces/MergeResult.md) - [NativeOAuthConfig](interfaces/NativeOAuthConfig.md) - [OAuthConfig](interfaces/OAuthConfig.md) diff --git a/docs/src/js/interfaces/BranchColumnChange.md b/docs/src/js/interfaces/BranchColumnChange.md new file mode 100644 index 000000000..4502eaebc --- /dev/null +++ b/docs/src/js/interfaces/BranchColumnChange.md @@ -0,0 +1,33 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BranchColumnChange + +# Interface: BranchColumnChange + +A column whose definition differs between main and the branch. + +## Properties + +### branch + +```ts +branch: BranchColumnSummary; +``` + +*** + +### main + +```ts +main: BranchColumnSummary; +``` + +*** + +### name + +```ts +name: string; +``` diff --git a/docs/src/js/interfaces/BranchColumnSummary.md b/docs/src/js/interfaces/BranchColumnSummary.md new file mode 100644 index 000000000..4dc9fa90e --- /dev/null +++ b/docs/src/js/interfaces/BranchColumnSummary.md @@ -0,0 +1,33 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BranchColumnSummary + +# Interface: BranchColumnSummary + +Summary of a column in a branch diff. + +## Properties + +### dataType + +```ts +dataType: string; +``` + +*** + +### name + +```ts +name: string; +``` + +*** + +### nullable + +```ts +nullable: boolean; +``` diff --git a/docs/src/js/interfaces/BranchDiff.md b/docs/src/js/interfaces/BranchDiff.md new file mode 100644 index 000000000..224be1991 --- /dev/null +++ b/docs/src/js/interfaces/BranchDiff.md @@ -0,0 +1,129 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BranchDiff + +# Interface: BranchDiff + +Read-only comparison of a branch against main. + +## Properties + +### addedColumns + +```ts +addedColumns: BranchColumnSummary[]; +``` + +*** + +### addedIndexes + +```ts +addedIndexes: BranchIndexSummary[]; +``` + +*** + +### baseMoved + +```ts +baseMoved: boolean; +``` + +*** + +### branchVersion + +```ts +branchVersion: number; +``` + +*** + +### changedColumns + +```ts +changedColumns: BranchColumnChange[]; +``` + +*** + +### fromBranch + +```ts +fromBranch: string; +``` + +*** + +### mainVersion + +```ts +mainVersion: number; +``` + +*** + +### mergeBlockers + +```ts +mergeBlockers: MergeBlocker[]; +``` + +*** + +### mergeable + +```ts +mergeable: boolean; +``` + +*** + +### parentVersion + +```ts +parentVersion: number; +``` + +*** + +### removedColumns + +```ts +removedColumns: BranchColumnSummary[]; +``` + +*** + +### removedIndexes + +```ts +removedIndexes: BranchIndexSummary[]; +``` + +*** + +### rowCountBranch + +```ts +rowCountBranch: number; +``` + +*** + +### rowCountMain + +```ts +rowCountMain: number; +``` + +*** + +### rowSummary + +```ts +rowSummary: BranchRowCountSummary; +``` diff --git a/docs/src/js/interfaces/BranchIndexSummary.md b/docs/src/js/interfaces/BranchIndexSummary.md new file mode 100644 index 000000000..c9aeb1ef5 --- /dev/null +++ b/docs/src/js/interfaces/BranchIndexSummary.md @@ -0,0 +1,41 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BranchIndexSummary + +# Interface: BranchIndexSummary + +Summary of an index in a branch diff. + +## Properties + +### columns + +```ts +columns: string[]; +``` + +*** + +### indexName + +```ts +indexName: string; +``` + +*** + +### indexType? + +```ts +optional indexType: string; +``` + +*** + +### status + +```ts +status: string; +``` diff --git a/docs/src/js/interfaces/BranchRowCountSummary.md b/docs/src/js/interfaces/BranchRowCountSummary.md new file mode 100644 index 000000000..f5c03077a --- /dev/null +++ b/docs/src/js/interfaces/BranchRowCountSummary.md @@ -0,0 +1,57 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BranchRowCountSummary + +# Interface: BranchRowCountSummary + +Row-level comparison between main and the branch. + +## Properties + +### deltaAvailable + +```ts +deltaAvailable: boolean; +``` + +*** + +### inputsChanged + +```ts +inputsChanged: number; +``` + +*** + +### newOnBase + +```ts +newOnBase: number; +``` + +*** + +### newOnBranch + +```ts +newOnBranch: number; +``` + +*** + +### staleRecompute + +```ts +staleRecompute: number; +``` + +*** + +### unchanged + +```ts +unchanged: number; +``` diff --git a/docs/src/js/interfaces/MergeBlocker.md b/docs/src/js/interfaces/MergeBlocker.md new file mode 100644 index 000000000..6c8f84b37 --- /dev/null +++ b/docs/src/js/interfaces/MergeBlocker.md @@ -0,0 +1,25 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / MergeBlocker + +# Interface: MergeBlocker + +A reason why a branch cannot currently be merged. + +## Properties + +### code + +```ts +code: string; +``` + +*** + +### message + +```ts +message: string; +``` diff --git a/docs/src/js/interfaces/MergeBranchResult.md b/docs/src/js/interfaces/MergeBranchResult.md new file mode 100644 index 000000000..61c2d0d33 --- /dev/null +++ b/docs/src/js/interfaces/MergeBranchResult.md @@ -0,0 +1,46 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / MergeBranchResult + +# Interface: MergeBranchResult + +Result of previewing or attempting a branch merge. + +## Properties + +### diff + +```ts +diff: BranchDiff; +``` + +*** + +### mainVersionAfter? + +```ts +optional mainVersionAfter: number; +``` + +*** + +### preview + +```ts +preview: MergePreview; +``` + +*** + +### status + +```ts +status: + | "unknown" + | "rejected" + | "ready" + | "notImplemented" + | "merged"; +``` diff --git a/docs/src/js/interfaces/MergePreview.md b/docs/src/js/interfaces/MergePreview.md new file mode 100644 index 000000000..0d9717289 --- /dev/null +++ b/docs/src/js/interfaces/MergePreview.md @@ -0,0 +1,17 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / MergePreview + +# Interface: MergePreview + +Changes that would be, or were, promoted by a branch merge. + +## Properties + +### promotedColumns + +```ts +promotedColumns: string[]; +``` diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index d665fe6a9..a59dd4e71 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -225,6 +225,113 @@ describe("remote connection", () => { ); }); + it("diffs and merges remote branches", async () => { + const sampleDiff = { + fromBranch: "exp", + parentVersion: 1, + mainVersion: 2, + branchVersion: 3, + baseMoved: false, + rowCountMain: 3, + rowCountBranch: 3, + rowSummary: { + unchanged: 3, + newOnBase: 0, + newOnBranch: 0, + staleRecompute: 0, + inputsChanged: 0, + deltaAvailable: false, + }, + addedColumns: [{ name: "tag", dataType: "utf8", nullable: true }], + removedColumns: [], + changedColumns: [], + addedIndexes: [], + removedIndexes: [], + mergeable: true, + mergeBlockers: [], + }; + const mergeBodies: Record[] = []; + + await withMockDatabase( + (req, res) => { + const path = req.url ?? ""; + if (path.endsWith("/describe/")) { + res.writeHead(200, { "Content-Type": "application/json" }).end( + JSON.stringify({ + name: "t", + version: 2, + schema: { fields: [] }, + }), + ); + return; + } + + let raw = ""; + req.on("data", (chunk) => { + raw += chunk; + }); + req.on("end", () => { + const body = raw ? JSON.parse(raw) : {}; + if (path.endsWith("/branches/diff/")) { + // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format + expect(body).toEqual({ from_branch: "exp" }); + res + .writeHead(200, { "Content-Type": "application/json" }) + .end(JSON.stringify(sampleDiff)); + return; + } + if (path.endsWith("/branches/merge/")) { + mergeBodies.push(body); + const dryRun = body["dry_run"] === true; + const response = { + status: dryRun ? "ready" : "rejected", + diff: dryRun + ? sampleDiff + : { + ...sampleDiff, + mergeable: false, + mergeBlockers: [ + { code: "baseMoved", message: "main has advanced" }, + ], + }, + preview: { promotedColumns: dryRun ? ["tag"] : [] }, + }; + res + .writeHead(dryRun ? 200 : 409, { + "Content-Type": "application/json", + }) + .end(JSON.stringify(response)); + return; + } + res.writeHead(404).end(); + }); + }, + async (db) => { + const table = await db.openTable("t"); + const branches = await table.branches(); + + await expect(branches.diff("exp")).resolves.toEqual(sampleDiff); + + const rejected = await branches.merge("exp"); + expect(rejected.status).toBe("rejected"); + expect(rejected.diff.mergeBlockers).toEqual([ + { code: "baseMoved", message: "main has advanced" }, + ]); + + const preview = await branches.merge("exp", true); + expect(preview.status).toBe("ready"); + expect(preview.preview.promotedColumns).toEqual(["tag"]); + }, + ); + + expect(mergeBodies).toEqual([ + // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format + { from_branch: "exp", dry_run: false }, + // biome-ignore lint/style/useNamingConvention: snake_case mandated by the server wire format + { from_branch: "exp", dry_run: true }, + ]); + }); + describe("TlsConfig", () => { it("should create TlsConfig with all fields", () => { const tlsConfig: TlsConfig = { diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index a54529153..cb91c30aa 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -124,6 +124,14 @@ export { export { Table, Branches, + BranchColumnSummary, + BranchColumnChange, + BranchIndexSummary, + BranchRowCountSummary, + MergeBlocker, + BranchDiff, + MergePreview, + MergeBranchResult, AddDataOptions, UpdateOptions, OptimizeOptions, diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 07222a192..d814e3e55 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -1329,6 +1329,76 @@ export interface FieldMetadataUpdate { replace?: boolean; } +/** Summary of a column in a branch diff. */ +export interface BranchColumnSummary { + name: string; + dataType: string; + nullable: boolean; +} + +/** A column whose definition differs between main and the branch. */ +export interface BranchColumnChange { + name: string; + main: BranchColumnSummary; + branch: BranchColumnSummary; +} + +/** Summary of an index in a branch diff. */ +export interface BranchIndexSummary { + indexName: string; + columns: string[]; + indexType?: string; + status: string; +} + +/** Row-level comparison between main and the branch. */ +export interface BranchRowCountSummary { + unchanged: number; + newOnBase: number; + newOnBranch: number; + staleRecompute: number; + inputsChanged: number; + deltaAvailable: boolean; +} + +/** A reason why a branch cannot currently be merged. */ +export interface MergeBlocker { + code: string; + message: string; +} + +/** Read-only comparison of a branch against main. */ +export interface BranchDiff { + fromBranch: string; + parentVersion: number; + mainVersion: number; + branchVersion: number; + baseMoved: boolean; + rowCountMain: number; + rowCountBranch: number; + rowSummary: BranchRowCountSummary; + addedColumns: BranchColumnSummary[]; + removedColumns: BranchColumnSummary[]; + changedColumns: BranchColumnChange[]; + addedIndexes: BranchIndexSummary[]; + removedIndexes: BranchIndexSummary[]; + mergeable: boolean; + mergeBlockers: MergeBlocker[]; +} + +/** Changes that would be, or were, promoted by a branch merge. */ +export interface MergePreview { + promotedColumns: string[]; +} + +/** Result of previewing or attempting a branch merge. */ +export interface MergeBranchResult { + status: "ready" | "rejected" | "notImplemented" | "merged" | "unknown"; + diff: BranchDiff; + preview: MergePreview; + mainVersionAfter?: number; +} + /** * Branch manager for a {@link Table}. * @@ -1381,4 +1451,28 @@ export class Branches { async delete(name: string): Promise { return await this.#inner.delete(name); } + + /** Compare a branch against main without modifying either branch. */ + async diff(fromBranch: string): Promise { + return (await this.#inner.diff(fromBranch)) as unknown as BranchDiff; + } + + /** + * Merge a branch into main. + * + * Set `dryRun` to `true` to preview the merge. A rejected merge resolves + * with `status: "rejected"` instead of throwing. + * + * @param fromBranch Branch to merge from. + * @param dryRun When true, only preview the merge. Defaults to false. + */ + async merge( + fromBranch: string, + dryRun: boolean = false, + ): Promise { + return (await this.#inner.merge( + fromBranch, + dryRun, + )) as unknown as MergeBranchResult; + } } diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 8252a8c3b..d1489034f 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -1355,4 +1355,28 @@ impl Branches { pub async fn delete(&self, name: String) -> napi::Result<()> { self.inner.delete_branch(&name).await.default_error() } + + #[napi(ts_return_type = "Promise>")] + pub async fn diff(&self, from_branch: String) -> napi::Result { + let diff = self.inner.diff_branch(&from_branch).await.default_error()?; + serde_json::to_value(diff).map_err(|err| { + napi::Error::from_reason(format!("failed to serialize branch diff: {err}")) + }) + } + + #[napi(ts_return_type = "Promise>")] + pub async fn merge( + &self, + from_branch: String, + dry_run: Option, + ) -> napi::Result { + let result = self + .inner + .merge_branch(&from_branch, dry_run.unwrap_or(false)) + .await + .default_error()?; + serde_json::to_value(result).map_err(|err| { + napi::Error::from_reason(format!("failed to serialize branch merge result: {err}")) + }) + } } diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 3b94c5d45..c922aa69b 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -319,6 +319,10 @@ class Branches: ) -> Table: ... async def checkout(self, name: str, version: Optional[int] = None) -> Table: ... async def delete(self, name: str) -> None: ... + async def diff(self, from_branch: str) -> Dict[str, Any]: ... + async def merge( + self, from_branch: str, dry_run: bool = False + ) -> Dict[str, Any]: ... class IndexConfig: name: str diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 6739d5864..0e181b457 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -6267,6 +6267,24 @@ class Branches: """Delete a branch.""" LOOP.run(self._table.branches.delete(name)) + def diff(self, from_branch: str) -> Dict[str, Any]: + """Diff a branch against main.""" + return LOOP.run(self._table.branches.diff(from_branch)) + + def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]: + """Merge a branch into main, or dry-run. + + Parameters + ---------- + from_branch: str + Branch to merge from. + dry_run: bool, default False + When True, only preview. When False, attempt the merge. + + A rejected merge returns ``status="rejected"`` instead of raising. + """ + return LOOP.run(self._table.branches.merge(from_branch, dry_run)) + def _wrap( self, async_table: "AsyncTable", version: Optional[int] = None ) -> "Table": @@ -6396,3 +6414,14 @@ class AsyncBranches: async def delete(self, name: str) -> None: """Delete a branch.""" await self._table.branches.delete(name) + + async def diff(self, from_branch: str) -> Dict[str, Any]: + """Diff a branch against main.""" + return await self._table.branches.diff(from_branch) + + async def merge(self, from_branch: str, dry_run: bool = False) -> Dict[str, Any]: + """Merge a branch into main, or dry-run. + + A rejected merge returns ``status="rejected"`` instead of raising. + """ + return await self._table.branches.merge(from_branch, dry_run) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 4f54f5961..477ee7496 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -236,6 +236,65 @@ def test_remote_table_branches_sync(): table.branches.delete("exp") +def test_remote_table_branch_merge_defaults_to_execute(): + merge_bodies = [] + diff = { + "fromBranch": "exp", + "parentVersion": 1, + "mainVersion": 2, + "branchVersion": 3, + "baseMoved": False, + "rowCountMain": 3, + "rowCountBranch": 3, + "rowSummary": { + "unchanged": 3, + "newOnBase": 0, + "newOnBranch": 0, + "staleRecompute": 0, + "inputsChanged": 0, + "deltaAvailable": False, + }, + "addedColumns": [], + "removedColumns": [], + "changedColumns": [], + "addedIndexes": [], + "removedIndexes": [], + "mergeable": True, + "mergeBlockers": [], + } + + def handler(request): + if request.path.endswith("/describe/"): + status = 200 + body = {"version": 2, "schema": {"fields": []}} + else: + content_len = int(request.headers.get("Content-Length")) + request_body = json.loads(request.rfile.read(content_len)) + merge_bodies.append(request_body) + dry_run = request_body["dry_run"] + status = 200 if dry_run else 409 + body = { + "status": "ready" if dry_run else "rejected", + "diff": diff, + "preview": {"promotedColumns": []}, + } + + request.send_response(status) + request.send_header("Content-Type", "application/json") + request.end_headers() + request.wfile.write(json.dumps(body).encode()) + + with mock_lancedb_connection(handler) as db: + branches = db.open_table("test").branches + assert branches.merge("exp")["status"] == "rejected" + assert branches.merge("exp", dry_run=True)["status"] == "ready" + + assert merge_bodies == [ + {"from_branch": "exp", "dry_run": False}, + {"from_branch": "exp", "dry_run": True}, + ] + + @pytest.mark.asyncio async def test_async_remote_open_table_branch_and_version(): async with mock_lancedb_connection_async(_branch_open_handler) as db: diff --git a/python/src/table.rs b/python/src/table.rs index 703c64c53..5d6910685 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1593,4 +1593,40 @@ impl Branches { Ok(()) }) } + + pub fn diff(self_: PyRef<'_, Self>, from_branch: String) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let diff = inner.diff_branch(&from_branch).await.infer_error()?; + Python::attach(|py| struct_to_wire_py(py, &diff)) + }) + } + + #[pyo3(signature = (from_branch, dry_run=false))] + pub fn merge( + self_: PyRef<'_, Self>, + from_branch: String, + dry_run: bool, + ) -> PyResult> { + let inner = self_.inner.clone(); + future_into_py(self_.py(), async move { + let result = inner + .merge_branch(&from_branch, dry_run) + .await + .infer_error()?; + Python::attach(|py| struct_to_wire_py(py, &result)) + }) + } +} + +/// Decode a serde value as the wire JSON object (camelCase keys). +fn struct_to_wire_py(py: Python<'_>, value: &impl serde::Serialize) -> PyResult> { + let json = py.import("json")?; + Ok(json + .call_method1( + "loads", + (serde_json::to_string(value) + .map_err(|e| PyRuntimeError::new_err(format!("failed to serialize json: {e}")))?,), + )? + .unbind()) } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 90e0d3dd6..2518ed65c 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -18,9 +18,11 @@ use crate::index::waiter::wait_for_index; use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest}; use crate::table::AddColumnsResult; use crate::table::AddResult; +use crate::table::BranchDiff; use crate::table::DeleteResult; use crate::table::DropColumnsResult; use crate::table::LsmWriteSpec; +use crate::table::MergeBranchResult; use crate::table::MergeResult; use crate::table::Tags; use crate::table::UpdateResult; @@ -1817,6 +1819,79 @@ impl BaseTable for RemoteTable { Ok(()) } + async fn diff_branch(&self, from_branch: &str) -> Result { + if from_branch.trim().is_empty() { + return Err(Error::InvalidInput { + message: "from_branch must be a non-empty string".into(), + }); + } + let request = self + .client + .post(&format!("/v1/table/{}/branches/diff/", self.identifier)) + .json(&serde_json::json!({ "from_branch": from_branch })); + let (request_id, response) = self.send(request, true).await?; + if response.status() == StatusCode::NOT_FOUND { + return Err(Error::TableNotFound { + name: format!("{} (branch: {})", self.name, from_branch), + source: format!("branch '{}' does not exist", from_branch).into(), + }); + } + let response = self.check_table_response(&request_id, response).await?; + let body = response.text().await.err_to_http(request_id.clone())?; + serde_json::from_str(&body).map_err(|err| Error::Http { + source: format!( + "Failed to parse diff_branch response: {}, body: {}", + err, body + ) + .into(), + request_id, + status_code: None, + }) + } + + async fn merge_branch(&self, from_branch: &str, dry_run: bool) -> Result { + if from_branch.trim().is_empty() { + return Err(Error::InvalidInput { + message: "from_branch must be a non-empty string".into(), + }); + } + let request = self + .client + .post(&format!("/v1/table/{}/branches/merge/", self.identifier)) + .json(&serde_json::json!({ + "from_branch": from_branch, + "dry_run": dry_run, + })); + // No retry. 409 rejected merge is final and carries a body. + let (request_id, response) = self.send(request, false).await?; + let status = response.status(); + if status == StatusCode::NOT_FOUND { + return Err(Error::TableNotFound { + name: format!("{} (branch: {})", self.name, from_branch), + source: format!("branch '{}' does not exist", from_branch).into(), + }); + } + // 200 and 409 both carry MergeBranchResult. + if status != StatusCode::OK && status != StatusCode::CONFLICT { + let body = response.text().await.unwrap_or_default(); + return Err(Error::Http { + source: format!("unexpected status {status} from merge_branch: {body}").into(), + request_id, + status_code: Some(status), + }); + } + let body = response.text().await.err_to_http(request_id.clone())?; + serde_json::from_str(&body).map_err(|err| Error::Http { + source: format!( + "Failed to parse merge_branch response: {}, body: {}", + err, body + ) + .into(), + request_id, + status_code: Some(status), + }) + } + fn current_branch(&self) -> Option { self.branch.clone() } @@ -8211,6 +8286,147 @@ mod tests { assert!(matches!(err, Error::TableNotFound { .. }), "got {err:?}"); } + fn sample_branch_diff_json() -> &'static str { + r#"{ + "fromBranch":"exp", + "parentVersion":1, + "mainVersion":1, + "branchVersion":2, + "baseMoved":false, + "rowCountMain":3, + "rowCountBranch":3, + "rowSummary":{ + "unchanged":3, + "newOnBase":0, + "newOnBranch":0, + "staleRecompute":0, + "inputsChanged":0, + "deltaAvailable":false + }, + "addedColumns":[{"name":"tag","dataType":"utf8","nullable":true}], + "removedColumns":[], + "changedColumns":[], + "addedIndexes":[], + "removedIndexes":[], + "mergeable":true, + "mergeBlockers":[] + }"# + } + + #[tokio::test] + async fn test_diff_branch() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/branches/diff/"); + let body = request_body_json(&request); + assert_eq!(body["from_branch"], "exp"); + http::Response::builder() + .status(200) + .body(sample_branch_diff_json()) + .unwrap() + }); + let diff = table.diff_branch("exp").await.unwrap(); + assert_eq!(diff.from_branch, "exp"); + assert!(diff.mergeable); + assert_eq!(diff.added_columns.len(), 1); + assert_eq!(diff.added_columns[0].name, "tag"); + } + + #[tokio::test] + async fn test_merge_branch_dry_run() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/branches/merge/"); + let body = request_body_json(&request); + assert_eq!(body["from_branch"], "exp"); + assert_eq!(body["dry_run"], true); + let resp = format!( + r#"{{"status":"ready","diff":{},"preview":{{"promotedColumns":["tag"]}}}}"#, + sample_branch_diff_json() + ); + http::Response::builder().status(200).body(resp).unwrap() + }); + let result = table.merge_branch("exp", true).await.unwrap(); + assert_eq!(result.status, crate::table::MergeBranchStatus::Ready); + assert_eq!(result.preview.promoted_columns, vec!["tag".to_string()]); + assert!(result.main_version_after.is_none()); + } + + #[tokio::test] + async fn test_merge_branch_rejected_returns_ok_with_body() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.url().path(), "/v1/table/my_table/branches/merge/"); + let body = request_body_json(&request); + assert_eq!(body["dry_run"], false); + let mut diff: serde_json::Value = + serde_json::from_str(sample_branch_diff_json()).unwrap(); + diff["mergeable"] = serde_json::json!(false); + diff["mergeBlockers"] = serde_json::json!([{ + "code": "baseMoved", + "message": "main has advanced" + }]); + let resp = serde_json::json!({ + "status": "rejected", + "diff": diff, + "preview": { "promotedColumns": [] } + }); + http::Response::builder() + .status(409) + .body(resp.to_string()) + .unwrap() + }); + let result = table.merge_branch("exp", false).await.unwrap(); + assert_eq!(result.status, crate::table::MergeBranchStatus::Rejected); + assert!(!result.diff.mergeable); + assert_eq!(result.diff.merge_blockers.len(), 1); + } + + #[tokio::test] + async fn test_merge_branch_unknown_blocker_code_parses() { + let table = Table::new_with_handler("my_table", |_| { + let mut diff: serde_json::Value = + serde_json::from_str(sample_branch_diff_json()).unwrap(); + diff["mergeable"] = serde_json::json!(false); + diff["mergeBlockers"] = serde_json::json!([{ + "code": "multipleCommits", + "message": "branch has more than one data commit" + }]); + let resp = serde_json::json!({ + "status": "rejected", + "diff": diff, + "preview": { "operation": "append", "rowsAdded": 2 } + }); + http::Response::builder() + .status(409) + .body(resp.to_string()) + .unwrap() + }); + let result = table.merge_branch("exp", false).await.unwrap(); + assert_eq!(result.status, crate::table::MergeBranchStatus::Rejected); + assert_eq!( + result.diff.merge_blockers[0].code, + crate::table::MergeBlockerCode::Unknown + ); + assert!(result.preview.promoted_columns.is_empty()); + } + + #[tokio::test] + async fn test_merge_branch_unexpected_2xx_is_error() { + let table = Table::new_with_handler("my_table", |_| { + http::Response::builder() + .status(204) + .body(String::new()) + .unwrap() + }); + let err = table.merge_branch("exp", false).await.unwrap_err(); + match err { + Error::Http { + status_code: Some(code), + .. + } => assert_eq!(code, reqwest::StatusCode::NO_CONTENT), + other => panic!("expected Http error, got {other:?}"), + } + } + #[tokio::test] async fn test_checkout_branch_validates_via_list() { let table = Table::new_with_handler("my_table", |request| { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 580b97010..1f7b2f75b 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -62,6 +62,7 @@ use self::dataset::DatasetConsistencyWrapper; use self::merge::MergeInsertBuilder; mod add_data; +pub mod branch_merge; mod create_index; pub mod datafusion; pub(crate) mod dataset; @@ -77,6 +78,10 @@ use crate::index::waiter::wait_for_index; #[cfg(feature = "remote")] pub(crate) use add_data::PreprocessingOutput; pub use add_data::{AddDataBuilder, AddDataMode, AddResult, NaNVectorBehavior}; +pub use branch_merge::{ + BranchDiff, ColumnChange, ColumnSummary, IndexSummary, MergeBlocker, MergeBlockerCode, + MergeBranchResult, MergeBranchStatus, MergePreview, RowCountSummary, +}; pub use chrono::Duration; pub use delete::DeleteResult; use futures::future::join_all; @@ -707,6 +712,19 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { async fn list_branches(&self) -> Result>; /// Delete a branch. async fn delete_branch(&self, name: &str) -> Result<()>; + /// Diff a branch against main. Remote only. + async fn diff_branch(&self, _from_branch: &str) -> Result { + Err(Error::NotSupported { + message: "diff_branch is only supported on remote tables".into(), + }) + } + /// Merge a branch into main, or dry-run. Remote only. + /// HTTP 409 still returns [`Ok`] with [`MergeBranchStatus::Rejected`]. + async fn merge_branch(&self, _from_branch: &str, _dry_run: bool) -> Result { + Err(Error::NotSupported { + message: "merge_branch is only supported on remote tables".into(), + }) + } /// The branch this handle is scoped to, or `None` for `main`. fn current_branch(&self) -> Option; /// Get the table definition. @@ -1953,6 +1971,21 @@ impl Table { self.inner.delete_branch(name).await } + /// Diff a branch against main. Remote only. + pub async fn diff_branch(&self, from_branch: &str) -> Result { + self.inner.diff_branch(from_branch).await + } + + /// Merge a branch into main, or dry-run. Remote only. + /// HTTP 409 still returns [`Ok`] with [`MergeBranchStatus::Rejected`]. + pub async fn merge_branch( + &self, + from_branch: &str, + dry_run: bool, + ) -> Result { + self.inner.merge_branch(from_branch, dry_run).await + } + /// The branch this handle is scoped to, or `None` for `main`. pub fn current_branch(&self) -> Option { self.inner.current_branch() diff --git a/rust/lancedb/src/table/branch_merge.rs b/rust/lancedb/src/table/branch_merge.rs new file mode 100644 index 000000000..ad81ab64c --- /dev/null +++ b/rust/lancedb/src/table/branch_merge.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Types for remote branch diff / merge against main. + +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ColumnSummary { + pub name: String, + pub data_type: String, + pub nullable: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ColumnChange { + pub name: String, + pub main: ColumnSummary, + pub branch: ColumnSummary, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct IndexSummary { + pub index_name: String, + pub columns: Vec, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub index_type: Option, + pub status: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct RowCountSummary { + pub unchanged: u64, + pub new_on_base: u64, + pub new_on_branch: u64, + pub stale_recompute: u64, + pub inputs_changed: u64, + pub delta_available: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub enum MergeBlockerCode { + BaseMoved, + RowCountMismatch, + RowsChanged, + ColumnRemoved, + ColumnChanged, + NoMergeableChanges, + NoColumnChanges, + InputColumnDependency, + ParentNotMain, + #[serde(other)] + Unknown, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct MergeBlocker { + pub code: MergeBlockerCode, + pub message: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct BranchDiff { + pub from_branch: String, + pub parent_version: u64, + pub main_version: u64, + pub branch_version: u64, + pub base_moved: bool, + pub row_count_main: u64, + pub row_count_branch: u64, + pub row_summary: RowCountSummary, + pub added_columns: Vec, + pub removed_columns: Vec, + pub changed_columns: Vec, + pub added_indexes: Vec, + pub removed_indexes: Vec, + pub mergeable: bool, + pub merge_blockers: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct MergePreview { + #[serde(default)] + pub promoted_columns: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub enum MergeBranchStatus { + Ready, + Rejected, + NotImplemented, + Merged, + #[serde(other)] + Unknown, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct MergeBranchResult { + pub status: MergeBranchStatus, + pub diff: BranchDiff, + pub preview: MergePreview, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub main_version_after: Option, +}