Merge pull request #34140 from bkirwi/batch-match

[persist] Check schemas in batch append

Author: bkirwi

src/persist-client/build.rs

@@ -52,6 +52,7 @@ fn main() {
         .type_attribute(".", "#[allow(missing_docs)]")
         .btree_map(["."])
         .bytes([
+            ".mz_persist_client.batch.ProtoBatch",
             ".mz_persist_client.internal.diff.ProtoStateFieldDiffs",
             ".mz_persist_client.internal.service.ProtoPushDiff",
             ".mz_persist_client.internal.state.ProtoEncodedSchemas",

src/persist-client/src/batch.proto

@@ -21,6 +21,9 @@ import "persist-client/src/internal/state.proto";
 message ProtoBatch {
   string shard_id = 1;
   string version = 2;
+  bytes key_schema = 5;
+  bytes val_schema = 6;
+
   mz_persist_client.internal.state.ProtoHollowBatch batch = 3;
 
   reserved 4;

src/persist-client/src/batch.rs

@@ -81,6 +81,9 @@ pub struct Batch<K, V, T, D> {
     /// The version of Materialize which wrote this batch.
     pub(crate) version: Version,
 
+    /// The encoded schemas of the data in the batch.
+    pub(crate) schemas: (Bytes, Bytes),
+
     /// A handle to the data represented by this batch.
     pub(crate) batch: HollowBatch<T>,
 
@@ -121,13 +124,15 @@ where
         blob: Arc<dyn Blob>,
         shard_metrics: Arc<ShardMetrics>,
         version: Version,
+        schemas: (Bytes, Bytes),
         batch: HollowBatch<T>,
     ) -> Self {
         Self {
             batch_delete_enabled,
             metrics,
             shard_metrics,
             version,
+            schemas,
             batch,
             blob,
             _phantom: PhantomData,
@@ -209,6 +214,8 @@ where
             shard_id: self.shard_metrics.shard_id.into_proto(),
             version: self.version.to_string(),
             batch: Some(self.batch.into_proto()),
+            key_schema: self.schemas.0.clone(),
+            val_schema: self.schemas.1.clone(),
         };
         self.mark_consumed();
         ret
@@ -729,6 +736,10 @@ where
             self.blob,
             shard_metrics,
             self.version,
+            (
+                K::encode_schema(&*self.write_schemas.key),
+                V::encode_schema(&*self.write_schemas.val),
+            ),
             HollowBatch::new(desc, run_parts, total_updates, run_meta, run_splits),
         );
 

src/persist-client/src/internal/machine.rs

@@ -1538,6 +1538,10 @@ pub mod datadriven {
                 Arc::clone(&self.client.blob),
                 self.client.metrics.shards.shard(&self.shard_id, "test"),
                 self.client.cfg.build_version.clone(),
+                (
+                    <String>::encode_schema(&*SCHEMAS.key),
+                    <()>::encode_schema(&*SCHEMAS.val),
+                ),
                 hollow,
             )
         }

src/persist-client/src/lib.rs

@@ -603,6 +603,7 @@ impl PersistClient {
             metrics: Arc::clone(&self.metrics),
             shard_metrics,
             version: Version::parse(&batch.version).expect("valid transmittable batch"),
+            schemas: (batch.key_schema, batch.val_schema),
             batch: batch
                 .batch
                 .into_rust_if_some("ProtoBatch::batch")

src/persist-client/src/write.rs

@@ -19,8 +19,8 @@ use differential_dataflow::trace::Description;
 use futures::StreamExt;
 use futures::stream::FuturesUnordered;
 use mz_dyncfg::Config;
-use mz_ore::instrument;
 use mz_ore::task::RuntimeExt;
+use mz_ore::{instrument, soft_panic_or_log};
 use mz_persist::location::Blob;
 use mz_persist_types::schema::SchemaId;
 use mz_persist_types::{Codec, Codec64};
@@ -32,7 +32,7 @@ use timely::PartialOrder;
 use timely::order::TotalOrder;
 use timely::progress::{Antichain, Timestamp};
 use tokio::runtime::Handle;
-use tracing::{Instrument, debug_span, info, warn};
+use tracing::{Instrument, debug_span, error, info, warn};
 use uuid::Uuid;
 
 use crate::batch::{
@@ -568,6 +568,23 @@ where
                     TODO: Error on very old versions once the leaked blob detector exists."
                 )
             }
+            fn assert_schema<A: Codec>(writer_schema: &A::Schema, batch_schema: &bytes::Bytes) {
+                if batch_schema.is_empty() {
+                    // Schema is either trivial or missing!
+                    return;
+                }
+                let batch_schema: A::Schema = A::decode_schema(batch_schema);
+                if *writer_schema != batch_schema {
+                    error!(
+                        ?writer_schema,
+                        ?batch_schema,
+                        "writer and batch schemas should be identical"
+                    );
+                    soft_panic_or_log!("writer and batch schemas should be identical");
+                }
+            }
+            assert_schema::<K>(&*self.write_schemas.key, &batch.schemas.0);
+            assert_schema::<V>(&*self.write_schemas.val, &batch.schemas.1);
         }
 
         let lower = expected_upper.clone();
@@ -822,6 +839,7 @@ where
             metrics: Arc::clone(&self.metrics),
             shard_metrics: Arc::clone(&self.machine.applier.shard_metrics),
             version: Version::parse(&batch.version).expect("valid transmittable batch"),
+            schemas: (batch.key_schema, batch.val_schema),
             batch: batch
                 .batch
                 .into_rust_if_some("ProtoBatch::batch")