feat: add graphman unassign/reassign command to graphql api

Author: shiyasmohd

core/graphman/src/commands/deployment/mod.rs

@@ -1,3 +1,5 @@
 pub mod info;
 pub mod pause;
+pub mod reassign;
 pub mod resume;
+pub mod unassign;

core/graphman/src/commands/deployment/reassign.rs

@@ -0,0 +1,96 @@
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use graph::components::store::DeploymentLocator;
+use graph::components::store::StoreEvent;
+use graph::prelude::EntityChange;
+use graph::prelude::NodeId;
+use graph_store_postgres::command_support::catalog;
+use graph_store_postgres::command_support::catalog::Site;
+use graph_store_postgres::connection_pool::ConnectionPool;
+use graph_store_postgres::NotificationSender;
+use thiserror::Error;
+
+use crate::deployment::DeploymentSelector;
+use crate::deployment::DeploymentVersionSelector;
+use crate::GraphmanError;
+
+pub struct ActiveDeployment {
+    locator: DeploymentLocator,
+    site: Site,
+}
+
+#[derive(Debug, Error)]
+pub enum ReassignDeploymentError {
+    #[error("deployment '{0}' is already assigned to '{1}'")]
+    AlreadyAssigned(String, String),
+
+    #[error(transparent)]
+    Common(#[from] GraphmanError),
+}
+
+pub fn load_deployment(
+    primary_pool: ConnectionPool,
+    deployment: &DeploymentSelector,
+) -> Result<ActiveDeployment, ReassignDeploymentError> {
+    let mut primary_conn = primary_pool.get().map_err(GraphmanError::from)?;
+
+    let locator = crate::deployment::load_deployment(
+        &mut primary_conn,
+        deployment,
+        &DeploymentVersionSelector::All,
+    )?
+    .locator();
+
+    let mut catalog_conn = catalog::Connection::new(primary_conn);
+
+    let site = catalog_conn
+        .locate_site(locator.clone())
+        .map_err(GraphmanError::from)?
+        .ok_or_else(|| {
+            GraphmanError::Store(anyhow!("deployment site not found for '{locator}'"))
+        })?;
+
+    Ok(ActiveDeployment { locator, site })
+}
+
+pub fn reassign_deployment(
+    primary_pool: ConnectionPool,
+    notification_sender: Arc<NotificationSender>,
+    deployment: ActiveDeployment,
+    node: &NodeId,
+) -> Result<(), ReassignDeploymentError> {
+    let primary_conn = primary_pool.get().map_err(GraphmanError::from)?;
+    let mut catalog_conn = catalog::Connection::new(primary_conn);
+
+    let changes: Vec<EntityChange> = match catalog_conn
+        .assigned_node(&deployment.site)
+        .map_err(GraphmanError::from)?
+    {
+        Some(curr) => {
+            if &curr == node {
+                vec![]
+            } else {
+                catalog_conn
+                    .reassign_subgraph(&deployment.site, &node)
+                    .map_err(GraphmanError::from)?
+            }
+        }
+        None => catalog_conn
+            .assign_subgraph(&deployment.site, &node)
+            .map_err(GraphmanError::from)?,
+    };
+
+    if changes.is_empty() {
+        return Err(ReassignDeploymentError::AlreadyAssigned(
+            deployment.locator.hash.to_string(),
+            node.to_string(),
+        ));
+    }
+
+    catalog_conn
+        .send_store_event(&notification_sender, &StoreEvent::new(changes))
+        .map_err(GraphmanError::from)?;
+
+    Ok(())
+}

core/graphman/src/commands/deployment/unassign.rs

@@ -0,0 +1,74 @@
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use graph::components::store::StoreEvent;
+use graph_store_postgres::command_support::catalog;
+use graph_store_postgres::command_support::catalog::Site;
+use graph_store_postgres::connection_pool::ConnectionPool;
+use graph_store_postgres::NotificationSender;
+use thiserror::Error;
+
+use crate::deployment::DeploymentSelector;
+use crate::deployment::DeploymentVersionSelector;
+use crate::GraphmanError;
+
+pub struct AssignedDeployment {
+    site: Site,
+}
+
+#[derive(Debug, Error)]
+pub enum UnassignDeploymentError {
+    #[error("deployment '{0}' is already unassigned")]
+    AlreadyUnassigned(String),
+
+    #[error(transparent)]
+    Common(#[from] GraphmanError),
+}
+
+pub fn load_assigned_deployment(
+    primary_pool: ConnectionPool,
+    deployment: &DeploymentSelector,
+) -> Result<AssignedDeployment, UnassignDeploymentError> {
+    let mut primary_conn = primary_pool.get().map_err(GraphmanError::from)?;
+
+    let locator = crate::deployment::load_deployment(
+        &mut primary_conn,
+        deployment,
+        &DeploymentVersionSelector::All,
+    )?
+    .locator();
+
+    let mut catalog_conn = catalog::Connection::new(primary_conn);
+
+    let site = catalog_conn
+        .locate_site(locator.clone())
+        .map_err(GraphmanError::from)?
+        .ok_or_else(|| {
+            GraphmanError::Store(anyhow!("deployment site not found for '{locator}'"))
+        })?;
+
+    match catalog_conn
+        .assigned_node(&site)
+        .map_err(GraphmanError::from)?
+    {
+        Some(_) => Ok(AssignedDeployment { site }),
+        None => Err(UnassignDeploymentError::AlreadyUnassigned(format!(
+            "{:?}",
+            deployment
+        ))),
+    }
+}
+
+pub fn unassign_deployment(
+    primary_pool: ConnectionPool,
+    notification_sender: Arc<NotificationSender>,
+    deployment: AssignedDeployment,
+) -> Result<(), GraphmanError> {
+    let primary_conn = primary_pool.get()?;
+    let mut catalog_conn = catalog::Connection::new(primary_conn);
+
+    let changes = catalog_conn.unassign_subgraph(&deployment.site)?;
+    catalog_conn.send_store_event(&notification_sender, &StoreEvent::new(changes))?;
+
+    Ok(())
+}

server/graphman/src/entities/empty_response.rs

@@ -5,11 +5,15 @@ use async_graphql::SimpleObject;
 #[derive(Clone, Debug, SimpleObject)]
 pub struct EmptyResponse {
     pub success: bool,
+    pub message: Option<String>,
 }
 
 impl EmptyResponse {
     /// Returns a successful response.
-    pub fn new() -> Self {
-        Self { success: true }
+    pub fn new(msg: Option<String>) -> Self {
+        Self {
+            success: true,
+            message: msg,
+        }
     }
 }

server/graphman/src/resolvers/deployment_mutation.rs

@@ -1,8 +1,11 @@
 use std::sync::Arc;
 
+use anyhow::anyhow;
 use async_graphql::Context;
 use async_graphql::Object;
 use async_graphql::Result;
+use graph::prelude::NodeId;
+use graph_store_postgres::command_support::catalog;
 use graph_store_postgres::graphman::GraphmanStore;
 
 use crate::entities::DeploymentSelector;
@@ -11,9 +14,10 @@ use crate::entities::ExecutionId;
 use crate::resolvers::context::GraphmanContext;
 
 mod pause;
+mod reassign;
 mod restart;
 mod resume;
-
+mod unassign;
 pub struct DeploymentMutation;
 
 /// Mutations related to one or multiple deployments.
@@ -30,7 +34,7 @@ impl DeploymentMutation {
 
         pause::run(&ctx, &deployment)?;
 
-        Ok(EmptyResponse::new())
+        Ok(EmptyResponse::new(None))
     }
 
     /// Resumes a deployment that has been previously paused.
@@ -44,7 +48,7 @@ impl DeploymentMutation {
 
         resume::run(&ctx, &deployment)?;
 
-        Ok(EmptyResponse::new())
+        Ok(EmptyResponse::new(None))
     }
 
     /// Pauses a deployment and resumes it after a delay.
@@ -65,4 +69,39 @@ impl DeploymentMutation {
 
         restart::run_in_background(ctx, store, deployment, delay_seconds).await
     }
+
+    /// Unassign a deployment
+    pub async fn unassign(
+        &self,
+        ctx: &Context<'_>,
+        deployment: DeploymentSelector,
+    ) -> Result<EmptyResponse> {
+        let ctx = GraphmanContext::new(ctx)?;
+        let deployment = deployment.try_into()?;
+
+        unassign::run(&ctx, &deployment)?;
+
+        Ok(EmptyResponse::new(None))
+    }
+
+    /// Assign or reassign a deployment
+    pub async fn reassign(
+        &self,
+        ctx: &Context<'_>,
+        deployment: DeploymentSelector,
+        node: String,
+    ) -> Result<EmptyResponse> {
+        let ctx = GraphmanContext::new(ctx)?;
+        let deployment = deployment.try_into()?;
+        let node = NodeId::new(node.clone()).map_err(|()| anyhow!("illegal node id `{}`", node))?;
+        reassign::run(&ctx, &deployment, &node)?;
+
+        let mirror = catalog::Mirror::primary_only(ctx.primary_pool);
+        let count = mirror.assignments(&node)?.len();
+        if count == 1 {
+            Ok(EmptyResponse::new(Some(format!("warning: this is the only deployment assigned to '{}'. Are you sure it is spelled correctly?",node.as_str()))))
+        } else {
+            Ok(EmptyResponse::new(None))
+        }
+    }
 }

server/graphman/src/resolvers/deployment_mutation/reassign.rs

@@ -0,0 +1,19 @@
+use async_graphql::Result;
+use graph::prelude::NodeId;
+use graphman::commands::deployment::reassign::load_deployment;
+use graphman::commands::deployment::reassign::reassign_deployment;
+use graphman::deployment::DeploymentSelector;
+
+use crate::resolvers::context::GraphmanContext;
+
+pub fn run(ctx: &GraphmanContext, deployment: &DeploymentSelector, node: &NodeId) -> Result<()> {
+    let deployment = load_deployment(ctx.primary_pool.clone(), deployment)?;
+    reassign_deployment(
+        ctx.primary_pool.clone(),
+        ctx.notification_sender.clone(),
+        deployment,
+        &node,
+    )?;
+
+    Ok(())
+}

server/graphman/src/resolvers/deployment_mutation/unassign.rs

@@ -0,0 +1,17 @@
+use async_graphql::Result;
+use graphman::commands::deployment::unassign::load_assigned_deployment;
+use graphman::commands::deployment::unassign::unassign_deployment;
+use graphman::deployment::DeploymentSelector;
+
+use crate::resolvers::context::GraphmanContext;
+
+pub fn run(ctx: &GraphmanContext, deployment: &DeploymentSelector) -> Result<()> {
+    let deployment = load_assigned_deployment(ctx.primary_pool.clone(), deployment)?;
+    unassign_deployment(
+        ctx.primary_pool.clone(),
+        ctx.notification_sender.clone(),
+        deployment,
+    )?;
+
+    Ok(())
+}