feat!: add engine name to DatanodeTableValue (#2395)

* feat: add engine name to DatanodeTableValue

* fix: by cr

src/cmd/src/cli/upgrade.rs

@@ -386,6 +386,7 @@ impl MigrateTableMetadata {
 
     async fn create_datanode_table_keys(&self, value: &TableGlobalValue) {
         let table_id = value.table_id();
+        let engine = value.table_info.meta.engine.as_str();
         let region_distribution: RegionDistribution =
             value.regions_id_map.clone().into_iter().collect();
 
@@ -394,7 +395,10 @@ impl MigrateTableMetadata {
             .map(|(datanode_id, regions)| {
                 let k = DatanodeTableKey::new(datanode_id, table_id);
                 info!("Creating DatanodeTableKey '{k}' => {regions:?}");
-                (k, DatanodeTableValue::new(table_id, regions))
+                (
+                    k,
+                    DatanodeTableValue::new(table_id, regions, engine.to_string()),
+                )
             })
             .collect::<Vec<_>>();
 

src/common/meta/src/key.rs

@@ -241,6 +241,7 @@ impl TableMetadataManager {
             .collect::<Vec<_>>();
         table_info.meta.region_numbers = region_numbers;
         let table_id = table_info.ident.table_id;
+        let engine = table_info.meta.engine.clone();
 
         // Creates table name.
         let table_name = TableNameKey::new(
@@ -260,9 +261,9 @@ impl TableMetadataManager {
 
         // Creates datanode table key value pairs.
         let distribution = region_distribution(&region_routes)?;
-        let create_datanode_table_txn = self
-            .datanode_table_manager()
-            .build_create_txn(table_id, distribution)?;
+        let create_datanode_table_txn =
+            self.datanode_table_manager()
+                .build_create_txn(table_id, &engine, distribution)?;
 
         // Creates table route.
         let table_route_value = TableRouteValue::new(region_routes);
@@ -439,6 +440,7 @@ impl TableMetadataManager {
     pub async fn update_table_route(
         &self,
         table_id: TableId,
+        engine: &str,
         current_table_route_value: TableRouteValue,
         new_region_routes: Vec<RegionRoute>,
     ) -> Result<()> {
@@ -449,6 +451,7 @@ impl TableMetadataManager {
 
         let update_datanode_table_txn = self.datanode_table_manager().build_update_txn(
             table_id,
+            engine,
             current_region_distribution,
             new_region_distribution,
         )?;
@@ -863,6 +866,7 @@ mod tests {
         let table_info: RawTableInfo =
             new_test_table_info(region_routes.iter().map(|r| r.region.id.region_number())).into();
         let table_id = table_info.ident.table_id;
+        let engine = table_info.meta.engine.as_str();
         let current_table_route_value = TableRouteValue::new(region_routes.clone());
         // creates metadata.
         table_metadata_manager
@@ -879,6 +883,7 @@ mod tests {
         table_metadata_manager
             .update_table_route(
                 table_id,
+                engine,
                 current_table_route_value.clone(),
                 new_region_routes.clone(),
             )
@@ -890,6 +895,7 @@ mod tests {
         table_metadata_manager
             .update_table_route(
                 table_id,
+                engine,
                 current_table_route_value.clone(),
                 new_region_routes.clone(),
             )
@@ -902,6 +908,7 @@ mod tests {
         table_metadata_manager
             .update_table_route(
                 table_id,
+                engine,
                 current_table_route_value.clone(),
                 new_region_routes.clone(),
             )
@@ -918,7 +925,7 @@ mod tests {
             new_region_route(4, 4),
         ]);
         assert!(table_metadata_manager
-            .update_table_route(table_id, wrong_table_route_value, new_region_routes)
+            .update_table_route(table_id, engine, wrong_table_route_value, new_region_routes)
             .await
             .is_err());
     }

src/common/meta/src/key/datanode_table.rs

@@ -85,14 +85,16 @@ impl TableMetaKey for DatanodeTableKey {
 pub struct DatanodeTableValue {
     pub table_id: TableId,
     pub regions: Vec<RegionNumber>,
+    pub engine: String,
     version: u64,
 }
 
 impl DatanodeTableValue {
-    pub fn new(table_id: TableId, regions: Vec<RegionNumber>) -> Self {
+    pub fn new(table_id: TableId, regions: Vec<RegionNumber>, engine: String) -> Self {
         Self {
             table_id,
             regions,
+            engine,
             version: 0,
         }
     }
@@ -143,13 +145,14 @@ impl DatanodeTableManager {
     pub fn build_create_txn(
         &self,
         table_id: TableId,
+        engine: &str,
         distribution: RegionDistribution,
     ) -> Result<Txn> {
         let txns = distribution
             .into_iter()
             .map(|(datanode_id, regions)| {
                 let key = DatanodeTableKey::new(datanode_id, table_id);
-                let val = DatanodeTableValue::new(table_id, regions);
+                let val = DatanodeTableValue::new(table_id, regions, engine.to_string());
 
                 Ok(TxnOp::Put(key.as_raw_key(), val.try_as_raw_value()?))
             })
@@ -164,6 +167,7 @@ impl DatanodeTableManager {
     pub(crate) fn build_update_txn(
         &self,
         table_id: TableId,
+        engine: &str,
         current_region_distribution: RegionDistribution,
         new_region_distribution: RegionDistribution,
     ) -> Result<Txn> {
@@ -184,14 +188,16 @@ impl DatanodeTableManager {
                 if *current_region != regions {
                     let key = DatanodeTableKey::new(datanode, table_id);
                     let raw_key = key.as_raw_key();
-                    let val = DatanodeTableValue::new(table_id, regions).try_as_raw_value()?;
+                    let val = DatanodeTableValue::new(table_id, regions, engine.to_string())
+                        .try_as_raw_value()?;
                     opts.push(TxnOp::Put(raw_key, val));
                 }
             } else {
                 // New datanodes
                 let key = DatanodeTableKey::new(datanode, table_id);
                 let raw_key = key.as_raw_key();
-                let val = DatanodeTableValue::new(table_id, regions).try_as_raw_value()?;
+                let val = DatanodeTableValue::new(table_id, regions, engine.to_string())
+                    .try_as_raw_value()?;
                 opts.push(TxnOp::Put(raw_key, val));
             }
         }
@@ -224,7 +230,6 @@ impl DatanodeTableManager {
 
 #[cfg(test)]
 mod tests {
-
     use super::*;
 
     #[test]
@@ -239,9 +244,10 @@ mod tests {
         let value = DatanodeTableValue {
             table_id: 42,
             regions: vec![1, 2, 3],
+            engine: Default::default(),
             version: 1,
         };
-        let literal = br#"{"table_id":42,"regions":[1,2,3],"version":1}"#;
+        let literal = br#"{"table_id":42,"regions":[1,2,3],"engine":"","version":1}"#;
 
         let raw_value = value.try_as_raw_value().unwrap();
         assert_eq!(raw_value, literal);

src/meta-srv/src/procedure/region_failover/update_metadata.rs

@@ -58,6 +58,7 @@ impl UpdateRegionMetadata {
         failed_region: &RegionIdent,
     ) -> Result<()> {
         let table_id = failed_region.table_ident.table_id;
+        let engine = failed_region.table_ident.engine.as_str();
 
         let table_route_value = ctx
             .table_metadata_manager
@@ -85,7 +86,7 @@ impl UpdateRegionMetadata {
         );
 
         ctx.table_metadata_manager
-            .update_table_route(table_id, table_route_value, new_region_routes)
+            .update_table_route(table_id, engine, table_route_value, new_region_routes)
             .await
             .context(error::UpdateTableRouteSnafu)?;