Wire in schema and index changes to {start,finish}_push (#26839)

GitOrigin-RevId: d479d6d401bdd9e078eee37ce25217224630be53

Author: Sujay Jayakar

crates/application/src/snapshot_import.rs

@@ -2182,7 +2182,11 @@ async fn prepare_table_for_import<RT: Runtime>(
                             )
                             .await?;
                         IndexModel::new(tx)
-                            .copy_indexes_to_table(table_name, table_id.tablet_id)
+                            .copy_indexes_to_table(
+                                TableNamespace::by_component_TODO(),
+                                table_name,
+                                table_id.tablet_id,
+                            )
                             .await?;
                         Ok(table_id)
                     }
@@ -3007,10 +3011,10 @@ a
             let mut tx = app.begin(identity.clone()).await?;
             let mut index_model = IndexModel::new(&mut tx);
             let index_id = index_model
-                .add_application_index(IndexMetadata::new_enabled(
-                    index_name.clone(),
-                    vec!["a".parse()?].try_into()?,
-                ))
+                .add_application_index(
+                    TableNamespace::test_user(),
+                    IndexMetadata::new_enabled(index_name.clone(), vec!["a".parse()?].try_into()?),
+                )
                 .await?;
             app.commit_test(tx).await?;
             index_id

crates/database/src/bootstrap_model/index.rs

@@ -61,7 +61,6 @@ use indexing::{
         Index,
     },
 };
-use itertools::Itertools;
 use value::{
     ResolvedDocumentId,
     TableMapping,
@@ -113,14 +112,15 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
     /// transaction has committed.
     pub async fn add_application_index(
         &mut self,
+        namespace: TableNamespace,
         index: IndexMetadata<TableName>,
     ) -> anyhow::Result<ResolvedDocumentId> {
         anyhow::ensure!(
             self.tx.identity().is_admin() || self.tx.identity().is_system(),
             unauthorized_error("add_index")
         );
         anyhow::ensure!(!index.name.is_system_owned(), "Can't change system indexes");
-        let application_indexes = self.get_application_indexes().await?;
+        let application_indexes = self.get_application_indexes(namespace).await?;
         // Indexes may exist in both a pending and an enabled state. If we're at or over
         // the index limit, we should still be able to add a new pending copy of
         // an enabled index with the expectation that the pending index will
@@ -137,8 +137,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
                 || index_names_in_table.len() < MAX_INDEXES_PER_TABLE,
             index_validation_error::too_many_indexes(index.name.table(), MAX_INDEXES_PER_TABLE)
         );
-        self._add_index(TableNamespace::by_component_TODO(), index)
-            .await
+        self._add_index(namespace, index).await
     }
 
     /// Add system index.
@@ -281,7 +280,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
             .await?;
 
         tracing::info!(
-            "Committed indexes: (added {}. dropped {})",
+            "Committed indexes for {namespace:?}: (added {}. dropped {})",
             index_diff.added.len(),
             index_diff.dropped.len(),
         );
@@ -312,7 +311,8 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
             added: vec![],
             dropped: dropped.clone(),
         };
-        self.apply_index_diff(&only_dropped_tables).await?;
+        self.apply_index_diff(namespace, &only_dropped_tables)
+            .await?;
 
         // Added indexes should have backfilled via build_indexes
         // (for < 0.14.0 CLIs) or in apply_config (for >= 0.14.0 CLIs).
@@ -339,15 +339,19 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
         Ok(())
     }
 
