fix: always include the Invariants writerFeature when upgrading to v7

The presence of a timestampntz column in CreateBuilder results in the
reader/writer being set to 3/7. The protocol reaquires that if the table
is writer v7, that `invariants` "must exist in the table protocl's
writerFeatures.

Sponsored-by: Scribd Inc
Signed-off-by: R. Tyler Croy <[email protected]>

Author: rtyler

crates/core/src/kernel/models/actions.rs

@@ -372,7 +372,10 @@ impl Protocol {
     /// Enable timestamp_ntz in the protocol
     pub fn enable_timestamp_ntz(mut self) -> Protocol {
         self = self.with_reader_features(vec![ReaderFeatures::TimestampWithoutTimezone]);
-        self = self.with_writer_features(vec![WriterFeatures::TimestampWithoutTimezone]);
+        self = self.with_writer_features(vec![
+            WriterFeatures::TimestampWithoutTimezone,
+            WriterFeatures::Invariants,
+        ]);
         self
     }
 }

crates/core/src/operations/create.rs

@@ -274,7 +274,9 @@ impl CreateBuilder {
             Protocol {
                 min_reader_version: 3,
                 min_writer_version: 7,
-                writer_features: Some(hashset! {WriterFeatures::TimestampWithoutTimezone}),
+                writer_features: Some(
+                    hashset! {WriterFeatures::TimestampWithoutTimezone, WriterFeatures::Invariants},
+                ),
                 reader_features: Some(hashset! {ReaderFeatures::TimestampWithoutTimezone}),
             }
         } else {
@@ -630,4 +632,43 @@ mod tests {
             .clone();
         assert_eq!(String::from("value"), value);
     }
+
+    #[tokio::test]
+    async fn test_invariants_when_timestampntz_exists() {
+        use crate::kernel::*;
+
+        let columns: Vec<StructField> = vec![
+            StructField::new("id", DataType::Primitive(PrimitiveType::Integer), false),
+            StructField::new(
+                "sunrise",
+                DataType::Primitive(PrimitiveType::TimestampNtz),
+                false,
+            ),
+        ];
+
+        let table = CreateBuilder::new()
+            .with_location("memory://")
+            .with_columns(columns)
+            .await
+            .expect("Failed to create the table");
+
+        // Ensure that the table can be created properly with the timestmapntz
+        assert_eq!(table.version(), 0);
+        let protocol = table.protocol().expect("Failed to load the protocol");
+
+        assert_eq!(protocol.min_reader_version, 3);
+        assert_eq!(protocol.min_writer_version, 7);
+
+        let writer_features = protocol.writer_features.as_ref().unwrap();
+        assert!(
+            writer_features.contains(&WriterFeatures::TimestampWithoutTimezone),
+            "The writer features must have TimestampeWithoutTimezone at writer:v7: {:#?}",
+            writer_features
+        );
+        assert!(
+            writer_features.contains(&WriterFeatures::Invariants),
+            "The writer features must have Invariants at writer:v7: {:#?}",
+            writer_features
+        );
+    }
 }

python/tests/test_alter.py

@@ -374,7 +374,7 @@ def test_add_timestamp_ntz_column(tmp_path: pathlib.Path, sample_table: pa.Table
     assert new_protocol.min_reader_version == 3
     assert new_protocol.min_writer_version == 7
     assert new_protocol.reader_features == ["timestampNtz"]
-    assert new_protocol.writer_features == ["timestampNtz"]
+    assert set(new_protocol.writer_features) == set(["timestampNtz", "invariants"])
 
 
 features = [

python/tests/test_checkpoint.py

@@ -203,7 +203,14 @@ def test_features_maintained_after_checkpoint(tmp_path: pathlib.Path):
     protocol_after_checkpoint = dt.protocol()
 
     assert protocol_after_checkpoint.reader_features == ["timestampNtz"]
-    assert current_protocol == protocol_after_checkpoint
+    assert (
+        protocol_after_checkpoint.min_reader_version
+        == current_protocol.min_reader_version
+    )
+    assert (
+        protocol_after_checkpoint.min_writer_version
+        == current_protocol.min_writer_version
+    )
 
 
 def test_features_null_on_below_v3_v7(tmp_path: pathlib.Path):

python/tests/test_create.py

@@ -30,7 +30,11 @@ def test_create_roundtrip_metadata(tmp_path: pathlib.Path, sample_data: pa.Table
     }
     assert dt.history()[0]["userName"] == "John Doe"
 
-    assert {*dt.protocol().writer_features} == {"appendOnly", "timestampNtz"}  # type: ignore
+    assert {*dt.protocol().writer_features} == {
+        "appendOnly",
+        "timestampNtz",
+        "invariants",
+    }  # type: ignore
 
 
 def test_create_modes(tmp_path: pathlib.Path, sample_data: pa.Table):

python/tests/test_writer.py

@@ -1760,7 +1760,7 @@ def test_write_timestamp_ntz_nested(tmp_path: pathlib.Path, array: pa.array):
     assert protocol.min_reader_version == 3
     assert protocol.min_writer_version == 7
     assert protocol.reader_features == ["timestampNtz"]
-    assert protocol.writer_features == ["timestampNtz"]
+    assert set(protocol.writer_features) == set(["invariants", "timestampNtz"])
 
 
 def test_write_timestamp_ntz_on_table_with_features_not_enabled(tmp_path: pathlib.Path):