graph, node, store: Change how restore deterimes the shard

If --shard is not given, consult deployment rules first, only fall back to
the primary if that doesn't yield a result

Author: lutter

graph/src/data/store/mod.rs

@@ -33,6 +33,12 @@ pub mod ethereum;
 /// Conversion of values to/from SQL
 pub mod sql;
 
+lazy_static! {
+    /// The name of the default node is `"default"`. This is used when no
+    /// node is specified for a deployment.
+    pub static ref DEFAULT_NODE_ID: NodeId = NodeId::new("default").unwrap();
+}
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct NodeId(String);
 

node/src/bin/manager.rs

@@ -352,7 +352,10 @@ pub enum Command {
     Restore {
         /// Path to the dump directory
         directory: String,
-        /// The database shard to restore into (default: primary)
+        /// The database shard to restore into. When not given, use the
+        /// deployment rules to determine the shard, or default to the
+        /// primary shard if no rules match. This option is required when
+        /// using `--add`
         #[clap(long)]
         shard: Option<String>,
         /// Subgraph name for deployment rule matching and node assignment.

node/src/manager/commands/restore.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 use graph::{bail, prelude::anyhow::Result};
 
 use graph::prelude::SubgraphName;
-use graph_store_postgres::{RestoreMode, Shard, SubgraphStore, PRIMARY_SHARD};
+use graph_store_postgres::{RestoreMode, Shard, SubgraphStore};
 
 pub async fn run(
     subgraph_store: Arc<SubgraphStore>,
@@ -28,11 +28,6 @@ pub async fn run(
         );
     }
 
-    let shard = match shard {
-        Some(s) => Shard::new(s)?,
-        None => PRIMARY_SHARD.clone(),
-    };
-
     let name = name
         .map(|n| SubgraphName::new(n.clone()).map_err(|_| anyhow::anyhow!("invalid name `{n}`")))
         .transpose()?;
@@ -47,6 +42,8 @@ pub async fn run(
         RestoreMode::Default
     };
 
+    let shard = shard.map(Shard::new).transpose()?;
+
     subgraph_store
         .restore(&directory, shard, name, mode)
         .await?;

store/postgres/src/subgraph_store.rs

@@ -24,6 +24,7 @@ use graph::{
     },
     data::{
         query::QueryTarget,
+        store::DEFAULT_NODE_ID,
         subgraph::{schema::DeploymentCreate, status, DeploymentFeatures},
     },
     internal_error,
@@ -737,35 +738,6 @@ 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,
@@ -1481,7 +1453,7 @@ impl Inner {
     pub async fn restore(
         &self,
         dir: &std::path::Path,
-        shard: Shard,
+        shard: Option<Shard>,
         name: Option<SubgraphName>,
         mode: RestoreMode,
     ) -> Result<(), StoreError> {
@@ -1490,11 +1462,6 @@ impl Inner {
         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 {
@@ -1517,9 +1484,16 @@ impl Inner {
         };
 
         // 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?;
+        // deployment and how to place it.
+        let (placed_shard, node) = self
+            .place(&name, &metadata.network, DEFAULT_NODE_ID.clone())
+            .await?;
+        let shard = shard.unwrap_or(placed_shard);
+
+        // Validate that the target shard exists before making any DB changes
+        self.stores
+            .get(&shard)
+            .ok_or_else(|| StoreError::UnknownShard(shard.to_string()))?;
 
         let mut pconn = self.primary_conn().await?;
         let action = pconn