Flow API: Custom resume endpoint, polished mirror status endpoint, auth to all endpoints

docker-compose-dev.yml

@@ -131,6 +131,7 @@ services:
       - 8113:8113
     environment:
       <<: [*catalog-config, *flow-worker-env, *minio-config]
+      PEERDB_PASSWORD:
     depends_on:
       temporal-admin-tools:
         condition: service_healthy

docker-compose.yml

@@ -117,6 +117,7 @@ services:
       - 8113:8113
     environment:
       <<: [*catalog-config, *flow-worker-env, *minio-config]
+      PEERDB_PASSWORD:
     depends_on:
       temporal-admin-tools:
         condition: service_healthy

flow/cmd/api.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/PeerDB-io/peer-flow/generated/protos"
 	"github.com/PeerDB-io/peer-flow/logger"
+	"github.com/PeerDB-io/peer-flow/middleware"
 	"github.com/PeerDB-io/peer-flow/peerdbenv"
 	"github.com/PeerDB-io/peer-flow/shared"
 	peerflow "github.com/PeerDB-io/peer-flow/workflows"
@@ -115,7 +116,16 @@ func APIMain(ctx context.Context, args *APIServerParams) error {
 		return fmt.Errorf("unable to create Temporal client: %w", err)
 	}
 
-	grpcServer := grpc.NewServer()
+	healthMethods := []string{
+		grpc_health_v1.Health_Check_FullMethodName,
+		grpc_health_v1.Health_Watch_FullMethodName,
+	}
+	grpcServer := grpc.NewServer(
+		grpc.ChainUnaryInterceptor(
+			middleware.CreateRequestLoggingInterceptor(healthMethods),
+			middleware.CreateAuthServerInterceptor(ctx, peerdbenv.PeerDBPassword(), healthMethods),
+		),
+	)
 
 	catalogConn, err := peerdbenv.GetCatalogConnectionPoolFromEnv(ctx)
 	if err != nil {

flow/cmd/custom_sync.go

@@ -0,0 +1,109 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+
+	"github.com/PeerDB-io/peer-flow/generated/protos"
+	"github.com/PeerDB-io/peer-flow/peerdbenv"
+	peerflow "github.com/PeerDB-io/peer-flow/workflows"
+)
+
+const peerdbPauseGuideDocLink = "https://docs.peerdb.io/features/pause-mirror"
+
+func (h *FlowRequestHandler) CustomSyncFlow(
+	ctx context.Context, req *protos.CreateCustomSyncRequest,
+) (*protos.CreateCustomSyncResponse, error) {
+	errResponse := &protos.CreateCustomSyncResponse{
+		FlowJobName:   req.FlowJobName,
+		NumberOfSyncs: 0,
+		ErrorMessage:  "error while processing request",
+		Ok:            false,
+	}
+
+	if req.FlowJobName == "" {
+		errResponse.ErrorMessage = "Mirror name cannot be empty."
+		return errResponse, nil
+	}
+
+	if req.NumberOfSyncs <= 0 || req.NumberOfSyncs > peerflow.MaxSyncsPerCdcFlow {
+		slog.Error("Invalid sync number request",
+			slog.Any("requested_number_of_syncs", req.NumberOfSyncs))
+		errResponse.ErrorMessage = fmt.Sprintf("Sync number request must be between 1 and %d (inclusive). Requested number: %d",
+			peerflow.MaxSyncsPerCdcFlow, req.NumberOfSyncs)
+		return errResponse, nil
+	}
+
+	mirrorExists, err := h.CheckIfMirrorNameExists(ctx, req.FlowJobName)
+	if err != nil {
+		slog.Error("Server error: unable to check if mirror exists", slog.Any("error", err))
+		errResponse.ErrorMessage = "Server error: unable to check if mirror " + req.FlowJobName + " exists."
+		return errResponse, nil
+	}
+	if !mirrorExists {
+		slog.Error("Mirror does not exist", slog.Any("mirror_name", req.FlowJobName))
+		errResponse.ErrorMessage = fmt.Sprintf("Mirror %s does not exist", req.FlowJobName)
+		return errResponse, nil
+	}
+
+	mirrorStatusResponse, _ := h.MirrorStatus(ctx, &protos.MirrorStatusRequest{
+		FlowJobName: req.FlowJobName,
+	})
+	if mirrorStatusResponse.ErrorMessage != "" {
+		slog.Error("Server error: unable to check the status of mirror",
+			slog.Any("mirror", req.FlowJobName),
+			slog.Any("error", mirrorStatusResponse.ErrorMessage))
+		errResponse.ErrorMessage = fmt.Sprintf("Server error: unable to check the status of mirror %s: %s",
+			req.FlowJobName, mirrorStatusResponse.ErrorMessage)
+		return errResponse, nil
+	}
+
+	if mirrorStatusResponse.CurrentFlowState != protos.FlowStatus_STATUS_PAUSED {
+		slog.Error("Mirror is not paused", slog.Any("mirror", req.FlowJobName))
+		errResponse.ErrorMessage = fmt.Sprintf(`Requested mirror %s is not paused. This is a requirement.
+		The mirror can be paused via PeerDB UI. Please follow %s`,
+			req.FlowJobName, peerdbPauseGuideDocLink)
+		return errResponse, nil
+	}
+
+	// Parallel sync-normalise should not be enabled
+	parallelSyncNormaliseEnabled := peerdbenv.PeerDBEnableParallelSyncNormalize()
+	if parallelSyncNormaliseEnabled {
+		errResponse.ErrorMessage = "Parallel sync-normalise is enabled. Please contact PeerDB support to disable it to proceed."
+		return errResponse, nil
+	}
+	// ---- REQUEST VALIDATED ----
+
+	// Resume mirror with custom sync number
+	_, err = h.FlowStateChange(ctx, &protos.FlowStateChangeRequest{
+		FlowJobName:        req.FlowJobName,
+		RequestedFlowState: protos.FlowStatus_STATUS_RUNNING,
+		FlowConfigUpdate: &protos.FlowConfigUpdate{
+			Update: &protos.FlowConfigUpdate_CdcFlowConfigUpdate{
+				CdcFlowConfigUpdate: &protos.CDCFlowConfigUpdate{
+					NumberOfSyncs: req.NumberOfSyncs,
+				},
+			},
+		},
+	})
+	if err != nil {
+		slog.Error("Unable to kick off custom sync for mirror",
+			slog.Any("mirror", req.FlowJobName),
+			slog.Any("error", err))
+		errResponse.ErrorMessage = fmt.Sprintf("Unable to kick off sync for mirror %s:%s",
+			req.FlowJobName, err.Error())
+		return errResponse, nil
+	}
+
+	slog.Info("Custom sync started for mirror",
+		slog.String("mirror", req.FlowJobName),
+		slog.Int("number_of_syncs", int(req.NumberOfSyncs)))
+
+	return &protos.CreateCustomSyncResponse{
+		FlowJobName:   req.FlowJobName,
+		NumberOfSyncs: req.NumberOfSyncs,
+		ErrorMessage:  "",
+		Ok:            true,
+	}, nil
+}

