add no_upload flag

risingwavelabs · Jun 7, 2023 · 21466a0 · 21466a0
1 parent e9d8fcb
commit 21466a0
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 8 deletions.
diff --git a/src/storage/src/hummock/store/state_store.rs b/src/storage/src/hummock/store/state_store.rs
@@ -66,6 +66,8 @@ pub struct LocalHummockStorage {
     /// Read handle.
     read_version: Arc<RwLock<HummockReadVersion>>,
 
+    no_upload: bool,
+
     /// Event sender.
     event_sender: mpsc::UnboundedSender<HummockEvent>,
 
@@ -404,9 +406,11 @@ impl LocalHummockStorage {
         self.update(VersionUpdate::Staging(StagingData::ImmMem(imm.clone())));
 
         // insert imm to uploader
-        self.event_sender
-            .send(HummockEvent::ImmToUploader(imm))
-            .unwrap();
+        if self.no_upload {
+            self.event_sender
+                .send(HummockEvent::ImmToUploader(imm))
+                .unwrap();
+        }
 
         timer.observe_duration();
 
@@ -437,6 +441,7 @@ impl LocalHummockStorage {
             table_id: option.table_id,
             is_consistent_op: option.is_consistent_op,
             table_option: option.table_option,
+            no_upload: option.no_upload,
             instance_guard,
             read_version,
             event_sender,

diff --git a/src/storage/src/store.rs b/src/storage/src/store.rs
@@ -375,16 +375,42 @@ pub struct NewLocalOptions {
     /// `update` and `delete` should match the original stored value.
     pub is_consistent_op: bool,
     pub table_option: TableOption,
+
+    /// Used to indicate that we should not upload the data.
+    pub no_upload: bool,
 }
 
 impl NewLocalOptions {
+    pub fn new(table_id: TableId, is_consistent_op: bool, table_option: TableOption) -> Self {
+        NewLocalOptions {
+            table_id,
+            is_consistent_op,
+            table_option,
+            no_upload: true,
+        }
+    }
+
+    pub fn new_temporary(
+        table_id: TableId,
+        is_consistent_op: bool,
+        table_option: TableOption,
+    ) -> Self {
+        NewLocalOptions {
+            table_id,
+            is_consistent_op,
+            table_option,
+            no_upload: false,
+        }
+    }
+
     pub fn for_test(table_id: TableId) -> Self {
         Self {
             table_id,
             is_consistent_op: false,
             table_option: TableOption {
                 retention_seconds: None,
             },
+            no_upload: false,
         }
     }
 }
diff --git a/src/stream/src/common/table/state_table.rs b/src/stream/src/common/table/state_table.rs
@@ -189,11 +189,11 @@ where
 
         let table_option = TableOption::build_table_option(table_catalog.get_properties());
         let local_state_store = store
-            .new_local(NewLocalOptions {
+            .new_local(NewLocalOptions::new(
                 table_id,
                 is_consistent_op,
                 table_option,
-            })
+            ))
             .await;
 
         let pk_data_types = pk_indices
@@ -384,11 +384,11 @@ where
         is_consistent_op: bool,
     ) -> Self {
         let local_state_store = store
-            .new_local(NewLocalOptions {
+            .new_local(NewLocalOptions::new(
                 table_id,
                 is_consistent_op,
-                table_option: TableOption::default(),
-            })
+                TableOption::default(),
+            ))
             .await;
 
         let pk_data_types = pk_indices