Support upsert mode pg -> pg

flow/connectors/postgres/qrep.go

@@ -524,7 +524,8 @@ func (c *PostgresConnector) SyncQRepRecords(
 	switch syncMode {
 	case protos.QRepSyncMode_QREP_SYNC_MODE_MULTI_INSERT:
 		stagingTableSync := &QRepStagingTableSync{connector: c}
-		return stagingTableSync.SyncQRepRecords(config.FlowJobName, dstTable, partition, stream)
+		return stagingTableSync.SyncQRepRecords(
+			config.FlowJobName, dstTable, partition, stream, config.WriteMode)
 	case protos.QRepSyncMode_QREP_SYNC_MODE_STORAGE_AVRO:
 		return 0, fmt.Errorf("[postgres] SyncQRepRecords not implemented for storage avro sync mode")
 	default:

flow/connectors/postgres/qrep_sync_method.go

@@ -3,11 +3,13 @@ package connpostgres
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/PeerDB-io/peer-flow/connectors/utils/metrics"
 	"github.com/PeerDB-io/peer-flow/generated/protos"
 	"github.com/PeerDB-io/peer-flow/model"
+	util "github.com/PeerDB-io/peer-flow/utils"
 	"github.com/jackc/pgx/v5"
 	log "github.com/sirupsen/logrus"
 	"google.golang.org/protobuf/encoding/protojson"
@@ -31,6 +33,7 @@ func (s *QRepStagingTableSync) SyncQRepRecords(
 	dstTableName *SchemaTable,
 	partition *protos.QRepPartition,
 	stream *model.QRecordStream,
+	writeMode *protos.QRepWriteMode,
 ) (int, error) {
 	partitionID := partition.PartitionId
 	startTime := time.Now()
@@ -66,19 +69,100 @@ func (s *QRepStagingTableSync) SyncQRepRecords(
 	// Step 2: Insert records into the destination table.
 	copySource := model.NewQRecordBatchCopyFromSource(stream)
 
-	// Perform the COPY FROM operation
 	syncRecordsStartTime := time.Now()
-	syncedRows, err := tx.CopyFrom(
-		context.Background(),
-		pgx.Identifier{dstTableName.Schema, dstTableName.Table},
-		schema.GetColumnNames(),
-		copySource,
-	)
+	var numRowsSynced int64
 
-	if err != nil {
-		return -1, fmt.Errorf("failed to copy records into destination table: %v", err)
+	if writeMode == nil ||
+		writeMode.WriteType == protos.QRepWriteType_QREP_WRITE_MODE_APPEND {
+		// Perform the COPY FROM operation
+		numRowsSynced, err = tx.CopyFrom(
+			context.Background(),
+			pgx.Identifier{dstTableName.Schema, dstTableName.Table},
+			schema.GetColumnNames(),
+			copySource,
+		)
+		if err != nil {
+			return -1, fmt.Errorf("failed to copy records into destination table: %v", err)
+		}
+	} else {
+		// Step 2.1: Create a temp staging table
+		stagingTableName := fmt.Sprintf("_peerdb_staging_%s", util.RandomString(8))
+		stagingTableIdentifier := pgx.Identifier{dstTableName.Schema, stagingTableName}
+		dstTableIdentifier := pgx.Identifier{dstTableName.Schema, dstTableName.Table}
+
+		createStagingTableStmt := fmt.Sprintf(
+			"CREATE UNLOGGED TABLE %s (LIKE %s);",
+			stagingTableIdentifier.Sanitize(),
+			dstTableIdentifier.Sanitize(),
+		)
+
+		log.Infof("Creating staging table %s - '%s'", stagingTableName, createStagingTableStmt)
+		_, err = tx.Exec(context.Background(), createStagingTableStmt)
+
+		if err != nil {
+			return -1, fmt.Errorf("failed to create staging table: %v", err)
+		}
+
+		// Step 2.2: Insert records into the staging table
+		numRowsSynced, err = tx.CopyFrom(
+			context.Background(),
+			stagingTableIdentifier,
+			schema.GetColumnNames(),
+			copySource,
+		)
+		if err != nil {
+			return -1, fmt.Errorf("failed to copy records into staging table: %v", err)
+		}
+
+		// construct the SET clause for the upsert operation
+		upsertMatchColsList := writeMode.UpsertKeyColumns
+		upsertMatchCols := make(map[string]bool)
+		for _, col := range upsertMatchColsList {
+			upsertMatchCols[col] = true
+		}
+
+		setClause := ""
+		for _, col := range schema.GetColumnNames() {
+			_, ok := upsertMatchCols[col]
+			if !ok {
+				setClause += fmt.Sprintf("%s = EXCLUDED.%s,", col, col)
+			}
+		}
+
+		setClause = strings.TrimSuffix(setClause, ",")
+		selectStr := strings.Join(schema.GetColumnNames(), ", ")
+
+		// Step 2.3: Perform the upsert operation, ON CONFLICT UPDATE
+		upsertStmt := fmt.Sprintf(
+			"INSERT INTO %s (%s) SELECT %s FROM %s ON CONFLICT (%s) DO UPDATE SET %s;",
+			dstTableIdentifier.Sanitize(),
+			selectStr,
+			selectStr,
+			stagingTableIdentifier.Sanitize(),
+			strings.Join(writeMode.UpsertKeyColumns, ", "),
+			setClause,
+		)
+		log.Infof("Performing upsert operation: %s", upsertStmt)
+		res, err := tx.Exec(context.Background(), upsertStmt)
+		if err != nil {
+			return -1, fmt.Errorf("failed to perform upsert operation: %v", err)
+		}
+
+		numRowsSynced = res.RowsAffected()
+
+		// Step 2.4: Drop the staging table
+		dropStagingTableStmt := fmt.Sprintf(
+			"DROP TABLE %s;",
+			stagingTableIdentifier.Sanitize(),
+		)
+		log.Infof("Dropping staging table %s", stagingTableName)
+		_, err = tx.Exec(context.Background(), dropStagingTableStmt)
+		if err != nil {
+			return -1, fmt.Errorf("failed to drop staging table: %v", err)
+		}
 	}
-	metrics.LogQRepSyncMetrics(s.connector.ctx, flowJobName, syncedRows, time.Since(syncRecordsStartTime))
+
+	metrics.LogQRepSyncMetrics(s.connector.ctx, flowJobName, numRowsSynced, time.Since(syncRecordsStartTime))
 
 	// marshal the partition to json using protojson
 	pbytes, err := protojson.Marshal(partition)