-    pub async fn apply_index_diff(&mut self, diff: &LegacyIndexDiff) -> anyhow::Result<()> {
+    pub async fn apply_index_diff(
+        &mut self,
+        namespace: TableNamespace,
+        diff: &LegacyIndexDiff,
+    ) -> anyhow::Result<()> {
         if !(self.tx.identity().is_admin() || self.tx.identity().is_system()) {
             anyhow::bail!(unauthorized_error("modify_indexes"));
         }
         for index in &diff.dropped {
             self.drop_index(index.id()).await?;
         }
         for index in &diff.added {
-            self.add_application_index(index.clone()).await?;
+            self.add_application_index(namespace, index.clone()).await?;
         }
 
         Ok(())
@@ -398,7 +402,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
 
         let mut remaining_indexes: HashMap<IndexName, Vec<ParsedDocument<DeveloperIndexMetadata>>> =
             HashMap::new();
-        for index in self.get_application_indexes().await? {
+        for index in self.get_application_indexes(namespace).await? {
             remaining_indexes
                 .entry(index.name.clone())
                 .or_default()
@@ -513,7 +517,8 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
         // Dropped will be removed in apply_config when the rest of the schema is
         // committed.
         if !only_new_and_mutated.is_empty() {
-            self.apply_index_diff(&only_new_and_mutated).await?;
+            self.apply_index_diff(namespace, &only_new_and_mutated)
+                .await?;
         }
         Ok(diff)
     }
@@ -527,7 +532,7 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
         if diff.is_empty() {
             return Ok(diff);
         }
-        self.apply_index_diff(&diff).await?;
+        self.apply_index_diff(namespace, &diff).await?;
         Ok(diff)
     }
 
@@ -775,8 +780,9 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
     /// name (but different configurations).
     pub async fn get_system_indexes(
         &mut self,
+        namespace: TableNamespace,
     ) -> anyhow::Result<Vec<ParsedDocument<DeveloperIndexMetadata>>> {
-        self.get_indexes(IndexCategory::System).await
+        self.get_indexes(IndexCategory::System, namespace).await
     }
 
     /// Returns all registered indexes that aren't system owned including both
@@ -786,23 +792,32 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
     /// name (but different configurations).
     pub async fn get_application_indexes(
         &mut self,
+        namespace: TableNamespace,
     ) -> anyhow::Result<Vec<ParsedDocument<DeveloperIndexMetadata>>> {
-        self.get_indexes(IndexCategory::Application).await
+        self.get_indexes(IndexCategory::Application, namespace)
+            .await
     }
 
     async fn get_indexes(
         &mut self,
         category: IndexCategory,
+        namespace: TableNamespace,
     ) -> anyhow::Result<Vec<ParsedDocument<DeveloperIndexMetadata>>> {
-        let indexes_with_id = self.get_all_indexes().await?;
+        let indexes = self.get_all_indexes().await?;
         let table_mapping = self.tx.table_mapping();
-        indexes_with_id
-            .into_iter()
-            .filter(|doc| category.belongs(doc, table_mapping))
-            .map(|doc: ParsedDocument<TabletIndexMetadata>| {
-                doc.map(|metadata| metadata.map_table(&table_mapping.tablet_to_name()))
-            })
-            .try_collect()
+        let mut result = vec![];
+        for doc in indexes {
+            if !category.belongs(&doc, table_mapping) {
+                continue;
+            }
+            let tablet_id = *doc.name.table();
+            if table_mapping.tablet_namespace(tablet_id)? != namespace {
+                continue;
+            }
+            let doc = doc.map(|metadata| metadata.map_table(&table_mapping.tablet_to_name()))?;
+            result.push(doc);
+        }
+        Ok(result)
     }
 
     pub async fn drop_index(&mut self, index_id: ResolvedDocumentId) -> anyhow::Result<()> {
@@ -812,10 +827,14 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
         Ok(())
     }
 
-    pub async fn drop_system_index(&mut self, index_name: IndexName) -> anyhow::Result<()> {
+    pub async fn drop_system_index(
+        &mut self,
+        namespace: TableNamespace,
+        index_name: IndexName,
+    ) -> anyhow::Result<()> {
         anyhow::ensure!(index_name.table().is_system());
         let system_index = self
-            .get_system_indexes()
+            .get_system_indexes(namespace)
             .await?
             .into_iter()
             .find(|index| index.name == index_name);
@@ -828,14 +847,15 @@ impl<'a, RT: Runtime> IndexModel<'a, RT> {
 
     pub async fn copy_indexes_to_table(
         &mut self,
+        namespace: TableNamespace,
         source_table: &TableName,
         target_table: TabletId,
     ) -> anyhow::Result<()> {
         // Copy over enabled indexes from existing active table, if any.
         let Some(active_table_id) = self
             .tx
             .table_mapping()
-            .namespace(TableNamespace::by_component_TODO())
+            .namespace(namespace)
             .id_if_exists(source_table)
         else {
             return Ok(());

crates/database/src/bootstrap_model/system_metadata.rs

@@ -5,9 +5,12 @@ use common::{
 };
 use value::{
     ConvexObject,
+    InternalId,
     ResolvedDocumentId,
+    TableIdentifier,
     TableName,
     TableNamespace,
+    TabletIdAndTableNumber,
 };
 
 use crate::{
@@ -57,8 +60,37 @@ impl<'a, RT: Runtime> SystemMetadataModel<'a, RT> {
         if !(self.tx.identity.is_system() || self.tx.identity.is_admin()) {
             anyhow::bail!(unauthorized_error("insert_metadata"));
         }
-        let table_id = self
-            .tx
+        let table_id = self.lookup_table_id(table)?;
+        let id = self.tx.id_generator.generate(&table_id);
+        let creation_time = self.tx.next_creation_time.increment()?;
+        let document = ResolvedDocument::new(id, creation_time, value)?;
+        self.tx.insert_document(document).await
+    }
+
+    pub fn allocate_internal_id(&mut self) -> anyhow::Result<InternalId> {
+        Ok(self.tx.id_generator.generate_internal())
+    }
+
+    /// Create a new document with a predetermined internal ID.
+    pub async fn insert_with_internal_id(
+        &mut self,
+        table: &TableName,
+        internal_id: InternalId,
+        value: ConvexObject,
+    ) -> anyhow::Result<ResolvedDocumentId> {
+        anyhow::ensure!(table.is_system());
+        if !(self.tx.identity.is_system() || self.tx.identity.is_admin()) {
+            anyhow::bail!(unauthorized_error("insert_metadata"));
+        }
+        let table_id = self.lookup_table_id(table)?;
+        let document_id = table_id.id(internal_id);
+        let creation_time = self.tx.next_creation_time.increment()?;
+        let document = ResolvedDocument::new(document_id, creation_time, value)?;
+        self.tx.insert_document(document).await
+    }
+
+    fn lookup_table_id(&mut self, table: &TableName) -> anyhow::Result<TabletIdAndTableNumber> {
+        self.tx
             .table_mapping()
             .namespace(self.namespace)
             .id(table)
@@ -71,11 +103,7 @@ impl<'a, RT: Runtime> SystemMetadataModel<'a, RT> {
                 } else {
                     format!("Failed to find system table {table}")
                 }
-            })?;
-        let id = self.tx.id_generator.generate(&table_id);
-        let creation_time = self.tx.next_creation_time.increment()?;
-        let document = ResolvedDocument::new(id, creation_time, value)?;
-        self.tx.insert_document(document).await
+            })
     }
 
     /// Creates a new document with given value in the specified table without

crates/database/src/search_and_vector_bootstrap.rs

@@ -826,7 +826,7 @@ mod tests {
         let index_metadata = backfilling_vector_index()?;
         let mut tx = fixtures.db.begin_system().await?;
         IndexModel::new(&mut tx)
-            .add_application_index(index_metadata.clone())
+            .add_application_index(TableNamespace::test_user(), index_metadata.clone())
             .await?;
         fixtures.db.commit(tx).await?;
         reopen_db(&rt, &fixtures).await?;
@@ -842,7 +842,7 @@ mod tests {
         let index_metadata = backfilling_vector_index()?;
         let mut tx = db.begin_system().await?;
         IndexModel::new(&mut tx)
-            .add_application_index(index_metadata.clone())
+            .add_application_index(TableNamespace::test_user(), index_metadata.clone())
             .await?;
         db.commit(tx).await?;
         backfill_vector_indexes(rt.clone(), db.clone(), reader, storage.clone()).await?;
@@ -1188,7 +1188,7 @@ mod tests {
             btreeset! {"filterField".parse()?},
         );
         IndexModel::new(&mut tx)
-            .add_application_index(index)
+            .add_application_index(TableNamespace::test_user(), index)
             .await?;
         db.commit(tx).await?;
 

crates/database/src/table_iteration.rs

@@ -788,10 +788,10 @@ mod tests {
         let index_fields = IndexedFields::try_from(vec![field.clone()])?;
         let mut tx = database.begin(Identity::system()).await?;
         IndexModel::new(&mut tx)
-            .add_application_index(IndexMetadata::new_enabled(
-                index_name.clone(),
-                index_fields.clone(),
-            ))
+            .add_application_index(
+                TableNamespace::test_user(),
+                IndexMetadata::new_enabled(index_name.clone(), index_fields.clone()),
+            )
             .await?;
         database.commit(tx).await?;
         IndexWorker::new_terminating(