store: Add restore orchestration in SubgraphStore and DeploymentStore

Wire up the restore pipeline through SubgraphStore::restore() and
DeploymentStore::restore(). Uses plan_restore() to determine whether to
create or replace the deployment site, validates the target shard exists,
resolves the subgraph name for deployment rule matching, and assigns the
restored deployment to a node.

Changes:
- DeploymentStore::restore() coordinates schema creation, data import,
  and finalization
- Inner::restore() handles conflict resolution, site allocation, and
  node assignment via deployment rules
- Expose create_site() and find_active_site() on primary::Connection
- Make create_site() accept an `active` parameter

Author: lutter

store/postgres/src/deployment_store.rs

@@ -30,7 +30,7 @@ use std::collections::{BTreeMap, HashMap};
 use std::convert::Into;
 use std::ops::Bound;
 use std::ops::{Deref, Range};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::str::FromStr;
 use std::sync::{atomic::AtomicUsize, Arc, Mutex};
 use std::time::{Duration, Instant};
@@ -909,6 +909,21 @@ impl DeploymentStore {
             .dump(&mut conn, index_list, dir, &site.network, entity_count)
             .await
     }
+
+    pub(crate) async fn restore(
+        &self,
+        site: Arc<Site>,
+        dir: &Path,
+        metadata: &crate::relational::dump::Metadata,
+    ) -> Result<(), StoreError> {
+        let mut conn = self.pool.get_permitted().await?;
+        let layout =
+            crate::relational::restore::create_schema(&mut conn, site, metadata, dir).await?;
+        crate::relational::restore::import_data(&mut conn, &layout, metadata, dir, &self.logger)
+            .await?;
+        crate::relational::restore::finalize(&mut conn, &layout, metadata, &self.logger).await?;
+        Ok(())
+    }
 }
 
 /// Methods that back the trait `WritableStore`, but have small variations in their signatures

store/postgres/src/parquet/convert.rs

@@ -238,7 +238,6 @@ impl ColumnBuilder {
 }
 
 /// A row extracted from a Parquet file, ready for insertion into an entity table.
