Toggle the staged state if the dev config switches staged on the fly (#39997)

Example scenario

1. dev pushes with {staged:true}
2. Before backfilling completes, dev pushes without staged

Here we want to toggle the backfilling state without restarting the
backfill, thus ensuring that the CLI blocks until its completed.

GitOrigin-RevId: 78a421b63cbf82e4ac9bba78b5dd602ec71a0255

Author: nipunn1313

crates/common/src/bootstrap_model/index/database_index/index_state.rs

@@ -35,6 +35,18 @@ impl DatabaseIndexState {
             Self::Enabled => false,
         }
     }
+
+    pub fn set_staged(&mut self, staged_new: bool) {
+        match self {
+            Self::Backfilling(index_state) => {
+                index_state.staged = staged_new;
+            },
+            Self::Backfilled { staged } => {
+                *staged = staged_new;
+            },
+            Self::Enabled => {},
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize)]

crates/common/src/bootstrap_model/index/index_config.rs

@@ -57,18 +57,17 @@ pub enum IndexConfig {
 impl IndexConfig {
     pub fn is_staged(&self) -> bool {
         match self {
-            Self::Database {
-                developer_config: _,
-                on_disk_state,
-            } => on_disk_state.is_staged(),
-            Self::Text {
-                developer_config: _,
-                on_disk_state,
-            } => on_disk_state.is_staged(),
-            Self::Vector {
-                developer_config: _,
-                on_disk_state,
-            } => on_disk_state.is_staged(),
+            Self::Database { on_disk_state, .. } => on_disk_state.is_staged(),
+            Self::Text { on_disk_state, .. } => on_disk_state.is_staged(),
+            Self::Vector { on_disk_state, .. } => on_disk_state.is_staged(),
+        }
+    }
+
+    pub fn set_staged(&mut self, staged: bool) {
+        match self {
+            Self::Database { on_disk_state, .. } => on_disk_state.set_staged(staged),
+            Self::Text { on_disk_state, .. } => on_disk_state.set_staged(staged),
+            Self::Vector { on_disk_state, .. } => on_disk_state.set_staged(staged),
         }
     }
 

crates/common/src/bootstrap_model/index/text_index/index_state.rs

@@ -38,6 +38,18 @@ impl TextIndexState {
             Self::SnapshottedAt(_) => false,
         }
     }
+
+    pub fn set_staged(&mut self, staged_new: bool) {
+        match self {
+            Self::Backfilling(index_state) => {
+                index_state.staged = staged_new;
+            },
+            Self::Backfilled { staged, .. } => {
+                *staged = staged_new;
+            },
+            Self::SnapshottedAt(_) => {},
+        }
+    }
 }
 
 #[derive(Serialize, Deserialize)]

crates/common/src/bootstrap_model/index/vector_index/index_state.rs

@@ -37,6 +37,18 @@ impl VectorIndexState {
         }
     }
 
+    pub fn set_staged(&mut self, staged_new: bool) {
+        match self {
+            VectorIndexState::Backfilling(backfill_state) => {
+                backfill_state.staged = staged_new;
+            },
+            VectorIndexState::Backfilled { staged, .. } => {
+                *staged = staged_new;
+            },
+            VectorIndexState::SnapshottedAt(_) => {},
+        }
+    }
+
     pub fn segments(&self) -> anyhow::Result<&Vec<FragmentedVectorSegment>> {
         match self {
             VectorIndexState::Backfilling(backfill_state) => Ok(&backfill_state.segments),

crates/database/src/bootstrap_model/index.rs

@@ -602,7 +602,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
 
     fn identical_or_replaced(
         &mut self,
-        existing_index: ParsedDocument<DeveloperIndexMetadata>,
+        mut existing_index: ParsedDocument<DeveloperIndexMetadata>,
         new_index: DeveloperIndexMetadata,
     ) -> IndexComparison {
         let existing_fields = DeveloperIndexConfig::from(existing_index.config.clone());
@@ -612,12 +612,13 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
                 IndexComparison::Identical(existing_index)
             } else {
                 // Staged status changed - use the previous on-disk state
-                // staged -> enabled
+                // Toggle the staged status and return enabled/disabled to
+                // inform callsite.
                 if existing_index.config.is_staged() {
-                    // Should we modify the index's staged flag here or when we actually go to
-                    // change its state?
+                    existing_index.config.set_staged(false);
                     IndexComparison::Enabled(existing_index)
                 } else {
+                    existing_index.config.set_staged(true);
                     IndexComparison::Disabled(existing_index)
                 }
             }
@@ -708,6 +709,18 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
         for index in &diff.added {
             self.add_application_index(namespace, index.clone()).await?;
         }
+
+        // Replace all the enabled/disabled ones in-place to update their `staged`
+        // status. The actual enabling/disabling will be done at finish_push
+        // (apply_config).
+        for index in diff.enabled.iter().chain(diff.disabled.iter()) {
+            let (id, value) = index.clone().into_id_and_value();
+            let new_config = value.config.try_into()?;
+            SystemMetadataModel::new_global(self.tx)
+                .patch(id, patch_value!("config" =>  Some(new_config))?)
+                .await?;
+        }
+
         Ok(diff)
     }