Implement granular status in cdc and snapshot

Author: ilidemi

flow/activities/flowable.go

@@ -11,6 +11,8 @@ import (
 
 	"github.com/PeerDB-io/peerdb/flow/alerting"
 	"github.com/PeerDB-io/peerdb/flow/connectors"
+	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
+	"github.com/PeerDB-io/peerdb/flow/connectors/utils/monitoring"
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
 	"github.com/PeerDB-io/peerdb/flow/internal"
 	"github.com/PeerDB-io/peerdb/flow/shared"
@@ -36,7 +38,7 @@ type SnapshotActivity struct {
 }
 
 // closes the slot signal
-func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName string) error {
+func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName string) {
 	a.SnapshotStatesMutex.Lock()
 	defer a.SnapshotStatesMutex.Unlock()
 
@@ -48,7 +50,29 @@ func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName s
 		delete(a.SlotSnapshotStates, flowJobName)
 	}
 	a.Alerter.LogFlowEvent(ctx, flowJobName, "Ended Snapshot Flow Job")
+}
+
+func (a *SnapshotActivity) InitializeSnapshot(
+	ctx context.Context, flowName string,
+) (int32, error) {
+	ctx = context.WithValue(ctx, shared.FlowNameKey, flowName)
+	logger := internal.LoggerFromCtx(ctx)
+
+	snapshotID, err := monitoring.InitializeSnapshot(ctx, logger, a.CatalogPool, flowName)
+	if err != nil {
+		return -1, a.Alerter.LogFlowErrorNoStatus(ctx, flowName, err)
+	}
+	return snapshotID, nil
+}
 
+func (a *SnapshotActivity) FinishSnapshot(
+	ctx context.Context, flowName string, snapshotID int32,
+) error {
+	ctx = context.WithValue(ctx, shared.FlowNameKey, flowName)
+	logger := internal.LoggerFromCtx(ctx)
+	if err := monitoring.FinishSnapshot(ctx, logger, a.CatalogPool, flowName, snapshotID); err != nil {
+		return a.Alerter.LogFlowSnapshotError(ctx, flowName, snapshotID, err)
+	}
 	return nil
 }
 
@@ -63,7 +87,8 @@ func (a *SnapshotActivity) SetupReplication(
 
 	conn, err := connectors.GetByNameAs[connectors.CDCPullConnectorCore](ctx, nil, a.CatalogPool, config.PeerName)
 	if err != nil {
-		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("failed to get connector: %w", err))
+		conErr := fmt.Errorf("failed to get connector: %w", err)
+		return nil, a.Alerter.LogFlowSnapshotError(ctx, config.FlowJobName, config.SnapshotId, conErr)
 	}
 
 	logger.Info("waiting for slot to be created...")
@@ -72,7 +97,8 @@ func (a *SnapshotActivity) SetupReplication(
 	if err != nil {
 		connectors.CloseConnector(ctx, conn)
 		// it is important to close the connection here as it is not closed in CloseSlotKeepAlive
-		return nil, a.Alerter.LogFlowError(ctx, config.FlowJobName, fmt.Errorf("slot error: %w", err))
+		slotErr := fmt.Errorf("slot error: %w", err)
+		return nil, a.Alerter.LogFlowSnapshotError(ctx, config.FlowJobName, config.SnapshotId, slotErr)
 	} else if slotInfo.Conn == nil && slotInfo.SlotName == "" {
 		connectors.CloseConnector(ctx, conn)
 		logger.Info("replication setup without slot")
@@ -99,20 +125,22 @@ func (a *SnapshotActivity) SetupReplication(
 	}, nil
 }
 
-func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, peer string, env map[string]string) error {
+func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, peer string, env map[string]string, snapshotID int32) error {
 	shutdown := heartbeatRoutine(ctx, func() string {
 		return "maintaining transaction snapshot"
 	})
 	defer shutdown()
 	conn, err := connectors.GetByNameAs[connectors.CDCPullConnector](ctx, nil, a.CatalogPool, peer)
 	if err != nil {
-		return a.Alerter.LogFlowError(ctx, sessionID, err)
+		getErr := fmt.Errorf("failed to get connector: %w", err)
+		return a.Alerter.LogFlowSnapshotError(ctx, sessionID, snapshotID, getErr)
 	}
 	defer connectors.CloseConnector(ctx, conn)
 
 	exportSnapshotOutput, tx, err := conn.ExportTxSnapshot(ctx, env)
 	if err != nil {
-		return err
+		exportErr := fmt.Errorf("failed to export tx snapshot: %w", err)
+		return a.Alerter.LogFlowSnapshotError(ctx, sessionID, snapshotID, exportErr)
 	}
 
 	a.SnapshotStatesMutex.Lock()
@@ -134,7 +162,10 @@ func (a *SnapshotActivity) MaintainTx(ctx context.Context, sessionID string, pee
 			a.SnapshotStatesMutex.Lock()
 			delete(a.TxSnapshotStates, sessionID)
 			a.SnapshotStatesMutex.Unlock()
-			return conn.FinishExport(tx)
+			if err := conn.FinishExport(tx); err != nil {
+				finishErr := fmt.Errorf("failed to finish export: %w", err)
+				return a.Alerter.LogFlowSnapshotError(ctx, sessionID, snapshotID, finishErr)
+			}
 		}
 		time.Sleep(time.Minute)
 	}
