diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 2ba96115f..0077e95cf 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2973,6 +2973,11 @@ impl BaseTable for RemoteTable { &self, updates: &[FieldMetadataUpdate], ) -> Result { + // Reject reserved generated-column key updates before mutability checks, + // request construction, header-provider invocation, retry, or HTTP. + crate::table::schema_evolution::reject_reserved_generated_column_metadata_key_updates( + updates, + )?; self.check_mutable().await?; let mut body = serde_json::json!({ "updates": updates }); self.apply_branch_body(&mut body); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index c30ea1b87..0321513a1 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -95,6 +95,8 @@ mod merge_insert_generated_column_reject_contract; #[cfg(test)] mod schema_metadata_updates_dependency_contract; #[cfg(test)] +mod update_field_metadata_generated_column_guard_contract; +#[cfg(test)] mod update_generated_column_invalidation_contract; use crate::index::waiter::wait_for_index; @@ -5828,19 +5830,20 @@ mod tests { dependency_epoch: u64, materialized_epoch: u64, ) -> i32 { - use crate::function::GENERATED_COLUMN_METADATA_KEY; - let snapshot = table.generated_column_binding_snapshot().await.unwrap(); let field_id = snapshot.field(column).expect(column).field_id(); let json = status_projection_definition(field_id, dependency_epoch, materialized_epoch) .to_metadata_json() .unwrap(); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new(column).set(GENERATED_COLUMN_METADATA_KEY, json) - ]) - .await - .unwrap(); + schema_evolution::install_raw_generated_column_metadata_for_tests( + table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + column, + json, + ) + .await + .unwrap(); field_id } diff --git a/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs b/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs index 323858f3e..34c8c9623 100644 --- a/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs +++ b/rust/lancedb/src/table/append_generated_column_invalidation_contract.rs @@ -28,7 +28,6 @@ use crate::function::{ }; use crate::query::{ExecutableQuery, QueryBase, Select}; use crate::table::datafusion::BaseTableAdapter; -use crate::table::schema_evolution::FieldMetadataUpdate; use crate::table::{AddDataMode, Table, WriteOptions}; const GEN_OUT: &str = "gen_out"; @@ -108,12 +107,15 @@ async fn create_table_with_complete_literal_generated(name: &str) -> Fixture { INITIAL_MATERIALIZED_EPOCH, ); let json = definition.to_metadata_json().unwrap(); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, json) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_OUT, + json, + ) + .await + .unwrap(); assert_eq!( table.generated_column_status(GEN_OUT).await.unwrap(), @@ -704,13 +706,16 @@ async fn malformed_generated_metadata_rejects_append_before_mutation_and_redacts r#"{{"format_version":1,"output_field_id":{field_id},"function_call":{MALFORMED_MARKER},"dependency_epoch":1,"materialized_epoch":1}}"# ); assert!(raw.contains(MALFORMED_MARKER)); - fixture - .table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, raw.clone()) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_OUT, + raw.clone(), + ) + .await + .unwrap(); let version_before = fixture.table.version().await.unwrap(); let rows_before = ordinary_values(&fixture.table).await; diff --git a/rust/lancedb/src/table/delete_generated_column_invalidation_contract.rs b/rust/lancedb/src/table/delete_generated_column_invalidation_contract.rs index 8beb90a5f..216f96341 100644 --- a/rust/lancedb/src/table/delete_generated_column_invalidation_contract.rs +++ b/rust/lancedb/src/table/delete_generated_column_invalidation_contract.rs @@ -27,7 +27,6 @@ use crate::function::{ }; use crate::query::{ExecutableQuery, QueryBase, Select}; use crate::table::Table; -use crate::table::schema_evolution::FieldMetadataUpdate; const ID: &str = "id"; const INPUT_A: &str = "input_a"; @@ -212,27 +211,23 @@ async fn create_row_set_generated_fixture(name: &str) -> Fixture { let literal = literal_only_definition(gen_literal_field_id); let literal_argument = literal_only_argument(); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_DIRECT).set( - GENERATED_COLUMN_METADATA_KEY, - direct.to_metadata_json().unwrap(), - ), - FieldMetadataUpdate::new(GEN_TRANSITIVE).set( - GENERATED_COLUMN_METADATA_KEY, - transitive.to_metadata_json().unwrap(), - ), - FieldMetadataUpdate::new(GEN_INDEPENDENT).set( - GENERATED_COLUMN_METADATA_KEY, - independent.to_metadata_json().unwrap(), - ), - FieldMetadataUpdate::new(GEN_LITERAL).set( - GENERATED_COLUMN_METADATA_KEY, - literal.to_metadata_json().unwrap(), - ), - ]) + let native = table + .as_native() + .expect("generated-column fixture planting requires a Native table"); + for (column, definition) in [ + (GEN_DIRECT, &direct), + (GEN_TRANSITIVE, &transitive), + (GEN_INDEPENDENT, &independent), + (GEN_LITERAL, &literal), + ] { + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + native, + column, + definition.to_metadata_json().unwrap(), + ) .await .unwrap(); + } let planted_direct = read_generated_definition(&table, GEN_DIRECT).await; let planted_transitive = read_generated_definition(&table, GEN_TRANSITIVE).await; @@ -944,12 +939,16 @@ async fn string_malformed_generated_metadata_rejects_before_mutation() { fixture.gen_independent_field_id, MALFORMED_MARKER ); assert!(planted_raw.contains(MALFORMED_MARKER)); - fixture - .table - .update_field_metadata(&[FieldMetadataUpdate::new(GEN_INDEPENDENT) - .set(GENERATED_COLUMN_METADATA_KEY, planted_raw.clone())]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_INDEPENDENT, + planted_raw.clone(), + ) + .await + .unwrap(); assert_eq!( read_raw_generated_metadata(&fixture.table, GEN_INDEPENDENT).await, planted_raw, @@ -1007,12 +1006,16 @@ async fn expr_malformed_generated_metadata_rejects_before_mutation() { fixture.gen_literal_field_id, MALFORMED_MARKER ); assert!(planted_raw.contains(MALFORMED_MARKER)); - fixture - .table - .update_field_metadata(&[FieldMetadataUpdate::new(GEN_LITERAL) - .set(GENERATED_COLUMN_METADATA_KEY, planted_raw.clone())]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_LITERAL, + planted_raw.clone(), + ) + .await + .unwrap(); assert_eq!( read_raw_generated_metadata(&fixture.table, GEN_LITERAL).await, planted_raw, diff --git a/rust/lancedb/src/table/merge_insert_generated_column_reject_contract.rs b/rust/lancedb/src/table/merge_insert_generated_column_reject_contract.rs index ec5f73fec..09cb642d1 100644 --- a/rust/lancedb/src/table/merge_insert_generated_column_reject_contract.rs +++ b/rust/lancedb/src/table/merge_insert_generated_column_reject_contract.rs @@ -25,7 +25,6 @@ use crate::function::{ }; use crate::query::{ExecutableQuery, QueryBase, Select}; use crate::table::Table; -use crate::table::schema_evolution::FieldMetadataUpdate; const ID: &str = "id"; const ORDINARY: &str = "ordinary"; @@ -193,13 +192,16 @@ async fn create_table_with_complete_literal_generated(name: &str) -> Fixture { let field_id = snapshot.field(GEN_OUT).expect(GEN_OUT).field_id(); let definition = literal_only_definition(field_id); let json = definition.to_metadata_json().unwrap(); - fixture - .table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, json) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_OUT, + json, + ) + .await + .unwrap(); assert_eq!( fixture .table @@ -426,12 +428,16 @@ async fn malformed_generated_metadata_rejects_merge_insert_before_mutation_and_r ); assert!(planted_raw.contains(MALFORMED_MARKER)); assert!(planted_raw.contains(FN_ID)); - fixture - .table - .update_field_metadata(&[FieldMetadataUpdate::new(GEN_OUT) - .set(GENERATED_COLUMN_METADATA_KEY, planted_raw.clone())]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_OUT, + planted_raw.clone(), + ) + .await + .unwrap(); assert_eq!( read_raw_generated_metadata(&fixture.table).await, planted_raw, diff --git a/rust/lancedb/src/table/query.rs b/rust/lancedb/src/table/query.rs index 14ee51707..28ab8b812 100644 --- a/rust/lancedb/src/table/query.rs +++ b/rust/lancedb/src/table/query.rs @@ -1327,10 +1327,8 @@ mod tests { ) { use crate::function::{ Function, FunctionArgument, FunctionCall, FunctionId, FunctionOutput, - FunctionParameter, FunctionSignature, GENERATED_COLUMN_METADATA_KEY, - GeneratedColumnDefinition, + FunctionParameter, FunctionSignature, GeneratedColumnDefinition, }; - use crate::table::FieldMetadataUpdate; use arrow_array::StringArray; let snapshot = table.generated_column_binding_snapshot().await.unwrap(); @@ -1363,12 +1361,15 @@ mod tests { .unwrap() .to_metadata_json() .unwrap(); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new(column).set(GENERATED_COLUMN_METADATA_KEY, json) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + column, + json, + ) + .await + .unwrap(); } fn assert_incomplete_runtime_error(err: &Error, label: &str) { @@ -1470,7 +1471,6 @@ mod tests { async fn generated_column_query_runtime_malformed_rejects_before_namespace() { use crate::function::GENERATED_COLUMN_METADATA_KEY; use crate::query::Select; - use crate::table::FieldMetadataUpdate; let (table, native_table, namespace_client) = runtime_namespace_table("runtime_ns_malformed").await; @@ -1479,12 +1479,15 @@ mod tests { let raw = format!( r#"{{"format_version":1,"output_field_id":0,"function_call":{MARKER},"dependency_epoch":1,"materialized_epoch":1}}"# ); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new("gen_out").set(GENERATED_COLUMN_METADATA_KEY, raw.clone()) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + "gen_out", + raw.clone(), + ) + .await + .unwrap(); let query = AnyQuery::Query(QueryRequest { select: Select::Columns(vec!["gen_out".to_string()]), diff --git a/rust/lancedb/src/table/schema_evolution.rs b/rust/lancedb/src/table/schema_evolution.rs index ce208111a..322cc5724 100644 --- a/rust/lancedb/src/table/schema_evolution.rs +++ b/rust/lancedb/src/table/schema_evolution.rs @@ -13,7 +13,61 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; use super::NativeTable; -use crate::Result; +use crate::function::GENERATED_COLUMN_METADATA_KEY; +use crate::{Error, Result}; + +/// Shared rejection for general-purpose field-metadata updates that name the +/// reserved generated-column definition key. +/// +/// Generated-column definitions are table-schema state owned by +/// create/change/refresh Jobs. The public `update_field_metadata` API must not +/// create, replace, or remove that reserved key. Both Native and Remote +/// `BaseTable` implementations call this helper so direct trait calls cannot +/// bypass the syntax guard. +pub(crate) fn reject_reserved_generated_column_metadata_key_updates( + updates: &[FieldMetadataUpdate], +) -> Result<()> { + for update in updates { + if update.metadata.contains_key(GENERATED_COLUMN_METADATA_KEY) { + return Err(reserved_generated_column_metadata_not_supported()); + } + } + Ok(()) +} + +fn reserved_generated_column_metadata_not_supported() -> Error { + Error::NotSupported { + message: "generated column definitions are owned by create/change/refresh Jobs \ + and cannot be created, replaced, or removed through update_field_metadata" + .into(), + } +} + +/// Native-only state-aware guard: whole-map `replace()` on a field whose exact +/// Dataset snapshot metadata already contains the reserved generated-column key +/// would wipe that Job-owned definition even when the replacement map omits the +/// key. Detects raw key presence without decoding the payload. +fn reject_replace_that_would_remove_generated_column_metadata( + dataset: &lance::Dataset, + updates: &[FieldMetadataUpdate], +) -> Result<()> { + let schema = dataset.schema(); + for update in updates { + if !update.replace { + continue; + } + let Some(fields) = schema.resolve_case_insensitive(&update.path) else { + continue; + }; + let field = fields + .last() + .expect("resolve_case_insensitive returns a non-empty path"); + if field.metadata.contains_key(GENERATED_COLUMN_METADATA_KEY) { + return Err(reserved_generated_column_metadata_not_supported()); + } + } + Ok(()) +} /// The result of an add columns operation. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] @@ -144,8 +198,10 @@ pub(crate) async fn execute_update_field_metadata( table: &NativeTable, updates: &[FieldMetadataUpdate], ) -> Result { + reject_reserved_generated_column_metadata_key_updates(updates)?; table.dataset.ensure_mutable()?; let mut dataset = (*table.dataset.get().await?).clone(); + reject_replace_that_would_remove_generated_column_metadata(&dataset, updates)?; let mut builder = dataset.update_field_metadata(); for update in updates { @@ -163,6 +219,33 @@ pub(crate) async fn execute_update_field_metadata( Ok(UpdateFieldMetadataResult { version }) } +/// Test-only raw installer for generated-column field metadata on Native tables. +/// +/// Uses the Lance metadata commit path directly so contract fixtures can plant +/// reserved-key bytes without going through the public `update_field_metadata` +/// guard. Absent from non-test builds. +#[cfg(test)] +pub(crate) async fn install_raw_generated_column_metadata_for_tests( + table: &NativeTable, + path: impl AsRef, + raw: impl Into, +) -> Result { + table.dataset.ensure_mutable()?; + let mut dataset = (*table.dataset.get().await?).clone(); + let path = path.as_ref(); + let raw = raw.into(); + dataset + .update_field_metadata() + .update( + path, + [(GENERATED_COLUMN_METADATA_KEY.to_string(), Some(raw))], + )? + .await?; + let version = dataset.version().version; + table.dataset.update(dataset); + Ok(UpdateFieldMetadataResult { version }) +} + #[cfg(test)] mod tests { use arrow_array::{Int32Array, StringArray, record_batch}; diff --git a/rust/lancedb/src/table/update_field_metadata_generated_column_guard_contract.rs b/rust/lancedb/src/table/update_field_metadata_generated_column_guard_contract.rs new file mode 100644 index 000000000..4801d9f67 --- /dev/null +++ b/rust/lancedb/src/table/update_field_metadata_generated_column_guard_contract.rs @@ -0,0 +1,627 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Runtime contract tests for the B4f reserved generated-metadata update guard. +//! +//! Pins that the general-purpose [`crate::table::Table::update_field_metadata`] +//! API cannot create, replace, or remove `GENERATED_COLUMN_METADATA_KEY`, and +//! that Native `replace()` cannot wipe an existing generated definition by +//! omitting the reserved key. Remote explicit-key attempts must reject before +//! transport. + +use std::sync::Arc; + +use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use futures::TryStreamExt; +use tempfile::TempDir; + +use crate::connection::ConnectBuilder; +use crate::error::Error; +use crate::function::{ + Function, FunctionArgument, FunctionCall, FunctionId, FunctionOutput, FunctionParameter, + FunctionSignature, GENERATED_COLUMN_METADATA_KEY, GeneratedColumnDefinition, +}; +use crate::query::{ExecutableQuery, QueryBase, Select}; +use crate::table::Table; +use crate::table::schema_evolution::FieldMetadataUpdate; + +const GEN_OUT: &str = "gen_out"; +const ORDINARY: &str = "ordinary"; +const CATEGORY: &str = "category"; +const FN_ID: &str = "fn.exact.b4f.guard.literal"; +const MALFORMED_MARKER: &str = "SENSITIVE_B4F_GUARD_METADATA_MARKER_7c1e_d04b"; + +struct Fixture { + _tmp: TempDir, + table: Table, + uri: String, +} + +fn literal_definition( + output_field_id: i32, + dependency_epoch: u64, + materialized_epoch: u64, +) -> GeneratedColumnDefinition { + let function = Function::new( + FunctionId::try_new(FN_ID).unwrap(), + FunctionSignature::try_new( + vec![FunctionParameter::new("label", DataType::Utf8)], + FunctionOutput::new(DataType::Int32, true), + ) + .unwrap(), + ); + let call = FunctionCall::try_new( + &function, + vec![( + "label".to_string(), + FunctionArgument::try_literal( + Arc::new(StringArray::from(vec![Some("b4f-guard")])) as arrow_array::ArrayRef + ) + .unwrap(), + )], + ) + .unwrap(); + GeneratedColumnDefinition::try_new(output_field_id, call, dependency_epoch, materialized_epoch) + .unwrap() +} + +async fn create_ordinary_table(name: &str) -> Fixture { + let tmp = tempfile::tempdir().unwrap(); + let uri = tmp.path().to_str().unwrap().to_string(); + let conn = ConnectBuilder::new(&uri).execute().await.unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new(GEN_OUT, DataType::Int32, true), + Field::new(ORDINARY, DataType::Utf8, true), + Field::new(CATEGORY, DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some("seed")])), + Arc::new(StringArray::from(vec![Some("A")])), + ], + ) + .unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + Fixture { + _tmp: tmp, + table, + uri, + } +} + +async fn plant_generated_raw(table: &Table, column: &str, raw: String) { + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + column, + raw, + ) + .await + .expect("fixture raw generated-column metadata install must succeed"); +} + +async fn plant_valid_generated(table: &Table, column: &str) -> String { + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + let field_id = snapshot.field(column).expect(column).field_id(); + let raw = literal_definition(field_id, 3, 3) + .to_metadata_json() + .unwrap(); + plant_generated_raw(table, column, raw.clone()).await; + raw +} + +async fn read_raw_generated_metadata(table: &Table, column: &str) -> Option { + let snapshot = table.generated_column_binding_snapshot().await.unwrap(); + snapshot + .field(column) + .expect(column) + .field() + .metadata() + .get(GENERATED_COLUMN_METADATA_KEY) + .cloned() +} + +async fn ordinary_values(table: &Table) -> Vec { + let batches = table + .query() + .select(Select::columns(&[ORDINARY])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.unwrap().to_string()) + }) + .collect() +} + +async fn reopen(uri: &str, name: &str) -> Table { + ConnectBuilder::new(uri) + .execute() + .await + .unwrap() + .open_table(name) + .execute() + .await + .unwrap() +} + +fn assert_not_supported_redacted(err: &Error, label: &str, forbidden_substrings: &[&str]) { + match err { + Error::NotSupported { message } => { + let rendered = format!("{err}\n{err:?}\n{message}"); + assert!( + !rendered.contains(GENERATED_COLUMN_METADATA_KEY), + "{label}: leaked metadata wire key: {rendered}" + ); + assert!( + !rendered.contains(FN_ID), + "{label}: leaked Function ID: {rendered}" + ); + assert!( + !rendered.contains(MALFORMED_MARKER), + "{label}: leaked malformed marker: {rendered}" + ); + for needle in forbidden_substrings { + assert!( + !rendered.contains(needle), + "{label}: leaked forbidden substring `{needle}`: {rendered}" + ); + } + assert!( + message.to_lowercase().contains("generated") + || message.to_lowercase().contains("job"), + "{label}: message must describe Job-owned generated-column boundary: {message}" + ); + } + other => panic!("{label}: expected Error::NotSupported, got {other:?}"), + } +} + +#[tokio::test] +async fn native_ordinary_field_explicit_reserved_key_set_rejects_and_preserves_state() { + let fixture = create_ordinary_table("b4f_ordinary_set").await; + let version_before = fixture.table.version().await.unwrap(); + let rows_before = ordinary_values(&fixture.table).await; + let schema_before = fixture.table.schema().await.unwrap(); + let category_md_before = schema_before + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .clone(); + + let snapshot = fixture + .table + .generated_column_binding_snapshot() + .await + .unwrap(); + let field_id = snapshot.field(CATEGORY).expect(CATEGORY).field_id(); + let payload = literal_definition(field_id, 1, 1) + .to_metadata_json() + .unwrap(); + + let err = fixture + .table + .update_field_metadata(&[ + FieldMetadataUpdate::new(CATEGORY).set(GENERATED_COLUMN_METADATA_KEY, payload.clone()) + ]) + .await + .expect_err("explicit reserved-key set on ordinary field must reject"); + assert_not_supported_redacted(&err, "ordinary reserved set", &[&payload]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!(ordinary_values(&fixture.table).await, rows_before); + let schema_after = fixture.table.schema().await.unwrap(); + assert_eq!( + schema_after.field_with_name(CATEGORY).unwrap().metadata(), + &category_md_before + ); + assert!( + read_raw_generated_metadata(&fixture.table, CATEGORY) + .await + .is_none() + ); +} + +#[tokio::test] +async fn native_generated_field_explicit_remove_rejects_and_preserves_raw() { + let fixture = create_ordinary_table("b4f_gen_remove").await; + let planted = plant_valid_generated(&fixture.table, GEN_OUT).await; + let version_before = fixture.table.version().await.unwrap(); + + let err = fixture + .table + .update_field_metadata(&[ + FieldMetadataUpdate::new(GEN_OUT).remove(GENERATED_COLUMN_METADATA_KEY) + ]) + .await + .expect_err("explicit reserved-key remove must reject"); + assert_not_supported_redacted(&err, "generated remove", &[&planted]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + assert_eq!( + read_raw_generated_metadata(&fixture.table, GEN_OUT) + .await + .as_deref(), + Some(planted.as_str()) + ); + let fresh = reopen(&fixture.uri, "b4f_gen_remove").await; + assert_eq!( + read_raw_generated_metadata(&fresh, GEN_OUT) + .await + .as_deref(), + Some(planted.as_str()) + ); +} + +#[tokio::test] +async fn native_generated_field_explicit_replacement_rejects_and_preserves_raw() { + let fixture = create_ordinary_table("b4f_gen_replace_value").await; + let planted = plant_valid_generated(&fixture.table, GEN_OUT).await; + let version_before = fixture.table.version().await.unwrap(); + let replacement = literal_definition( + fixture + .table + .generated_column_binding_snapshot() + .await + .unwrap() + .field(GEN_OUT) + .unwrap() + .field_id(), + 9, + 1, + ) + .to_metadata_json() + .unwrap(); + assert_ne!(replacement, planted); + + let err = fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(GEN_OUT) + .set(GENERATED_COLUMN_METADATA_KEY, replacement.clone())]) + .await + .expect_err("explicit reserved-key replacement must reject"); + assert_not_supported_redacted(&err, "generated replace value", &[&planted, &replacement]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + let fresh = reopen(&fixture.uri, "b4f_gen_replace_value").await; + assert_eq!( + read_raw_generated_metadata(&fresh, GEN_OUT) + .await + .as_deref(), + Some(planted.as_str()) + ); +} + +#[tokio::test] +async fn native_generated_field_replace_with_ordinary_metadata_rejects_and_preserves_raw() { + let fixture = create_ordinary_table("b4f_gen_replace_map").await; + let planted = plant_valid_generated(&fixture.table, GEN_OUT).await; + let version_before = fixture.table.version().await.unwrap(); + + let err = fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(GEN_OUT) + .replace() + .set("unit", "label")]) + .await + .expect_err("replace() that would wipe generated definition must reject"); + assert_not_supported_redacted(&err, "generated replace map", &[&planted]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + let fresh = reopen(&fixture.uri, "b4f_gen_replace_map").await; + assert_eq!( + read_raw_generated_metadata(&fresh, GEN_OUT) + .await + .as_deref(), + Some(planted.as_str()) + ); +} + +#[tokio::test] +async fn native_mixed_batch_rejects_atomically_no_partial_commit() { + let fixture = create_ordinary_table("b4f_mixed_batch").await; + let planted = plant_valid_generated(&fixture.table, GEN_OUT).await; + let version_before = fixture.table.version().await.unwrap(); + + let err = fixture + .table + .update_field_metadata(&[ + FieldMetadataUpdate::new(CATEGORY).set("unit", "label"), + FieldMetadataUpdate::new(GEN_OUT).remove(GENERATED_COLUMN_METADATA_KEY), + ]) + .await + .expect_err("mixed batch with forbidden update must reject all-or-none"); + assert_not_supported_redacted(&err, "mixed forbidden second", &[&planted]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + let schema = fixture.table.schema().await.unwrap(); + assert!( + !schema + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .contains_key("unit"), + "ordinary metadata must not partially commit" + ); + + let err = fixture + .table + .update_field_metadata(&[ + FieldMetadataUpdate::new(GEN_OUT).set(GENERATED_COLUMN_METADATA_KEY, planted.clone()), + FieldMetadataUpdate::new(CATEGORY).set("unit", "label"), + ]) + .await + .expect_err("mixed batch with forbidden update first must reject all-or-none"); + assert_not_supported_redacted(&err, "mixed forbidden first", &[&planted]); + + let fresh = reopen(&fixture.uri, "b4f_mixed_batch").await; + assert_eq!(fresh.version().await.unwrap(), version_before); + assert_eq!( + read_raw_generated_metadata(&fresh, GEN_OUT) + .await + .as_deref(), + Some(planted.as_str()) + ); + let fresh_schema = fresh.schema().await.unwrap(); + assert!( + !fresh_schema + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .contains_key("unit") + ); +} + +#[tokio::test] +async fn native_malformed_generated_raw_replace_rejects_redacted_and_preserves_raw() { + let fixture = create_ordinary_table("b4f_malformed_replace").await; + let field_id = fixture + .table + .generated_column_binding_snapshot() + .await + .unwrap() + .field(GEN_OUT) + .unwrap() + .field_id(); + let planted_raw = format!( + r#"{{"format_version":1,"output_field_id":{field_id},"function_call":{MALFORMED_MARKER},"dependency_epoch":1,"materialized_epoch":1}}"# + ); + assert!(planted_raw.contains(MALFORMED_MARKER)); + plant_generated_raw(&fixture.table, GEN_OUT, planted_raw.clone()).await; + let version_before = fixture.table.version().await.unwrap(); + + let err = fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(GEN_OUT) + .replace() + .set("unit", "label")]) + .await + .expect_err("malformed generated raw must still block replace()"); + assert_not_supported_redacted(&err, "malformed replace", &[&planted_raw]); + + assert_eq!(fixture.table.version().await.unwrap(), version_before); + let fresh = reopen(&fixture.uri, "b4f_malformed_replace").await; + assert_eq!( + read_raw_generated_metadata(&fresh, GEN_OUT) + .await + .as_deref(), + Some(planted_raw.as_str()) + ); +} + +#[tokio::test] +async fn native_ordinary_metadata_merge_set_remove_replace_still_work() { + let fixture = create_ordinary_table("b4f_ordinary_controls").await; + + fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(CATEGORY) + .set("unit", "label") + .set("pii", "false")]) + .await + .unwrap(); + let md = fixture + .table + .schema() + .await + .unwrap() + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .clone(); + assert_eq!(md.get("unit").map(String::as_str), Some("label")); + assert_eq!(md.get("pii").map(String::as_str), Some("false")); + + fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(CATEGORY) + .set("source", "import") + .remove("pii")]) + .await + .unwrap(); + let md = fixture + .table + .schema() + .await + .unwrap() + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .clone(); + assert_eq!(md.get("unit").map(String::as_str), Some("label")); + assert_eq!(md.get("source").map(String::as_str), Some("import")); + assert!(!md.contains_key("pii")); + + fixture + .table + .update_field_metadata(&[FieldMetadataUpdate::new(CATEGORY) + .replace() + .set("only", "kept")]) + .await + .unwrap(); + let md = fixture + .table + .schema() + .await + .unwrap() + .field_with_name(CATEGORY) + .unwrap() + .metadata() + .clone(); + assert_eq!(md.len(), 1); + assert_eq!(md.get("only").map(String::as_str), Some("kept")); + assert!( + read_raw_generated_metadata(&fixture.table, CATEGORY) + .await + .is_none() + ); +} + +#[cfg(feature = "remote")] +mod remote_explicit_key_guard { + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use async_trait::async_trait; + + use super::*; + use crate::Error; + use crate::remote::{ClientConfig, HeaderProvider}; + + #[derive(Debug)] + struct CountingHeaderProvider { + calls: Arc, + } + + #[async_trait] + impl HeaderProvider for CountingHeaderProvider { + async fn get_headers(&self) -> crate::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(HashMap::from([( + "X-Test-Header".to_string(), + "must-not-be-requested".to_string(), + )])) + } + } + + fn panic_handler( + calls: Arc, + ) -> impl Fn(reqwest::Request) -> http::Response + Clone + Send + Sync + 'static { + move |_request| { + calls.fetch_add(1, Ordering::SeqCst); + panic!("remote reserved-key update must not invoke the HTTP handler"); + } + } + + #[tokio::test] + async fn remote_explicit_set_rejects_before_handler_and_header_provider() { + let handler_calls = Arc::new(AtomicUsize::new(0)); + let header_calls = Arc::new(AtomicUsize::new(0)); + let config = ClientConfig { + header_provider: Some(Arc::new(CountingHeaderProvider { + calls: header_calls.clone(), + }) as Arc), + ..Default::default() + }; + let table = Table::new_with_handler_and_config( + "my_table", + panic_handler(handler_calls.clone()), + config, + ); + + let err = table + .update_field_metadata(&[FieldMetadataUpdate::new(CATEGORY) + .set(GENERATED_COLUMN_METADATA_KEY, r#"{"format_version":1}"#)]) + .await + .expect_err("remote explicit reserved-key set must reject"); + assert!( + matches!(err, Error::NotSupported { .. }), + "expected NotSupported, got {err:?}" + ); + assert_eq!(handler_calls.load(Ordering::SeqCst), 0); + assert_eq!(header_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn remote_explicit_remove_rejects_before_handler_and_header_provider() { + let handler_calls = Arc::new(AtomicUsize::new(0)); + let header_calls = Arc::new(AtomicUsize::new(0)); + let config = ClientConfig { + header_provider: Some(Arc::new(CountingHeaderProvider { + calls: header_calls.clone(), + }) as Arc), + ..Default::default() + }; + let table = Table::new_with_handler_and_config( + "my_table", + panic_handler(handler_calls.clone()), + config, + ); + + let err = table + .update_field_metadata(&[ + FieldMetadataUpdate::new(CATEGORY).remove(GENERATED_COLUMN_METADATA_KEY) + ]) + .await + .expect_err("remote explicit reserved-key remove must reject"); + assert!( + matches!(err, Error::NotSupported { .. }), + "expected NotSupported, got {err:?}" + ); + assert_eq!(handler_calls.load(Ordering::SeqCst), 0); + assert_eq!(header_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn remote_ordinary_metadata_update_sends_exact_body_and_succeeds() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!( + request.url().path(), + "/v1/table/my_table/update_field_metadata/" + ); + let body = request + .body() + .expect("ordinary update must send a body") + .as_bytes() + .expect("body is in-memory"); + let parsed: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!( + parsed, + serde_json::json!({ + "updates": [{ + "path": "category", + "metadata": { "unit": "label" }, + "replace": false + }] + }) + ); + http::Response::builder() + .status(200) + .body(r#"{"version": 7}"#.to_string()) + .unwrap() + }); + + let result = table + .update_field_metadata(&[FieldMetadataUpdate::new(CATEGORY).set("unit", "label")]) + .await + .unwrap(); + assert_eq!(result.version, 7); + } +} diff --git a/rust/lancedb/src/table/update_generated_column_invalidation_contract.rs b/rust/lancedb/src/table/update_generated_column_invalidation_contract.rs index 6165c1266..c89f27621 100644 --- a/rust/lancedb/src/table/update_generated_column_invalidation_contract.rs +++ b/rust/lancedb/src/table/update_generated_column_invalidation_contract.rs @@ -25,7 +25,6 @@ use crate::function::{ }; use crate::query::{ExecutableQuery, QueryBase, Select}; use crate::table::Table; -use crate::table::schema_evolution::FieldMetadataUpdate; const ID: &str = "id"; const INPUT_A: &str = "input_a"; @@ -162,23 +161,22 @@ async fn create_dependent_generated_fixture(name: &str) -> Fixture { input_b_field_id, ); - table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_DIRECT).set( - GENERATED_COLUMN_METADATA_KEY, - direct.to_metadata_json().unwrap(), - ), - FieldMetadataUpdate::new(GEN_TRANSITIVE).set( - GENERATED_COLUMN_METADATA_KEY, - transitive.to_metadata_json().unwrap(), - ), - FieldMetadataUpdate::new(GEN_UNRELATED).set( - GENERATED_COLUMN_METADATA_KEY, - unrelated.to_metadata_json().unwrap(), - ), - ]) + let native = table + .as_native() + .expect("generated-column fixture planting requires a Native table"); + for (column, definition) in [ + (GEN_DIRECT, &direct), + (GEN_TRANSITIVE, &transitive), + (GEN_UNRELATED, &unrelated), + ] { + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + native, + column, + definition.to_metadata_json().unwrap(), + ) .await .unwrap(); + } let planted_direct = read_generated_definition(&table, GEN_DIRECT).await; let planted_transitive = read_generated_definition(&table, GEN_TRANSITIVE).await; @@ -898,13 +896,16 @@ async fn unrelated_update_rejects_malformed_generated_metadata_before_mutation() fixture.gen_unrelated_field_id, MALFORMED_MARKER ); assert!(raw.contains(MALFORMED_MARKER)); - fixture - .table - .update_field_metadata(&[ - FieldMetadataUpdate::new(GEN_UNRELATED).set(GENERATED_COLUMN_METADATA_KEY, raw) - ]) - .await - .unwrap(); + crate::table::schema_evolution::install_raw_generated_column_metadata_for_tests( + fixture + .table + .as_native() + .expect("generated-column fixture planting requires a Native table"), + GEN_UNRELATED, + raw, + ) + .await + .unwrap(); let version_before = fixture.table.version().await.unwrap(); let rows_before = safe_row_projection(&fixture.table).await;