flow/cmd/mirror_status.go

@@ -19,58 +19,87 @@ func (h *FlowRequestHandler) MirrorStatus(
 	ctx context.Context,
 	req *protos.MirrorStatusRequest,
 ) (*protos.MirrorStatusResponse, error) {
-	slog.Info("Mirror status endpoint called", slog.String(string(shared.FlowNameKey), req.FlowJobName))
-	cdcFlow, err := h.isCDCFlow(ctx, req.FlowJobName)
-	if err != nil {
-		slog.Error("unable to query flow", slog.Any("error", err))
-		return &protos.MirrorStatusResponse{
-			ErrorMessage: "unable to query flow: " + err.Error(),
-		}, nil
-	}
+	slog.Info("Mirror status endpoint called",
+		slog.Bool("includeFlowInfo", req.IncludeFlowInfo),
+		slog.String(string(shared.FlowNameKey), req.FlowJobName))
 
 	workflowID, err := h.getWorkflowID(ctx, req.FlowJobName)
 	if err != nil {
-		return nil, err
+		return &protos.MirrorStatusResponse{
+			FlowJobName:      req.FlowJobName,
+			CurrentFlowState: protos.FlowStatus_STATUS_UNKNOWN,
+			ErrorMessage:     "unable to get the workflow ID of mirror " + req.FlowJobName,
+			Ok:               false,
+		}, nil
 	}
 
 	currState, err := h.getWorkflowStatus(ctx, workflowID)
 	if err != nil {
 		return &protos.MirrorStatusResponse{
-			ErrorMessage: "unable to get flow state: " + err.Error(),
+			FlowJobName:      req.FlowJobName,
+			CurrentFlowState: protos.FlowStatus_STATUS_UNKNOWN,
+			ErrorMessage:     "unable to get the running status of mirror " + req.FlowJobName,
+			Ok:               false,
 		}, nil
 	}
 
-	if cdcFlow {
-		cdcStatus, err := h.CDCFlowStatus(ctx, req)
+	if req.IncludeFlowInfo {
+		cdcFlow, err := h.isCDCFlow(ctx, req.FlowJobName)
 		if err != nil {
+			slog.Error("unable to query flow", slog.Any("error", err))
 			return &protos.MirrorStatusResponse{
-				ErrorMessage: "unable to query flow: " + err.Error(),
+				FlowJobName:      req.FlowJobName,
+				CurrentFlowState: protos.FlowStatus_STATUS_UNKNOWN,
+				ErrorMessage:     "unable to determine if mirror" + req.FlowJobName + "is of type CDC.",
+				Ok:               false,
 			}, nil
 		}
+		if cdcFlow {
+			cdcStatus, err := h.CDCFlowStatus(ctx, req)
+			if err != nil {
+				return &protos.MirrorStatusResponse{
+					FlowJobName:      req.FlowJobName,
+					CurrentFlowState: protos.FlowStatus_STATUS_UNKNOWN,
+					ErrorMessage:     "unable to obtain CDC information for mirror " + req.FlowJobName,
+					Ok:               false,
+				}, nil
+			}
 
-		return &protos.MirrorStatusResponse{
-			FlowJobName: req.FlowJobName,
-			Status: &protos.MirrorStatusResponse_CdcStatus{
-				CdcStatus: cdcStatus,
-			},
-			CurrentFlowState: currState,
-		}, nil
-	} else {
-		qrepStatus, err := h.QRepFlowStatus(ctx, req)
-		if err != nil {
 			return &protos.MirrorStatusResponse{
-				ErrorMessage: "unable to query flow: " + err.Error(),
+				FlowJobName: req.FlowJobName,
+				Status: &protos.MirrorStatusResponse_CdcStatus{
+					CdcStatus: cdcStatus,
+				},
+				CurrentFlowState: currState,
+				Ok:               true,
 			}, nil
-		}
+		} else {
+			qrepStatus, err := h.QRepFlowStatus(ctx, req)
+			if err != nil {
+				return &protos.MirrorStatusResponse{
+					FlowJobName:      req.FlowJobName,
+					CurrentFlowState: protos.FlowStatus_STATUS_UNKNOWN,
+					ErrorMessage:     "unable to obtain snapshot information for mirror " + req.FlowJobName,
+					Ok:               false,
+				}, nil
+			}
 
-		return &protos.MirrorStatusResponse{
-			FlowJobName: req.FlowJobName,
-			Status: &protos.MirrorStatusResponse_QrepStatus{
-				QrepStatus: qrepStatus,
-			},
-			CurrentFlowState: currState,
-		}, nil
+			return &protos.MirrorStatusResponse{
+				FlowJobName: req.FlowJobName,
+				Status: &protos.MirrorStatusResponse_QrepStatus{
+					QrepStatus: qrepStatus,
+				},
+				CurrentFlowState: currState,
+				Ok:               true,
+			}, nil
+		}
 	}
+
+	return &protos.MirrorStatusResponse{
+		FlowJobName:      req.FlowJobName,
+		CurrentFlowState: currState,
+		Ok:               true,
+	}, nil
 }
 
 func (h *FlowRequestHandler) CDCFlowStatus(

flow/cmd/validate_mirror.go

@@ -127,7 +127,7 @@ func (h *FlowRequestHandler) CheckIfMirrorNameExists(ctx context.Context, mirror
 	var nameExists pgtype.Bool
 	err := h.pool.QueryRow(ctx, "SELECT EXISTS(SELECT * FROM flows WHERE name = $1)", mirrorName).Scan(&nameExists)
 	if err != nil {
-		return true, fmt.Errorf("failed to check if mirror name exists: %v", err)
+		return false, fmt.Errorf("failed to check if mirror name exists: %v", err)
 	}
 
 	return nameExists.Bool, nil

flow/middleware/auth.go

@@ -0,0 +1,45 @@
+package middleware
+
+import (
+	"context"
+	"encoding/base64"
+	"log/slog"
+	"strings"
+
+	"golang.org/x/crypto/bcrypt"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/status"
+
+	"github.com/PeerDB-io/peer-flow/logger"
+)
+
+// Authorize checks the authorization metadata and compares the incoming bearer token with the plaintext
+func Authorize(ctx context.Context, plaintext string) (context.Context, error) {
+	md, _ := metadata.FromIncomingContext(ctx)
+	if len(md["authorization"]) == 0 {
+		return nil, status.Errorf(codes.Unauthenticated, "Authorization token is required")
+	}
+	headerValue := md["authorization"][0]
+	bearerPrefix := "Bearer "
+	if !strings.HasPrefix(headerValue, bearerPrefix) {
+		return nil, status.Errorf(codes.Unauthenticated, "Unsupported authorization type")
+	}
+	base64Token := strings.TrimPrefix(headerValue, bearerPrefix)
+	if base64Token == "" {
+		return nil, status.Errorf(codes.Unauthenticated, "Authorization token is required")
+	}
+	// Always a good practice to have the actual token in base64
+	tokenBytes, err := base64.StdEncoding.DecodeString(base64Token)
+	if err != nil {
+		logger.LoggerFromCtx(ctx).Warn("Error decoding token", slog.String("token", base64Token), slog.Any("error", err))
+		return nil, status.Errorf(codes.Unauthenticated, "Authentication failed")
+	}
+	token := string(tokenBytes)
+	err = bcrypt.CompareHashAndPassword([]byte(token), []byte(plaintext))
+	if err != nil {
+		logger.LoggerFromCtx(ctx).Warn("Error validating token", slog.String("token", token), slog.Any("error", err))
+		return nil, status.Errorf(codes.Unauthenticated, "Authentication failed")
+	}
+	return ctx, nil
+}