From 4a881ef795fc20cc93fe3918250209d9a4cbdb78 Mon Sep 17 00:00:00 2001 From: "R. Tyler Croy" Date: Wed, 23 Oct 2024 00:13:57 +0000 Subject: [PATCH] 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 --- crates/core/src/kernel/models/actions.rs | 5 ++- crates/core/src/operations/create.rs | 43 +++++++++++++++++++++++- python/tests/test_alter.py | 2 +- python/tests/test_checkpoint.py | 9 ++++- python/tests/test_create.py | 6 +++- python/tests/test_writer.py | 2 +- 6 files changed, 61 insertions(+), 6 deletions(-) diff --git a/crates/core/src/kernel/models/actions.rs b/crates/core/src/kernel/models/actions.rs index bfcd9ec757..83a30a0436 100644 --- a/crates/core/src/kernel/models/actions.rs +++ b/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 } } diff --git a/crates/core/src/operations/create.rs b/crates/core/src/operations/create.rs index ad0413722e..eff01924de 100644 --- a/crates/core/src/operations/create.rs +++ b/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 = 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 + ); + } } diff --git a/python/tests/test_alter.py b/python/tests/test_alter.py index acc37db822..2ff87e0696 100644 --- a/python/tests/test_alter.py +++ b/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 = [ diff --git a/python/tests/test_checkpoint.py b/python/tests/test_checkpoint.py index 5ce6656463..232e54b60c 100644 --- a/python/tests/test_checkpoint.py +++ b/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): diff --git a/python/tests/test_create.py b/python/tests/test_create.py index 4bb73183fe..0fda3a968a 100644 --- a/python/tests/test_create.py +++ b/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): diff --git a/python/tests/test_writer.py b/python/tests/test_writer.py index a6662c48d6..e5d17819f0 100644 --- a/python/tests/test_writer.py +++ b/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):