Merge branch 'main' into normalize-split

.github/workflows/ui-lint.yml

@@ -43,3 +43,4 @@ jobs:
           prettier_dir: ui
           eslint_args: "--max-warnings 0"
           eslint_extensions: js,ts,jsx,tsx
+          prettier_extensions: js,ts,jsx,tsx,json

flow/activities/flowable.go

@@ -217,7 +217,7 @@ func (a *FlowableActivity) StartFlow(ctx context.Context,
 
 	go a.recordSlotSizePeriodically(errCtx, srcConn, slotNameForMetrics, input.FlowConnectionConfigs.Source.Name)
 
-	shutdown := utils.HeartbeatRoutine(ctx, 10*time.Second, func() string {
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		jobName := input.FlowConnectionConfigs.FlowJobName
 		return fmt.Sprintf("transferring records for job - %s", jobName)
 	})
@@ -387,7 +387,7 @@ func (a *FlowableActivity) StartNormalize(
 	}
 	defer connectors.CloseConnector(dstConn)
 
-	shutdown := utils.HeartbeatRoutine(ctx, 15*time.Second, func() string {
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		return fmt.Sprintf("normalizing records from batch for job - %s", input.FlowConnectionConfigs.FlowJobName)
 	})
 	defer shutdown()
@@ -455,7 +455,7 @@ func (a *FlowableActivity) GetQRepPartitions(ctx context.Context,
 	}
 	defer connectors.CloseConnector(srcConn)
 
-	shutdown := utils.HeartbeatRoutine(ctx, 2*time.Minute, func() string {
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		return fmt.Sprintf("getting partitions for job - %s", config.FlowJobName)
 	})
 	defer shutdown()
@@ -544,6 +544,10 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 	defer connectors.CloseConnector(dstConn)
 
 	slog.InfoContext(ctx, fmt.Sprintf("replicating partition %s\n", partition.PartitionId))
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
+		return fmt.Sprintf("syncing partition - %s: %d of %d total.", partition.PartitionId, idx, total)
+	})
+	defer shutdown()
 
 	var stream *model.QRecordStream
 	bufferSize := shared.FetchAndChannelSize
@@ -593,11 +597,6 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 		}
 	}
 