-#[allow(dead_code)]
 pub struct RestoreRow {
     pub vid: i64,
     pub block: BlockNumber,
@@ -253,7 +252,6 @@ pub struct RestoreRow {
 ///
 /// The batch must follow the column layout produced by `schema::arrow_schema`:
 /// `vid`, block tracking columns, optional `causality_region`, then data columns.
-#[allow(dead_code)]
 pub fn record_batch_to_restore_rows(
     batch: &RecordBatch,
     table: &Table,
@@ -337,7 +335,6 @@ pub fn record_batch_to_restore_rows(
 }
 
 /// A row extracted from a `data_sources$` Parquet file.
-#[allow(dead_code)]
 pub struct DataSourceRestoreRow {
     pub vid: i64,
     pub block_range_start: BlockNumber,
@@ -356,7 +353,6 @@ pub struct DataSourceRestoreRow {
 ///
 /// The batch must follow the fixed column layout from
 /// `schema::data_sources_arrow_schema`.
-#[allow(dead_code)]
 pub fn record_batch_to_data_source_rows(
     batch: &RecordBatch,
 ) -> Result<Vec<DataSourceRestoreRow>, StoreError> {

store/postgres/src/parquet/reader.rs

@@ -10,7 +10,6 @@ use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
 /// Opens the file, reads all row groups, and returns them as a vector
 /// of `RecordBatch`es. The batches retain the schema embedded in the
 /// Parquet file.
-#[allow(dead_code)]
 pub fn read_batches(path: &Path) -> Result<Vec<RecordBatch>, StoreError> {
     let file = fs::File::open(path).map_err(|e| {
         StoreError::InternalError(format!(

store/postgres/src/primary.rs

@@ -1391,7 +1391,7 @@ impl Connection {
     /// function only performs the basic operations for creation, and the
     /// caller must check that other conditions (like whether there already
     /// is an active site for the deployment) are met
-    async fn create_site(
+    pub(crate) async fn create_site(
         &mut self,
         shard: Shard,
         deployment: DeploymentHash,
@@ -1467,6 +1467,13 @@ impl Connection {
             .map(|site| (site, site_was_created))
     }
 
+    pub async fn find_active_site(
+        &mut self,
+        subgraph: &DeploymentHash,
+    ) -> Result<Option<Site>, StoreError> {
+        queries::find_active_site(&mut self.conn, subgraph).await
+    }
+
     pub async fn assigned_node(&mut self, site: &Site) -> Result<Option<NodeId>, StoreError> {
         queries::assigned_node(&mut self.conn, site).await
     }

store/postgres/src/relational/dump.rs

@@ -167,7 +167,6 @@ pub(crate) struct Metadata {
 
 impl Metadata {
     /// Read and validate a dump's `metadata.json`.
-    #[allow(dead_code)]
     pub fn from_file(path: &Path) -> Result<Self, StoreError> {
         let content = fs::read_to_string(path).map_err(|e| {
             StoreError::InternalError(format!("failed to read {}: {e}", path.display()))

store/postgres/src/relational/restore.rs

@@ -220,7 +220,6 @@ async fn import_data_sources(
 /// - Obtaining the right shard connection
 ///
 /// Entity data import and finalization are handled separately.
-#[allow(dead_code)]
 pub async fn create_schema(
     conn: &mut AsyncPgConnection,
     site: Arc<Site>,
@@ -319,7 +318,6 @@ pub async fn create_schema(
 /// This is resumable: if interrupted, it can be called again and will
 /// skip already-imported rows by checking the current max(vid) in each
 /// table.
-#[allow(dead_code)]
 pub async fn import_data(
     conn: &mut AsyncPgConnection,
     layout: &Layout,
@@ -365,7 +363,6 @@ pub async fn import_data(
 /// This must be called after `import_data` has completed successfully.
 /// Setting the head block is the very last operation — it marks the
 /// deployment as "ready".
-#[allow(dead_code)]
 pub async fn finalize(
     conn: &mut AsyncPgConnection,
     layout: &Layout,

store/postgres/src/relational_queries.rs

@@ -2355,7 +2355,6 @@ impl<'a> InsertRow<'a> {
     /// Unlike `new()`, this looks up values by SQL column name rather than
     /// entity field name, since `RestoreRow.values` is keyed by SQL name.
     /// Fulltext columns are regenerated from their source fields.
-    #[allow(dead_code)]
     fn from_restore(
         columns: &[&'a Column],
         row: &'a RestoreRow,
@@ -2460,7 +2459,6 @@ impl<'a> InsertQuery<'a> {
     /// All data columns from the dump are present, so `unique_columns`
     /// includes every column in the table. Fulltext columns are
     /// regenerated from their source fields.
-    #[allow(dead_code)]
     pub fn for_restore(
         table: &'a Table,
         rows: &'a [RestoreRow],

store/postgres/src/subgraph_store.rs

@@ -18,8 +18,8 @@ use graph::{
     components::{
         server::index_node::VersionInfo,
         store::{
-            self, BlockPtrForNumber, BlockStore, DeploymentLocator, EnsLookup as EnsLookupTrait,
-            PruneReporter, PruneRequest, SubgraphFork,
+            self, BlockPtrForNumber, BlockStore, DeploymentLocator, DeploymentSchemaVersion,
+            EnsLookup as EnsLookupTrait, PruneReporter, PruneRequest, SubgraphFork,
         },
     },
     data::{
@@ -42,7 +42,9 @@ use graph::{derive::CheapClone, futures03::future::join_all, prelude::alloy::pri
 use crate::{
     catalog::Catalog,
     deployment::{OnSync, SubgraphHealth},
-    primary::{self, DeploymentId, Mirror as PrimaryMirror, Primary, Site},
+    primary::{
+        self, DeploymentId, Mirror as PrimaryMirror, Primary, RestoreAction, RestoreMode, Site,
+    },
     relational::{
         self,
         index::{IndexList, Method},
@@ -735,6 +737,35 @@ impl Inner {
         }
     }
 
+    /// Determine the target node for a deployment using the configured
+    /// deployment rules, ignoring the shard selection. Returns an error
+    /// if no rule matches.
+    async fn node_for_deployment(
+        &self,
+        name: &SubgraphName,
+        network: &str,
+    ) -> Result<NodeId, StoreError> {
+        let placement = self
+            .placer
+            .place(name.as_str(), network)
+            .map_err(|msg| internal_error!("illegal indexer name in deployment rule: {}", msg))?;
+
+        match placement {
+            Some((_, nodes)) if !nodes.is_empty() => {
+                if nodes.len() == 1 {
+                    Ok(nodes.into_iter().next().unwrap())
+                } else {
+                    let mut pconn = self.primary_conn().await?;
+                    // unwrap: nodes is not empty
+                    Ok(pconn.least_assigned_node(&nodes).await?.unwrap())
+                }
+            }
+            _ => Err(StoreError::InternalError(
+                "no deployment rule matches this deployment".into(),
+            )),
+        }
+    }
+
     pub async fn copy_deployment(
         &self,
         src: &DeploymentLocator,
@@ -1446,6 +1477,114 @@ impl Inner {
 
         store.dump(site, directory).await
     }
+
+    pub async fn restore(
+        &self,
+        dir: &std::path::Path,
+        shard: Shard,
+        name: Option<SubgraphName>,
+        mode: RestoreMode,
+    ) -> Result<(), StoreError> {
+        use crate::relational::dump::Metadata;
+
+        let metadata_path = dir.join("metadata.json");
+        let metadata = Metadata::from_file(&metadata_path)?;
+
+        // Validate that the target shard exists before making any DB changes
+        self.stores
+            .get(&shard)
+            .ok_or_else(|| StoreError::UnknownShard(shard.to_string()))?;
+
+        // Resolve the subgraph name for deployment rule matching. If not
+        // supplied, look up an existing name from the DB; error if none.
+        let name = match name {
+            Some(n) => n,
+            None => {
+                let names = self
+                    .mirror
+                    .subgraphs_by_deployment_hash(metadata.deployment.as_str())
+                    .await?;
+                let (name, _) = names.into_iter().next().ok_or_else(|| {
+                    StoreError::InternalError(
+                        "no subgraph name found for this deployment; use --name to specify one"
+                            .into(),
+                    )
+                })?;
+                SubgraphName::new(name).map_err(|n| {
+                    StoreError::InternalError(format!("invalid subgraph name `{n}` in database"))
+                })?
+            }
+        };
+
+        // Use deployment rules to determine which node should index this
+        // deployment. The rules also return candidate shards, but we ignore
+        // those since the shard is user-specified for restore.
+        let node = self.node_for_deployment(&name, &metadata.network).await?;
+
+        let mut pconn = self.primary_conn().await?;
+        let action = pconn
+            .plan_restore(&shard, &metadata.deployment, &mode)
+            .await?;
+
+        // Determine schema_version the same way allocate_site does
+        let schema_version = match metadata.graft_base.as_ref() {
+            Some(graft_base) => {
+                let base_site = pconn.find_active_site(graft_base).await?.ok_or_else(|| {
+                    StoreError::DeploymentNotFound("graft_base not found".to_string())
+                })?;
+                base_site.schema_version
+            }
+            None => DeploymentSchemaVersion::LATEST,
+        };
+
+        let site = match action {
+            RestoreAction::Create { active } => {
+                pconn
+                    .create_site(
+                        shard,
+                        metadata.deployment.clone(),
+                        metadata.network.clone(),
+                        schema_version,
+                        active,
+                    )
+                    .await?
+            }
+            RestoreAction::Replace { existing } => {
+                let was_active = existing.active;
+                let existing = Arc::new(existing);
+                let store = self.for_site(&existing)?;
+                store.drop_deployment(&existing).await?;
+                pconn.drop_site(&existing).await?;
+                // Drop and re-acquire the primary connection to avoid pool
+                // deadlock: drop_deployment above used a separate connection
+                // from the same pool, and create_site below needs one too.
+                drop(pconn);
+                let mut pconn = self.primary_conn().await?;
+                pconn
+                    .create_site(
+                        shard,
+                        metadata.deployment.clone(),
+                        metadata.network.clone(),
+                        schema_version,
+                        was_active,
+                    )
+                    .await?
+            }
+        };
+
+        let site = Arc::new(site);
+        let store = self.for_site(&site)?;
+        store.restore(site.cheap_clone(), dir, &metadata).await?;
+
+        // Assign the restored deployment to the node determined by
+        // deployment rules
+        let mut pconn = self.primary_conn().await?;
+        let changes = pconn.assign_subgraph(&site, &node).await?;
+        let event = StoreEvent::new(changes);
+        pconn.send_store_event(&self.sender, &event).await?;
+
+        Ok(())
+    }
 }
 
 const STATE_ENS_NOT_CHECKED: u8 = 0;