Refactor graphman rewind to use pause and resume logic (#5213)

* Refactor deployment pause and resume logic

* Refactor pause_and_resume to use DeploymentLocator, Add cloning functionality to the Context struct,

* Update node/src/manager/commands/rewind.rs

Co-authored-by: Krishnanand V P <[email protected]>

* Update node/src/bin/manager.rs

Co-authored-by: Krishnanand V P <[email protected]>

* Update node/src/bin/manager.rs

Co-authored-by: Krishnanand V P <[email protected]>

* Refactor cloning and pausing logic in manager.rs and rewind.rs

* Refactor code to use mutable connection in assign.rs and remove unused constant in rewind.rs

* removing cloning in manager.rs

* minor refactoring

* node: throw an error when a deployment search matches more than one in rewind

---------

Co-authored-by: Krishnanand V P <[email protected]>
Co-authored-by: incrypto32 <[email protected]>

Author: itsjerryokolo

node/src/bin/manager.rs

@@ -1176,15 +1176,24 @@ async fn main() -> anyhow::Result<()> {
         }
         Pause { deployment } => {
             let sender = ctx.notification_sender();
-            commands::assign::pause_or_resume(ctx.primary_pool(), &sender, &deployment, true)
+            let pool = ctx.primary_pool();
+            let locator = &deployment.locate_unique(&pool)?;
+            commands::assign::pause_or_resume(pool, &sender, locator, true)
         }
+
         Resume { deployment } => {
             let sender = ctx.notification_sender();
-            commands::assign::pause_or_resume(ctx.primary_pool(), &sender, &deployment, false)
+            let pool = ctx.primary_pool();
+            let locator = &deployment.locate_unique(&pool).unwrap();
+
+            commands::assign::pause_or_resume(pool, &sender, locator, false)
         }
         Restart { deployment, sleep } => {
             let sender = ctx.notification_sender();
-            commands::assign::restart(ctx.primary_pool(), &sender, &deployment, sleep)
+            let pool = ctx.primary_pool();
+            let locator = &deployment.locate_unique(&pool).unwrap();
+
+            commands::assign::restart(pool, &sender, locator, sleep)
         }
         Rewind {
             force,
@@ -1194,13 +1203,16 @@ async fn main() -> anyhow::Result<()> {
             deployments,
             start_block,
         } => {
+            let notification_sender = ctx.notification_sender();
             let (store, primary) = ctx.store_and_primary();
+
             commands::rewind::run(
                 primary,
                 store,
                 deployments,
                 block_hash,
                 block_number,
+                &notification_sender,
                 force,
                 sleep,
                 start_block,

node/src/manager/commands/assign.rs

@@ -1,3 +1,4 @@
+use graph::components::store::DeploymentLocator;
 use graph::prelude::{anyhow::anyhow, Error, NodeId, StoreEvent};
 use graph_store_postgres::{
     command_support::catalog, connection_pool::ConnectionPool, NotificationSender,
@@ -75,11 +76,9 @@ pub fn reassign(
 pub fn pause_or_resume(
     primary: ConnectionPool,
     sender: &NotificationSender,
-    search: &DeploymentSearch,
+    locator: &DeploymentLocator,
     should_pause: bool,
 ) -> Result<(), Error> {
-    let locator = search.locate_unique(&primary)?;
-
     let pconn = primary.get()?;
     let mut conn = catalog::Connection::new(pconn);
 
@@ -115,15 +114,15 @@ pub fn pause_or_resume(
 pub fn restart(
     primary: ConnectionPool,
     sender: &NotificationSender,
-    search: &DeploymentSearch,
+    locator: &DeploymentLocator,
     sleep: Duration,
 ) -> Result<(), Error> {
-    pause_or_resume(primary.clone(), sender, search, true)?;
+    pause_or_resume(primary.clone(), sender, locator, true)?;
     println!(
         "Waiting {}s to make sure pausing was processed",
         sleep.as_secs()
     );
     thread::sleep(sleep);
-    pause_or_resume(primary, sender, search, false)?;
+    pause_or_resume(primary, sender, locator, false)?;
     Ok(())
 }

node/src/manager/commands/rewind.rs

@@ -5,10 +5,11 @@ use std::{collections::HashSet, convert::TryFrom};
 
 use graph::anyhow::bail;
 use graph::components::store::{BlockStore as _, ChainStore as _};
-use graph::prelude::{anyhow, BlockNumber, BlockPtr, NodeId, SubgraphStore};
-use graph_store_postgres::BlockStore;
+use graph::prelude::{anyhow, BlockNumber, BlockPtr};
 use graph_store_postgres::{connection_pool::ConnectionPool, Store};
+use graph_store_postgres::{BlockStore, NotificationSender};
 
+use crate::manager::commands::assign::pause_or_resume;
 use crate::manager::deployment::{Deployment, DeploymentSearch};
 
 async fn block_ptr(
@@ -61,12 +62,11 @@ pub async fn run(
     searches: Vec<DeploymentSearch>,
     block_hash: Option<String>,
     block_number: Option<BlockNumber>,
+    sender: &NotificationSender,
     force: bool,
     sleep: Duration,
     start_block: bool,
 ) -> Result<(), anyhow::Error> {
-    const PAUSED: &str = "paused_";
-
     // Sanity check
     if !start_block && (block_hash.is_none() || block_number.is_none()) {
         bail!("--block-hash and --block-number must be specified when --start-block is not set");
@@ -75,15 +75,20 @@ pub async fn run(
     let subgraph_store = store.subgraph_store();
     let block_store = store.block_store();
 
-    let deployments = searches
-        .iter()
-        .map(|search| search.lookup(&primary))
-        .collect::<Result<Vec<_>, _>>()?
-        .into_iter()
-        .flatten()
-        .collect::<Vec<_>>();
+    let mut deployments = Vec::new();
+    for search in &searches {
+        let results = search.lookup(&primary)?;
+        if results.len() > 1 {
+            bail!(
+                "Multiple deployments found for the search : {}. Try using the id of the deployment (eg: sgd143) to uniquely identify the deployment.",
+                search
+            );
+        }
+        deployments.extend(results);
+    }
+
     if deployments.is_empty() {
-        println!("nothing to do");
+        println!("No deployments found");
         return Ok(());
     }
 
@@ -104,29 +109,17 @@ pub async fn run(
     };
 
     println!("Pausing deployments");
-    let mut paused = false;
     for deployment in &deployments {
-        if let Some(node) = &deployment.node_id {
-            if !node.starts_with(PAUSED) {
-                let loc = deployment.locator();
-                let node =
-                    NodeId::new(format!("{}{}", PAUSED, node)).expect("paused_ node id is valid");
-                subgraph_store.reassign_subgraph(&loc, &node)?;
-                println!("  ... paused {}", loc);
-                paused = true;
-            }
-        }
+        pause_or_resume(primary.clone(), &sender, &deployment.locator(), true)?;
     }
 
-    if paused {
-        // There's no good way to tell that a subgraph has in fact stopped
-        // indexing. We sleep and hope for the best.
-        println!(
-            "\nWaiting {}s to make sure pausing was processed",
-            sleep.as_secs()
-        );
-        thread::sleep(sleep);
-    }
+    // There's no good way to tell that a subgraph has in fact stopped
+    // indexing. We sleep and hope for the best.
+    println!(
+        "\nWaiting {}s to make sure pausing was processed",
+        sleep.as_secs()
+    );
+    thread::sleep(sleep);
 
     println!("\nRewinding deployments");
     for deployment in &deployments {
@@ -158,11 +151,7 @@ pub async fn run(
 
     println!("Resuming deployments");
     for deployment in &deployments {
-        if let Some(node) = &deployment.node_id {
-            let loc = deployment.locator();
-            let node = NodeId::new(node.clone()).expect("node id is valid");
-            subgraph_store.reassign_subgraph(&loc, &node)?;
-        }
+        pause_or_resume(primary.clone(), &sender, &deployment.locator(), false)?;
     }
     Ok(())
 }