-	shutdown := utils.HeartbeatRoutine(ctx, 1*time.Minute, func() string {
-		return fmt.Sprintf("syncing partition - %s: %d of %d total.", partition.PartitionId, idx, total)
-	})
-	defer shutdown()
-
 	rowsSynced, err := dstConn.SyncQRepRecords(config, partition, stream)
 	if err != nil {
 		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
@@ -635,10 +634,9 @@ func (a *FlowableActivity) ConsolidateQRepPartitions(ctx context.Context, config
 	} else if err != nil {
 		return err
 	}
-
 	defer connectors.CloseConnector(dstConn)
 
-	shutdown := utils.HeartbeatRoutine(ctx, 2*time.Minute, func() string {
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		return fmt.Sprintf("consolidating partitions for job - %s", config.FlowJobName)
 	})
 	defer shutdown()
@@ -802,7 +800,21 @@ func (a *FlowableActivity) QRepWaitUntilNewRows(ctx context.Context,
 	attemptCount := 1
 	for {
 		activity.RecordHeartbeat(ctx, fmt.Sprintf("no new rows yet, attempt #%d", attemptCount))
-		time.Sleep(waitBetweenBatches)
+		waitUntil := time.Now().Add(waitBetweenBatches)
+		for {
+			sleep := time.Until(waitUntil)
+			if sleep > 15*time.Second {
+				sleep = 15 * time.Second
+			}
+			time.Sleep(sleep)
+
+			activity.RecordHeartbeat(ctx, "heartbeat while waiting before next batch")
+			if err := ctx.Err(); err != nil {
+				return fmt.Errorf("cancelled while waiting for new rows: %w", err)
+			} else if time.Now().After(waitUntil) {
+				break
+			}
+		}
 
 		result, err := pgSrcConn.CheckForUpdatedMaxValue(config, last)
 		if err != nil {
@@ -830,6 +842,11 @@ func (a *FlowableActivity) RenameTables(ctx context.Context, config *protos.Rena
 	}
 	defer connectors.CloseConnector(dstConn)
 
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
+		return fmt.Sprintf("renaming tables for job - %s", config.FlowJobName)
+	})
+	defer shutdown()
+
 	if config.Peer.Type == protos.DBType_SNOWFLAKE {
 		sfConn, ok := dstConn.(*connsnowflake.SnowflakeConnector)
 		if !ok {
@@ -949,7 +966,7 @@ func (a *FlowableActivity) ReplicateXminPartition(ctx context.Context,
 		return nil
 	})
 
-	shutdown := utils.HeartbeatRoutine(ctx, 5*time.Minute, func() string {
+	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		return "syncing xmin."
 	})
 	defer shutdown()

flow/cmd/validate_peer.go

@@ -41,13 +41,13 @@ func (h *FlowRequestHandler) ValidatePeer(
 	defer conn.Close()
 
 	if req.Peer.Type == protos.DBType_POSTGRES {
-		version, err := conn.(*connpostgres.PostgresConnector).GetPostgresVersion()
+		isValid, version, err := conn.(*connpostgres.PostgresConnector).MajorVersionCheck(connpostgres.POSTGRES_12)
 		if err != nil {
 			slog.Error("/peer/validate: pg version check", slog.Any("error", err))
 			return nil, err
 		}
 
-		if version < 12 {
+		if !isValid {
 			return &protos.ValidatePeerResponse{
 				Status: protos.ValidatePeerStatus_INVALID,
 				Message: fmt.Sprintf("%s peer %s must be of version 12 or above. Current version: %d",

flow/connectors/bigquery/bigquery.go

@@ -828,6 +828,7 @@ func (c *BigQueryConnector) SetupNormalizedTables(
 		datasetTablesSet[*datasetTable] = struct{}{}
 		// log that table was created
 		c.logger.Info(fmt.Sprintf("created table %s", tableIdentifier))
+		utils.RecordHeartbeatWithRecover(c.ctx, fmt.Sprintf("created table %s", tableIdentifier))
 	}
 
 	return &protos.SetupNormalizedTableBatchOutput{

flow/connectors/bigquery/qrep_avro_sync.go

@@ -42,7 +42,7 @@ func (s *QRepAvroSyncMethod) SyncRecords(
 	stream *model.QRecordStream,
 	tableNameRowsMapping map[string]uint32,
 ) (*model.SyncResponse, error) {
-	activity.RecordHeartbeat(s.connector.ctx, time.Minute,
+	activity.RecordHeartbeat(s.connector.ctx,
 		fmt.Sprintf("Flow job %s: Obtaining Avro schema"+
 			" for destination table %s and sync batch ID %d",
 			req.FlowJobName, rawTableName, syncBatchID),
@@ -77,7 +77,7 @@ func (s *QRepAvroSyncMethod) SyncRecords(
 		return nil, fmt.Errorf("failed to update metadata: %v", err)
 	}
 
-	activity.RecordHeartbeat(s.connector.ctx, time.Minute,
+	activity.RecordHeartbeat(s.connector.ctx,
 		fmt.Sprintf("Flow job %s: performing insert and update transaction"+
 			" for destination table %s and sync batch ID %d",
 			req.FlowJobName, rawTableName, syncBatchID),
@@ -392,12 +392,10 @@ func (s *QRepAvroSyncMethod) writeToStage(
 	stream *model.QRecordStream,
 	flowName string,
 ) (int, error) {
-	shutdown := utils.HeartbeatRoutine(s.connector.ctx, time.Minute,
-		func() string {
-			return fmt.Sprintf("writing to avro stage for objectFolder %s and staging table %s",
-				objectFolder, stagingTable)
-		},
-	)
+	shutdown := utils.HeartbeatRoutine(s.connector.ctx, func() string {
+		return fmt.Sprintf("writing to avro stage for objectFolder %s and staging table %s",
+			objectFolder, stagingTable)
+	})
 	defer shutdown()
 
 	var avroFile *avro.AvroFile

flow/connectors/eventhub/eventhub.go

@@ -132,11 +132,8 @@ func (c *EventHubConnector) processBatch(
 	lastUpdatedOffset := int64(0)
 
 	numRecords := atomic.Uint32{}
-	shutdown := utils.HeartbeatRoutine(c.ctx, 10*time.Second, func() string {
-		return fmt.Sprintf(
-			"processed %d records for flow %s",
-			numRecords.Load(), flowJobName,
-		)
+	shutdown := utils.HeartbeatRoutine(c.ctx, func() string {
+		return fmt.Sprintf("processed %d records for flow %s", numRecords.Load(), flowJobName)
 	})
 	defer shutdown()
 

flow/connectors/postgres/cdc.go

@@ -244,7 +244,7 @@ func (p *PostgresCDCSource) consumeStream(
 		}
 	}()
 
-	shutdown := utils.HeartbeatRoutine(p.ctx, 10*time.Second, func() string {
+	shutdown := utils.HeartbeatRoutine(p.ctx, func() string {
 		jobName := p.flowJobName
 		currRecords := cdcRecordsStorage.Len()
 		return fmt.Sprintf("pulling records for job - %s, currently have %d records", jobName, currRecords)

flow/connectors/postgres/client.go

@@ -17,6 +17,14 @@ import (
 	"github.com/lib/pq/oid"
 )
 
+type PGVersion int
+
+const (
+	POSTGRES_12 PGVersion = 120000
+	POSTGRES_13 PGVersion = 130000
+	POSTGRES_15 PGVersion = 150000
+)
+
 const (
 	mirrorJobsTableIdentifier = "peerdb_mirror_jobs"
 	createMirrorJobsTableSQL  = `CREATE TABLE IF NOT EXISTS %s.%s(mirror_job_name TEXT PRIMARY KEY,
@@ -305,7 +313,7 @@ func (c *PostgresConnector) createSlotAndPublication(
 
 	if !s.PublicationExists {
 		// check and enable publish_via_partition_root
-		supportsPubViaRoot, err := c.majorVersionCheck(130000)
+		supportsPubViaRoot, _, err := c.MajorVersionCheck(POSTGRES_13)
 		if err != nil {
 			return fmt.Errorf("error checking Postgres version: %w", err)
 		}
@@ -472,14 +480,14 @@ func (c *PostgresConnector) jobMetadataExists(jobName string) (bool, error) {
 	return result.Bool, nil
 }
 
-func (c *PostgresConnector) majorVersionCheck(majorVersion int) (bool, error) {
+func (c *PostgresConnector) MajorVersionCheck(majorVersion PGVersion) (bool, int64, error) {
 	var version pgtype.Int8
 	err := c.pool.QueryRow(c.ctx, "SELECT current_setting('server_version_num')::INTEGER").Scan(&version)
 	if err != nil {
-		return false, fmt.Errorf("failed to get server version: %w", err)
+		return false, 0, fmt.Errorf("failed to get server version: %w", err)
 	}
 
-	return int(version.Int64) >= majorVersion, nil
+	return version.Int64 >= int64(majorVersion), version.Int64, nil
 }
 
 func (c *PostgresConnector) updateSyncMetadata(flowJobName string, lastCP int64, syncBatchID int64,
@@ -621,7 +629,12 @@ func (c *PostgresConnector) CheckReplicationPermissions(username string) error {
 	}
 
 	if !replicationRes {
-		return fmt.Errorf("postgres user does not have replication role")
+		// RDS case: check pg_settings for rds.logical_replication
+		var setting string
+		err := c.pool.QueryRow(c.ctx, "SELECT setting FROM pg_settings WHERE name = 'rds.logical_replication';").Scan(&setting)
+		if err != nil || setting != "on" {
+			return fmt.Errorf("postgres user does not have replication role")
+		}
 	}
 
 	// check wal_level
@@ -653,22 +666,3 @@ func (c *PostgresConnector) CheckReplicationPermissions(username string) error {
 
 	return nil
 }
-
-func (c *PostgresConnector) GetPostgresVersion() (int, error) {
-	if c.pool == nil {
-		return -1, fmt.Errorf("version check: pool is nil")
-	}
-
-	var versionRes string
-	err := c.pool.QueryRow(c.ctx, "SHOW server_version_num;").Scan(&versionRes)
-	if err != nil {
-		return -1, err
-	}
-
-	version, err := strconv.Atoi(versionRes)
-	if err != nil {
-		return -1, err
-	}
-
-	return version / 10000, nil
-}

flow/connectors/postgres/postgres.go

@@ -476,7 +476,7 @@ func (c *PostgresConnector) NormalizeRecords(req *model.NormalizeRecordsRequest)
 		}
 	}()
 
-	supportsMerge, err := c.majorVersionCheck(150000)
+	supportsMerge, _, err := c.MajorVersionCheck(POSTGRES_15)
 	if err != nil {
 		return nil, err
 	}

flow/connectors/postgres/qrep_query_executor.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
-	"time"
 
 	"github.com/PeerDB-io/peer-flow/connectors/utils"
 	"github.com/PeerDB-io/peer-flow/geo"
@@ -84,7 +83,7 @@ func (qe *QRepQueryExecutor) executeQueryInTx(tx pgx.Tx, cursorName string, fetc
 	q := fmt.Sprintf("FETCH %d FROM %s", fetchSize, cursorName)
 
 	if !qe.testEnv {
-		shutdown := utils.HeartbeatRoutine(qe.ctx, 1*time.Minute, func() string {
+		shutdown := utils.HeartbeatRoutine(qe.ctx, func() string {
 			qe.logger.Info(fmt.Sprintf("still running '%s'...", q))
 			return fmt.Sprintf("running '%s'", q)
 		})

flow/connectors/snowflake/qrep_avro_sync.go

@@ -270,7 +270,7 @@ func (s *SnowflakeAvroSyncHandler) putFileToStage(avroFile *avro.AvroFile, stage
 	activity.RecordHeartbeat(s.connector.ctx, "putting file to stage")
 	putCmd := fmt.Sprintf("PUT file://%s @%s", avroFile.FilePath, stage)
 
-	shutdown := utils.HeartbeatRoutine(s.connector.ctx, 10*time.Second, func() string {
+	shutdown := utils.HeartbeatRoutine(s.connector.ctx, func() string {
 		return fmt.Sprintf("putting file to stage %s", stage)
 	})
 	defer shutdown()

flow/connectors/snowflake/snowflake.go

@@ -414,6 +414,7 @@ func (c *SnowflakeConnector) SetupNormalizedTables(
 			return nil, fmt.Errorf("[sf] error while creating normalized table: %w", err)
 		}
 		tableExistsMapping[tableIdentifier] = false
+		utils.RecordHeartbeatWithRecover(c.ctx, fmt.Sprintf("created table %s", tableIdentifier))
 	}
 
 	return &protos.SetupNormalizedTableBatchOutput{

flow/connectors/utils/avro/avro_writer.go

@@ -8,7 +8,6 @@ import (
 	"log/slog"
 	"os"
 	"sync/atomic"
-	"time"
 
 	"github.com/PeerDB-io/peer-flow/connectors/utils"
 	"github.com/PeerDB-io/peer-flow/model"
@@ -131,7 +130,7 @@ func (p *peerDBOCFWriter) writeRecordsToOCFWriter(ocfWriter *goavro.OCFWriter) (
 	numRows := atomic.Uint32{}
 
 	if p.ctx != nil {
-		shutdown := utils.HeartbeatRoutine(p.ctx, 30*time.Second, func() string {
+		shutdown := utils.HeartbeatRoutine(p.ctx, func() string {
 			written := numRows.Load()
 			return fmt.Sprintf("[avro] written %d rows to OCF", written)
 		})

flow/connectors/utils/heartbeat.go

@@ -11,7 +11,6 @@ import (
 
 func HeartbeatRoutine(
 	ctx context.Context,
-	interval time.Duration,
 	message func() string,
 ) func() {
 	shutdown := make(chan struct{})
@@ -26,7 +25,7 @@ func HeartbeatRoutine(
 				return
 			case <-ctx.Done():
 				return
-			case <-time.After(interval):
+			case <-time.After(15 * time.Second):
 			}
 		}
 	}()

flow/workflows/cdc_flow.go

@@ -345,7 +345,7 @@ func CDCFlowWorkflowWithConfig(
 			state.TableNameSchemaMapping = correctedTableNameSchemaMapping
 			renameTablesCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
 				StartToCloseTimeout: 12 * time.Hour,
-				HeartbeatTimeout:    1 * time.Hour,
+				HeartbeatTimeout:    time.Minute,
 			})
 			renameTablesFuture := workflow.ExecuteActivity(renameTablesCtx, flowable.RenameTables, renameOpts)
 			if err := renameTablesFuture.Get(renameTablesCtx, nil); err != nil {