From 3a95caff32a4380f60c35e188bbbd2b6e172d38b Mon Sep 17 00:00:00 2001
From: xx01cyx <caoyuanxin0531@outlook.com>
Date: Fri, 17 Jun 2022 16:20:09 +0800
Subject: [PATCH] set vnode mapping for materialzied source

---
 src/meta/src/rpc/service/ddl_service.rs | 47 ++++++++++++++-----------
 1 file changed, 26 insertions(+), 21 deletions(-)

diff --git a/src/meta/src/rpc/service/ddl_service.rs b/src/meta/src/rpc/service/ddl_service.rs
index 003a4b38d9b2b..d7c4c5b51b0cb 100644
--- a/src/meta/src/rpc/service/ddl_service.rs
+++ b/src/meta/src/rpc/service/ddl_service.rs
@@ -296,27 +296,7 @@ where
                 .map_err(tonic_err)?;
             return Err(e.into());
         } else {
-            // Fill in vnode mapping in order to inform frontend of data distribution.
-            let vnode_mapping = self
-                .env
-                .hash_mapping_manager_ref()
-                .get_table_hash_mapping(&id);
-            match vnode_mapping {
-                Some(vnode_mapping) => {
-                    let (original_indices, data) = compress_data(&vnode_mapping);
-                    mview.mapping = Some(ParallelUnitMapping {
-                        table_id: id,
-                        original_indices,
-                        data,
-                    });
-                }
-                None => {
-                    return Err(ErrorCode::InternalError(
-                        "no data distribution found for materialized view".to_string(),
-                    ))
-                    .map_err(tonic_err);
-                }
-            }
+            self.set_mview_mapping(&mut mview).map_err(tonic_err)?;
         }
 
         // 4. Finally, update the catalog.
@@ -547,6 +527,8 @@ where
             // drop previously created source
             self.source_manager.drop_source(source_id).await?;
             return Err(e);
+        } else {
+            self.set_mview_mapping(&mut mview).map_err(tonic_err)?;
         }
 
         // Finally, update the catalog.
@@ -580,4 +562,27 @@ where
 
         Ok(version)
     }
+
+    /// Fill in mview's vnode mapping so that frontend will know the data distribution.
+    fn set_mview_mapping(&self, mview: &mut Table) -> RwResult<()> {
+        let vnode_mapping = self
+            .env
+            .hash_mapping_manager_ref()
+            .get_table_hash_mapping(&mview.id);
+        match vnode_mapping {
+            Some(vnode_mapping) => {
+                let (original_indices, data) = compress_data(&vnode_mapping);
+                mview.mapping = Some(ParallelUnitMapping {
+                    table_id: mview.id,
+                    original_indices,
+                    data,
+                });
+                Ok(())
+            }
+            None => Err(ErrorCode::InternalError(
+                "no data distribution found for materialized view".to_string(),
+            )
+            .into()),
+        }
+    }
 }