feat(maintenance): use proto field descriptor (#3531)

Author: iamKunalGupta

flow/activities/maintenance_activity.go

@@ -15,6 +15,8 @@ import (
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/sdk/activity"
 	"go.temporal.io/sdk/client"
+	proto2 "google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/PeerDB-io/peerdb/flow/alerting"
@@ -112,13 +114,30 @@ func (a *MaintenanceActivity) WaitForRunningSnapshotsAndIntermediateStates(
 	return mirrors, nil
 }
 
-var waitStatuses map[protos.FlowStatus]struct{} = map[protos.FlowStatus]struct{}{
-	protos.FlowStatus_STATUS_SNAPSHOT:  {},
-	protos.FlowStatus_STATUS_SETUP:     {},
-	protos.FlowStatus_STATUS_RESYNC:    {},
-	protos.FlowStatus_STATUS_UNKNOWN:   {},
-	protos.FlowStatus_STATUS_PAUSING:   {},
-	protos.FlowStatus_STATUS_MODIFYING: {},
+var waitStatuses = buildWaitStatuses()
+
+func buildWaitStatuses() map[protos.FlowStatus]struct{} {
+	waitStatuses := make(map[protos.FlowStatus]struct{})
+
+	for index := range protos.FlowStatus_name {
+		flowStatus := protos.FlowStatus(index)
+
+		// Get the enum value descriptor
+		enumValueDesc := flowStatus.Descriptor().Values().ByNumber(protoreflect.EnumNumber(index))
+		if enumValueDesc == nil {
+			continue
+		}
+
+		// Get the extension value from the enum value options
+		if proto2.HasExtension(enumValueDesc.Options(), protos.E_PeerdbMaintenanceWait) {
+			waitValue := proto2.GetExtension(enumValueDesc.Options(), protos.E_PeerdbMaintenanceWait)
+			if boolVal, ok := waitValue.(bool); ok && boolVal {
+				waitStatuses[flowStatus] = struct{}{}
+			}
+		}
+	}
+
+	return waitStatuses
 }
 
 func (a *MaintenanceActivity) checkAndWaitIfNeeded(

flow/activities/maintenance_test.go

@@ -0,0 +1,16 @@
+package activities
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+)
+
+func TestMaintenanceStatusesAreCorrectlyPopulated(t *testing.T) {
+	statuses := buildWaitStatuses()
+	t.Logf("statuses: %v", statuses)
+	assert.Contains(t, statuses, protos.FlowStatus_STATUS_SNAPSHOT, "Snapshot should be in statuses")
+	assert.NotContains(t, statuses, protos.FlowStatus_STATUS_RUNNING, "Running should not be in statuses")
+}

protos/flow.proto

@@ -1,8 +1,15 @@
 syntax = "proto3";
 
 import "google/protobuf/timestamp.proto";
+import "google/protobuf/descriptor.proto";
 import "peers.proto";
 
+extend google.protobuf.EnumValueOptions {
+  // If the flow is in this state, then maintenance should wait for it to transition to another non-wait state
+  // This is to make sure we have only either terminal states or running state which we can pause
+  optional bool peerdb_maintenance_wait = 16551843;
+}
+
 package peerdb_flow;
 
 message AlertInput {
@@ -400,20 +407,20 @@ message GetOpenConnectionsForUserResult {
 // UI reads current workflow status and also requests status changes using same enum
 // see flow/cmd/handler.go FlowStateChange
 enum FlowStatus {
-  STATUS_UNKNOWN = 0;
+  STATUS_UNKNOWN = 0 [(peerdb_maintenance_wait) = true];
   STATUS_RUNNING = 1;
   STATUS_PAUSED = 2;
-  STATUS_PAUSING = 3;
+  STATUS_PAUSING = 3 [(peerdb_maintenance_wait) = true];;
   // not reachable in QRep mirrors
-  STATUS_SETUP = 4;
+  STATUS_SETUP = 4 [(peerdb_maintenance_wait) = true];
   // not reachable in QRep mirrors
-  STATUS_SNAPSHOT = 5;
+  STATUS_SNAPSHOT = 5 [(peerdb_maintenance_wait) = true];;
   STATUS_TERMINATING = 6;
   STATUS_TERMINATED = 7;
   STATUS_COMPLETED = 8;
-  STATUS_RESYNC = 9;
+  STATUS_RESYNC = 9 [(peerdb_maintenance_wait) = true];;
   STATUS_FAILED = 10;
-  STATUS_MODIFYING = 11;
+  STATUS_MODIFYING = 11 [(peerdb_maintenance_wait) = true];;
 }
 
 message CDCFlowConfigUpdate {