@@ -166,6 +197,40 @@ func (a *SnapshotActivity) LoadTableSchema(
 	ctx context.Context,
 	flowName string,
 	tableName string,
+	snapshotID int32,
 ) (*protos.TableSchema, error) {
-	return internal.LoadTableSchemaFromCatalog(ctx, a.CatalogPool, flowName, tableName)
+	schema, err := internal.LoadTableSchemaFromCatalog(ctx, a.CatalogPool, flowName, tableName)
+	if err != nil {
+		loadErr := fmt.Errorf("failed to load schema from catalog: %w", err)
+		return nil, a.Alerter.LogFlowSnapshotError(ctx, flowName, snapshotID, loadErr)
+	}
+	return schema, nil
+}
+
+func (a *SnapshotActivity) ParseSchemaTable(
+	ctx context.Context,
+	flowName string,
+	tableName string,
+	snapshotID int32,
+) (*utils.SchemaTable, error) {
+	parsedTable, err := utils.ParseSchemaTable(tableName)
+	if err != nil {
+		parseErr := fmt.Errorf("failed to parse schema table: %w", err)
+		return nil, a.Alerter.LogFlowSnapshotError(ctx, flowName, snapshotID, parseErr)
+	}
+	return parsedTable, nil
+}
+
+func (a *SnapshotActivity) GetPeerType(
+	ctx context.Context,
+	flowName string,
+	peerName string,
+	snapshotID int32,
+) (protos.DBType, error) {
+	dbtype, err := connectors.LoadPeerType(ctx, a.CatalogPool, peerName)
+	if err != nil {
+		peerErr := fmt.Errorf("failed to get peer type: %w", err)
+		return 0, a.Alerter.LogFlowSnapshotError(ctx, flowName, snapshotID, peerErr)
+	}
+	return dbtype, nil
 }

flow/activities/flowable_core.go

@@ -433,19 +433,34 @@ func (a *Alerter) LogNonFlowEvent(ctx context.Context, eventType telemetry.Event
 	a.sendTelemetryMessage(ctx, logger, string(eventType)+":"+key, message, level)
 }
 
