chore: nit refactoring and fix

PeerDB-io · Nov 14, 2024 · 484bf66 · 484bf66
1 parent 4675941
commit 484bf66
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 23 deletions.
diff --git a/flow/workflows/maintenance_flow.go b/flow/workflows/maintenance_flow.go
@@ -15,24 +15,28 @@ import (
 	"github.com/PeerDB-io/peer-flow/shared"
 )
 
+func getMaintenanceWorkflowOptions(workflowIDPrefix string, taskQueueId shared.TaskQueueID) client.StartWorkflowOptions {
+	maintenanceWorkflowOptions := client.StartWorkflowOptions{
+		WorkflowIDReusePolicy:    tEnums.WORKFLOW_ID_REUSE_POLICY_ALLOW_DUPLICATE,
+		WorkflowIDConflictPolicy: tEnums.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING,
+		TaskQueue:                peerdbenv.PeerFlowTaskQueueName(taskQueueId),
+		ID:                       workflowIDPrefix,
+	}
+	if deploymentUid := peerdbenv.PeerDBDeploymentUID(); deploymentUid != "" {
+		maintenanceWorkflowOptions.ID += "-" + deploymentUid
+	}
+	return maintenanceWorkflowOptions
+}
+
 // RunStartMaintenanceWorkflow is a helper function to start the StartMaintenanceWorkflow with sane defaults
 func RunStartMaintenanceWorkflow(
 	ctx context.Context,
 	temporalClient client.Client,
 	input *protos.StartMaintenanceFlowInput,
 	taskQueueId shared.TaskQueueID,
 ) (client.WorkflowRun, error) {
-	startWorkflowOptions := client.StartWorkflowOptions{
-		// This is to ensure that maintenance workflows are deduped
-		WorkflowIDReusePolicy:    tEnums.WORKFLOW_ID_REUSE_POLICY_ALLOW_DUPLICATE,
-		WorkflowIDConflictPolicy: tEnums.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING,
-		TaskQueue:                peerdbenv.PeerFlowTaskQueueName(taskQueueId),
-	}
-	startWorkflowOptions.ID = "start-maintenance"
-	if deploymentUid := peerdbenv.PeerDBDeploymentUID(); deploymentUid != "" {
-		startWorkflowOptions.ID += "-" + deploymentUid
-	}
-	workflowRun, err := temporalClient.ExecuteWorkflow(ctx, startWorkflowOptions, StartMaintenanceWorkflow, input)
+	workflowOptions := getMaintenanceWorkflowOptions("start-maintenance", taskQueueId)
+	workflowRun, err := temporalClient.ExecuteWorkflow(ctx, workflowOptions, StartMaintenanceWorkflow, input)
 	if err != nil {
 		return nil, err
 	}
@@ -46,18 +50,8 @@ func RunEndMaintenanceWorkflow(
 	input *protos.EndMaintenanceFlowInput,
 	taskQueueId shared.TaskQueueID,
 ) (client.WorkflowRun, error) {
-	startWorkflowOptions := client.StartWorkflowOptions{
-		// This is to ensure that maintenance workflows are deduped
-		WorkflowIDReusePolicy:    tEnums.WORKFLOW_ID_REUSE_POLICY_ALLOW_DUPLICATE,
-		WorkflowIDConflictPolicy: tEnums.WORKFLOW_ID_CONFLICT_POLICY_USE_EXISTING,
-		TaskQueue:                peerdbenv.PeerFlowTaskQueueName(taskQueueId),
-	}
-	startWorkflowOptions.ID = "end-maintenance"
-	if deploymentUid := peerdbenv.PeerDBDeploymentUID(); deploymentUid != "" {
-		startWorkflowOptions.ID += "-" + deploymentUid
-	}
-
-	workflowRun, err := temporalClient.ExecuteWorkflow(ctx, startWorkflowOptions, EndMaintenanceWorkflow, &protos.EndMaintenanceFlowInput{})
+	workflowOptions := getMaintenanceWorkflowOptions("end-maintenance", taskQueueId)
+	workflowRun, err := temporalClient.ExecuteWorkflow(ctx, workflowOptions, EndMaintenanceWorkflow, &protos.EndMaintenanceFlowInput{})
 	if err != nil {
 		return nil, err
 	}

diff --git a/nexus/catalog/migrations/V40__maintenance_flows.sql b/nexus/catalog/migrations/V40__maintenance_flows.sql
@@ -14,6 +14,8 @@ CREATE TABLE IF NOT EXISTS maintenance.maintenance_flows
     to_version      TEXT
 );
 
+CREATE INDEX IF NOT EXISTS idx_maintenance_flows_state ON maintenance.maintenance_flows (state);
+
 CREATE TABLE IF NOT EXISTS maintenance.start_maintenance_outputs
 (
     id             SERIAL PRIMARY KEY,