fix: enforce manifest capability checks before side effects

Author: Xuanwo

rust/lance-namespace-impls/src/dir/manifest.rs

@@ -1933,6 +1933,7 @@ impl ManifestNamespace {
     /// concurrent upgrade in between is still caught.
     async fn ensure_manifest_writable(&self) -> Result<()> {
         let dataset_guard = self.manifest_dataset.get().await?;
+        ensure_can_write_manifest(dataset_guard.manifest())?;
         ensure_writable(dataset_guard.metadata())
     }
 
@@ -1953,10 +1954,11 @@ impl ManifestNamespace {
 
         loop {
             let dataset_guard = self.manifest_dataset.get_refreshed().await?;
+            ensure_can_write_manifest(dataset_guard.manifest())?;
             let dataset = Arc::new(dataset_guard.clone());
             drop(dataset_guard);
-            // Refuse to mutate a manifest written with a writer feature flag this
-            // build does not understand.
+            // The namespace format has its own capabilities in table metadata,
+            // separate from the Lance manifest capabilities checked above.
             ensure_writable(dataset.metadata())?;
             // Staged files, indices, the commit, and cleanup must all use the dataset's
             // own object store (see `commit_manifest_overwrite`).
@@ -3859,6 +3861,7 @@ mod tests {
         CreateNamespaceRequest, CreateTableRequest, DescribeTableRequest, DropTableRequest,
         ListTablesRequest, TableExistsRequest,
     };
+    use lance_table::feature_flags::FLAG_UNKNOWN;
     use lance_table::format::Fragment;
     use rstest::rstest;
     use std::collections::{HashMap, HashSet};
@@ -4391,6 +4394,64 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_manifest_rewrite_rejects_unknown_writer_flag_before_staging() {
+        let temp_dir = TempStdDir::default();
+        let temp_path = temp_dir.to_str().unwrap();
+        let manifest_ns = create_manifest_namespace(temp_path, false).await;
+        let data_paths_before = manifest_data_paths(&manifest_ns).await;
+        let original_version = {
+            let mut dataset = manifest_ns.manifest_dataset.get_mut().await.unwrap();
+            let mut manifest = dataset.manifest().clone();
+            manifest.writer_feature_flags |= FLAG_UNKNOWN << 1;
+            let version = manifest.version;
+            dataset.manifest = Arc::new(manifest);
+            version
+        };
+
+        let entries_before = dir_entry_names(temp_path);
+        let mut create_request = CreateTableRequest::new();
+        create_request.id = Some(vec!["new_table".to_string()]);
+        let error = manifest_ns
+            .create_table(create_request, Bytes::from(create_test_ipc_data()))
+            .await
+            .unwrap_err();
+        assert!(
+            error.to_string().to_lowercase().contains("upgrade"),
+            "expected an upgrade error, got: {error}"
+        );
+        assert_eq!(dir_entry_names(temp_path), entries_before);
+
+        let error = manifest_ns
+            .insert_into_manifest_with_metadata(
+                vec![ManifestEntry {
+                    object_id: "table".to_string(),
+                    object_type: ObjectType::Table,
+                    location: Some("table.lance".to_string()),
+                    metadata: None,
+                }],
+                None,
+            )
+            .await
+            .unwrap_err();
+
+        assert!(
+            error.to_string().to_lowercase().contains("upgrade"),
+            "expected an upgrade error, got: {error}"
+        );
+        assert_eq!(
+            manifest_ns
+                .manifest_dataset
+                .get()
+                .await
+                .unwrap()
+                .version()
+                .version,
+            original_version
+        );
+        assert_eq!(manifest_data_paths(&manifest_ns).await, data_paths_before);
+    }
+
     #[tokio::test]
     async fn test_manifest_noop_delete_uses_latest_snapshot() {
         let temp_dir = TempStdDir::default();

rust/lance/src/dataset.rs

@@ -149,7 +149,8 @@ use lance_core::box_error;
 use lance_index::scalar::lance_format::LanceIndexStore;
 use lance_namespace::models::{DeclareTableRequest, DescribeTableRequest};
 use lance_table::feature_flags::{
-    apply_feature_flags, can_read_dataset, validate_paired_feature_flags,
+    apply_feature_flags, ensure_can_read_manifest, ensure_can_write_manifest,
+    validate_paired_feature_flags,
 };
 use lance_table::io::deletion::{DELETIONS_DIR, relative_deletion_file_path};
 use lance_table::rowids::{RowIdSequence, write_row_ids};
@@ -766,16 +767,7 @@ impl Dataset {
             read_struct(object_reader.as_ref(), offset).await
         }?;
 
-        validate_paired_feature_flags(&manifest)?;
-
-        if !can_read_dataset(manifest.reader_feature_flags) {
-            let message = format!(
-                "This dataset cannot be read by this version of Lance. \
-                 Please upgrade Lance to read this dataset.\n Flags: {}",
-                manifest.reader_feature_flags
-            );
-            return Err(Error::not_supported_source(message.into()));
-        }
+        ensure_can_read_manifest(&manifest)?;
 
         // If indices were also in the last block, we can take the opportunity to
         // decode them now and cache them.
@@ -848,6 +840,7 @@ impl Dataset {
             e_tag: manifest_location.e_tag.as_deref(),
         };
         if let Some(cached) = metadata_cache.get_with_key(&manifest_key).await {
+            ensure_can_read_manifest(&cached)?;
             return Ok(cached);
         }
         let loaded =
@@ -1207,35 +1200,13 @@ impl Dataset {
             .resolve_latest_location(&self.base, &self.object_store)
             .await?;
 
-        // Check if manifest is in cache before reading from storage
-        let manifest_key = ManifestKey {
-            version: location.version,
-            e_tag: location.e_tag.as_deref(),
-        };
-        let cached_manifest = self.metadata_cache.get_with_key(&manifest_key).await;
-        if let Some(cached_manifest) = cached_manifest {
-            return Ok((cached_manifest, location));
-        }
-
         if self.already_checked_out(&location, self.manifest.branch.as_deref()) {
+            ensure_can_read_manifest(&self.manifest)?;
             return Ok((self.manifest.clone(), self.manifest_location.clone()));
         }
-        let mut manifest = read_manifest(&self.object_store, &location.path, location.size).await?;
-        if manifest.schema.has_dictionary_types() {
-            let reader = if let Some(size) = location.size {
-                self.object_store
-                    .open_with_size(&location.path, size as usize)
-                    .await?
-            } else {
-                self.object_store.open(&location.path).await?
-            };
-            populate_manifest_schema_dictionaries(&mut manifest, reader.as_ref()).await?;
-        }
-        let manifest_arc = Arc::new(manifest);
-        self.metadata_cache
-            .insert_with_key(&manifest_key, manifest_arc.clone())
-            .await;
-        Ok((manifest_arc, location))
+        let manifest =
+            Self::get_manifest(&self.object_store, &location, &self.uri, &self.session).await?;
+        Ok((manifest, location))
     }
 
     /// Read the transaction file for this version of the dataset.
@@ -3307,6 +3278,7 @@ impl Dataset {
 
         // Resolve source dataset and its manifest using checkout_version
         let src_ds = self.checkout_version(version).await?;
+        ensure_can_write_manifest(&src_ds.manifest)?;
         let src_paths = src_ds.collect_paths().await?;
 
         // Prepare target object store and base path

rust/lance/src/dataset/tests/dataset_io.rs

@@ -44,7 +44,7 @@ use lance_file::{
 };
 use lance_io::assert_io_eq;
 use lance_table::feature_flags;
-use lance_table::format::BasePath;
+use lance_table::format::{BasePath, Fragment};
 use object_store::ObjectStoreExt;
 
 use crate::index::DatasetIndexExt;
@@ -1412,6 +1412,51 @@ async fn test_restore_rejects_unknown_target_flags() {
     assert!(matches!(error, Error::NotSupported { .. }), "{error}");
 }
 
+#[tokio::test]
+async fn test_checkout_latest_rejects_unsupported_reader_before_caching() {
+    let test_uri = TempStrDir::default();
+    let data = gen_batch()
+        .col("i", array::step::<Int32Type>())
+        .into_reader_rows(RowCount::from(1), BatchCount::from(1));
+    let mut dataset = Dataset::write(data, &test_uri, None).await.unwrap();
+    let original_version = dataset.version().version;
+
+    let mut unsupported_manifest = dataset.manifest.as_ref().clone();
+    unsupported_manifest.version += 1;
+    unsupported_manifest.reader_feature_flags |= feature_flags::FLAG_UNKNOWN;
+    unsupported_manifest.writer_feature_flags |= feature_flags::FLAG_UNKNOWN;
+    let location = write_manifest_file(
+        dataset.object_store.as_ref(),
+        dataset.commit_handler.as_ref(),
+        &dataset.base,
+        &mut unsupported_manifest,
+        None,
+        &ManifestWriteConfig {
+            auto_set_feature_flags: false,
+            ..Default::default()
+        },
+        dataset.manifest_location.naming_scheme,
+        None,
+    )
+    .await
+    .unwrap();
+
+    let error = dataset.checkout_latest().await.unwrap_err();
+    assert!(matches!(error, Error::NotSupported { .. }), "{error}");
+    assert_eq!(dataset.version().version, original_version);
+    assert!(
+        dataset
+            .metadata_cache
+            .get_with_key(&ManifestKey {
+                version: location.version,
+                e_tag: location.e_tag.as_deref(),
+            })
+            .await
+            .is_none(),
+        "unsupported manifest must not be cached"
+    );
+}
+
 #[tokio::test]
 async fn test_rle_v2_v23_write_and_append() {
     let test_uri = TempStrDir::default();
@@ -1792,6 +1837,52 @@ async fn test_deep_clone(
     assert_eq!(count_files(store, &dst_root, "_deletions").await, 0);
 }
 
+#[tokio::test]
+async fn test_deep_clone_rejects_unsupported_writer_before_copying() {
+    let test_dir = TempStdDir::default();
+    let source_dir = test_dir.join("source");
+    let target_dir = test_dir.join("target");
+    let mut source = Dataset::write(
+        gen_batch()
+            .col("id", array::step::<Int32Type>())
+            .into_reader_rows(RowCount::from(32), BatchCount::from(1)),
+        source_dir.to_str().unwrap(),
+        None,
+    )
+    .await
+    .unwrap();
+
+    let mut unsupported_manifest = source.manifest.as_ref().clone();
+    unsupported_manifest.version += 1;
+    unsupported_manifest.writer_feature_flags |= feature_flags::FLAG_UNKNOWN << 1;
+    write_manifest_file(
+        source.object_store.as_ref(),
+        source.commit_handler.as_ref(),
+        &source.base,
+        &mut unsupported_manifest,
+        None,
+        &ManifestWriteConfig {
+            auto_set_feature_flags: false,
+            ..Default::default()
+        },
+        source.manifest_location.naming_scheme,
+        None,
+    )
+    .await
+    .unwrap();
+
+    let error = source
+        .deep_clone(
+            target_dir.to_str().unwrap(),
+            unsupported_manifest.version,
+            None,
+        )
+        .await
+        .unwrap_err();
+    assert!(matches!(error, Error::NotSupported { .. }));
+    assert!(!target_dir.exists());
+}
+
 #[tokio::test]
 async fn test_deep_clone_recognizes_ambiguous_commit_as_own() {
     use crate::utils::test::{AmbiguousCommitHandler, AmbiguousFailure};