-// logFlowErrorInternal pushes the error to the errors table and emits a metric as well as a telemetry message
-func (a *Alerter) logFlowErrorInternal(ctx context.Context, flowName, errorType string, inErr error, loggerFunc func(string, ...any)) {
-	logger := internal.LoggerFromCtx(ctx)
+// logFlowErrorImpl pushes the error to the errors table and emits a metric as well as a telemetry message
+func (a *Alerter) logFlowErrorImpl(ctx context.Context, flowName, errorType string, inErr error, logger log.Logger, loggerFunc func(string, ...any)) {
 	inErrWithStack := fmt.Sprintf("%+v", inErr)
 	loggerFunc(inErr.Error(), slog.String("stack", inErrWithStack))
-	if _, err := a.CatalogPool.Exec(
-		ctx, "INSERT INTO peerdb_stats.flow_errors(flow_name,error_message,error_type) VALUES($1,$2,$3)",
-		flowName, inErrWithStack, errorType,
-	); err != nil {
-		logger.Error("failed to insert flow error", slog.Any("error", err))
-		return
+	retryInterval := time.Second
+	for {
+		if _, err := a.CatalogPool.Exec(
+			ctx, "INSERT INTO peerdb_stats.flow_errors(flow_name,error_message,error_type) VALUES($1,$2,$3)",
+			flowName, inErrWithStack, errorType,
+		); err != nil {
+			insertErr := shared.LogError(logger, fmt.Errorf("failed to insert flow error: %w", err))
+			errInfo := ErrorInfo{Source: ErrorSourcePostgresCatalog, Code: "UNKNOWN"}
+			a.emitClassifiedError(ctx, logger, flowName, ErrorInternal, errInfo, insertErr, insertErr.Error(), loggerFunc)
+			time.Sleep(retryInterval)
+			retryInterval = min(retryInterval*2, time.Minute)
+			continue
+		}
+		break
 	}
 
+	errorClass, errInfo := GetErrorClass(ctx, inErr)
+	a.emitClassifiedError(ctx, logger, flowName, errorClass, errInfo, inErr, inErrWithStack, loggerFunc)
+}
+
+func (a *Alerter) emitClassifiedError(
+	ctx context.Context, logger log.Logger, flowName string, errorClass ErrorClass, errInfo ErrorInfo,
+	inErr error, inErrWithStack string, loggerFunc func(string, ...any),
+) {
 	var tags []string
 	if errors.Is(inErr, context.Canceled) {
 		tags = append(tags, string(shared.ErrTypeCanceled))
@@ -477,14 +492,11 @@ func (a *Alerter) logFlowErrorInternal(ctx context.Context, flowName, errorType
 	if errors.As(inErr, &sshErr) {
 		tags = append(tags, string(shared.ErrTypeNet))
 	}
-
-	errorClass, errInfo := GetErrorClass(ctx, inErr)
 	tags = append(tags, "errorClass:"+errorClass.String(), "errorAction:"+errorClass.ErrorAction().String())
-
 	if !internal.PeerDBTelemetryErrorActionBasedAlertingEnabled() || errorClass.ErrorAction() == NotifyTelemetry {
-		// Warnings alert us just like errors until there's a customer warning system
 		a.sendTelemetryMessage(ctx, logger, flowName, inErrWithStack, telemetry.ERROR, tags...)
 	}
+
 	loggerFunc(fmt.Sprintf("Emitting classified error '%s'", inErr.Error()),
 		slog.Any("error", inErr),
 		slog.Any("errorClass", errorClass),
@@ -500,15 +512,55 @@ func (a *Alerter) logFlowErrorInternal(ctx context.Context, flowName, errorType
 	a.otelManager.Metrics.ErrorEmittedGauge.Record(ctx, 1, errorAttributeSet)
 }
 
-func (a *Alerter) LogFlowError(ctx context.Context, flowName string, inErr error) error {
+func (a *Alerter) LogFlowErrorNoStatus(ctx context.Context, flowName string, inErr error) error {
+	logger := internal.LoggerFromCtx(ctx)
+	// TODO check that this one just logs without updating status
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
+	return inErr
+}
+
+func (a *Alerter) LogFlowSyncError(ctx context.Context, flowName string, batchID int64, inErr error) error {
+	logger := internal.LoggerFromCtx(ctx)
+	// TODO use batchID
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
+	return inErr
+}
+
+func (a *Alerter) LogFlowNormalizeError(ctx context.Context, flowName string, batchID int64, inErr error) error {
+	logger := internal.LoggerFromCtx(ctx)
+	// TODO use batchID
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
+	return inErr
+}
+
+func (a *Alerter) LogFlowSnapshotError(ctx context.Context, flowName string, snapshotID int32, inErr error) error {
+	logger := internal.LoggerFromCtx(ctx)
+	// TODO use snapshotID
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
+	return inErr
+}
+
+func (a *Alerter) LogFlowSnapshotQRepError(
+	ctx context.Context, flowName string, snapshotID int32, qRepRunID string, inErr error,
+) error {
+	logger := internal.LoggerFromCtx(ctx)
+	// TODO use snapshotID, qRepRunID
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
+	return inErr
+}
+
+func (a *Alerter) LogFlowSnapshotPartitionError(
+	ctx context.Context, flowName string, snapshotID int32, partitionID string, inErr error,
+) error {
 	logger := internal.LoggerFromCtx(ctx)
-	a.logFlowErrorInternal(ctx, flowName, "error", inErr, logger.Error)
+	// TODO use snapshotID, partitionID
+	a.logFlowErrorImpl(ctx, flowName, "error", inErr, logger, logger.Error)
 	return inErr
 }
 
 func (a *Alerter) LogFlowWarning(ctx context.Context, flowName string, inErr error) {
 	logger := internal.LoggerFromCtx(ctx)
-	a.logFlowErrorInternal(ctx, flowName, "warn", inErr, logger.Warn)
+	a.logFlowErrorImpl(ctx, flowName, "warn", inErr, logger, logger.Warn)
 }
 
 func (a *Alerter) LogFlowEvent(ctx context.Context, flowName string, info string) {

flow/activities/snapshot_activity.go

@@ -205,6 +205,30 @@ func GetErrorClass(ctx context.Context, err error) (ErrorClass, ErrorInfo) {
 		}
 	}
 
+	var dynamicConfError *exceptions.DynamicConfError
+	if errors.As(err, &dynamicConfError) {
+		return ErrorInternal, ErrorInfo{
+			Source: ErrorSourceOther,
+			Code:   "UNKNOWN",
+		}
+	}
+
+	var protoMarshalError *exceptions.ProtoMarshalError
+	if errors.As(err, &protoMarshalError) {
+		return ErrorInternal, ErrorInfo{
+			Source: ErrorSourceOther,
+			Code:   "UNKNOWN",
+		}
+	}
+
+	var protoUnmarshalError *exceptions.ProtoUnmarshalError
+	if errors.As(err, &protoUnmarshalError) {
+		return ErrorInternal, ErrorInfo{
+			Source: ErrorSourceOther,
+			Code:   "UNKNOWN",
+		}
+	}
+
 	if errors.Is(err, context.Canceled) {
 		// Generally happens during workflow cancellation
 		return ErrorIgnoreContextCancelled, ErrorInfo{

flow/alerting/alerting.go

@@ -155,7 +155,7 @@ func TestClickHousePushingToViewShouldBeMvError(t *testing.T) {
 		is not supported: while converting source column created_at to destination column created_at:
 		while pushing to view db_name.hello_mv`,
 	}
-	errorClass, errInfo := GetErrorClass(t.Context(), exceptions.NewNormalizationError(fmt.Errorf("error in WAL: %w", err)))
+	errorClass, errInfo := GetErrorClass(t.Context(), exceptions.NewNormalizationError("error in WAL: %w", err))
 	assert.Equal(t, ErrorNotifyMVOrView, errorClass, "Unexpected error class")
 	assert.Equal(t, ErrorInfo{
 		Source: ErrorSourceClickHouse,
@@ -191,8 +191,8 @@ func TestClickHouseChaoticNormalizeErrorShouldBeNotifyMVNow(t *testing.T) {
 				Left key __table1.column_2 type String. Right key __table2.column_1 type Int64`,
 	}
 	errorClass, errInfo := GetErrorClass(t.Context(),
-		exceptions.NewNormalizationError(fmt.Errorf(`Normalization Error: failed to normalize records:
-		 error while inserting into normalized table table_A: %w`, err)))
+		exceptions.NewNormalizationError(`Normalization Error: failed to normalize records:
+		 error while inserting into normalized table table_A: %w`, err))
 	assert.Equal(t, ErrorNotifyMVOrView, errorClass, "Unexpected error class")
 	assert.Equal(t, ErrorInfo{
 		Source: ErrorSourceClickHouse,
@@ -382,7 +382,7 @@ func TestClickHouseUnknownTableShouldBeDestinationModified(t *testing.T) {
 		Message: "Table abc does not exist.",
 	}
 	errorClass, errInfo := GetErrorClass(t.Context(),
-		exceptions.NewNormalizationError(fmt.Errorf("failed to normalize records: %w", err)))
+		exceptions.NewNormalizationError("failed to normalize records: %w", err))
 	assert.Equal(t, ErrorNotifyDestinationModified, errorClass, "Unexpected error class")
 	assert.Equal(t, ErrorInfo{
 		Source: ErrorSourceClickHouse,
@@ -398,7 +398,7 @@ func TestClickHouseUnkownTableWhilePushingToViewShouldBeNotifyMVNow(t *testing.T
 		Message: "Table abc does not exist. Maybe you meant abc2?: while executing 'FUNCTION func()': while pushing to view some_mv (some-uuid-here)",
 	}
 	errorClass, errInfo := GetErrorClass(t.Context(),
-		exceptions.NewNormalizationError(fmt.Errorf("failed to normalize records: %w", err)))
+		exceptions.NewNormalizationError("failed to normalize records: %w", err))
 	assert.Equal(t, ErrorNotifyMVOrView, errorClass, "Unexpected error class")
 	assert.Equal(t, ErrorInfo{
 		Source: ErrorSourceClickHouse,

flow/alerting/classifier.go

@@ -12,7 +12,6 @@ import (
 	tEnums "go.temporal.io/api/enums/v1"
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/sdk/client"
-	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/durationpb"
 
 	"github.com/PeerDB-io/peerdb/flow/alerting"
@@ -146,6 +145,11 @@ func (h *FlowRequestHandler) CreateCDCFlow(
 		return nil, fmt.Errorf("unable to update flow config in catalog: %w", err)
 	}
 
+	if err := h.createGranularStatusEntry(ctx, cfg); err != nil {
+		slog.Error("unable to create granular status entry", slog.Any("error", err))
+		return nil, fmt.Errorf("unable to create granular status entry: %w", err)
+	}
+
 	if _, err := h.temporalClient.ExecuteWorkflow(ctx, workflowOptions, peerflow.CDCFlowWorkflow, cfg, nil); err != nil {
 		slog.Error("unable to start PeerFlow workflow", slog.Any("error", err))
 		return nil, fmt.Errorf("unable to start PeerFlow workflow: %w", err)
@@ -163,6 +167,20 @@ func (h *FlowRequestHandler) updateFlowConfigInCatalog(
 	return internal.UpdateCDCConfigInCatalog(ctx, h.pool, slog.Default(), cfg)
 }
 
+func (h *FlowRequestHandler) createGranularStatusEntry(
+	ctx context.Context,
+	cfg *protos.FlowConnectionConfigs,
+) error {
+	if _, err := h.pool.Exec(ctx, `
+		INSERT INTO peerdb_stats.granular_status (flow_name, snapshot_succeeding,
+		sync_succeeding, normalize_succeeding, slot_lag_low) VALUES ($1, true, true, true, true)
+		`, cfg.FlowJobName,
+	); err != nil {
+		return fmt.Errorf("unable to insert into granular status table for flow %s", cfg.FlowJobName)
+	}
+	return nil
+}
+
 func (h *FlowRequestHandler) CreateQRepFlow(
 	ctx context.Context, req *protos.CreateQRepFlowRequest,
 ) (*protos.CreateQRepFlowResponse, error) {
@@ -220,7 +238,7 @@ func (h *FlowRequestHandler) updateQRepConfigInCatalog(
 	ctx context.Context,
 	cfg *protos.QRepConfig,
 ) error {
-	cfgBytes, err := proto.Marshal(cfg)
+	cfgBytes, err := internal.ProtoMarshal(cfg)
 	if err != nil {
 		return fmt.Errorf("unable to marshal qrep config: %w", err)
 	}