Merge branch 'main' into schema-changes-diff-from-schema

.github/workflows/flow.yml

@@ -12,7 +12,7 @@ jobs:
       matrix:
         runner: [ubicloud-standard-16-ubuntu-2204-arm]
     runs-on: ${{ matrix.runner }}
-    timeout-minutes: 40
+    timeout-minutes: 30
     services:
       catalog:
         image: imresamu/postgis:15-3.4-alpine
@@ -96,7 +96,7 @@ jobs:
           temporal operator search-attribute create --name MirrorName --type Text --namespace default
           ./peer-flow worker &
           ./peer-flow snapshot-worker &
-          go test -p 32 ./... -timeout 1200s
+          go test -p 32 ./... -timeout 900s
         working-directory: ./flow
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}

flow/activities/flowable.go

@@ -301,7 +301,7 @@ func (a *FlowableActivity) SyncFlow(
 	}
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
-		return "transferring records for job - " + flowName
+		return "transferring records for job"
 	})
 	defer shutdown()
 
@@ -467,7 +467,7 @@ func (a *FlowableActivity) StartNormalize(
 	defer connectors.CloseConnector(ctx, dstConn)
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
-		return "normalizing records from batch for job - " + input.FlowConnectionConfigs.FlowJobName
+		return "normalizing records from batch for job"
 	})
 	defer shutdown()
 
@@ -535,7 +535,7 @@ func (a *FlowableActivity) GetQRepPartitions(ctx context.Context,
 	defer connectors.CloseConnector(ctx, srcConn)
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
-		return "getting partitions for job - " + config.FlowJobName
+		return "getting partitions for job"
 	})
 	defer shutdown()
 
@@ -604,12 +604,6 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 	ctx = context.WithValue(ctx, shared.FlowNameKey, config.FlowJobName)
 	logger := activity.GetLogger(ctx)
 
-	err := monitoring.UpdateStartTimeForPartition(ctx, a.CatalogPool, runUUID, partition, time.Now())
-	if err != nil {
-		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
-		return fmt.Errorf("failed to update start time for partition: %w", err)
-	}
-
 	srcConn, err := connectors.GetQRepPullConnector(ctx, config.SourcePeer)
 	if err != nil {
 		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
@@ -624,6 +618,25 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 	}
 	defer connectors.CloseConnector(ctx, dstConn)
 
+	done, err := dstConn.IsQRepPartitionSynced(ctx, &protos.IsQRepPartitionSyncedInput{
+		FlowJobName: config.FlowJobName,
+		PartitionId: partition.PartitionId,
+	})
+	if err != nil {
+		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
+		return fmt.Errorf("failed to get fetch status of partition: %w", err)
+	}
+	if done {
+		logger.Info("no records to push for partition " + partition.PartitionId)
+		return nil
+	}
+
+	err = monitoring.UpdateStartTimeForPartition(ctx, a.CatalogPool, runUUID, partition, time.Now())
+	if err != nil {
+		a.Alerter.LogFlowError(ctx, config.FlowJobName, err)
+		return fmt.Errorf("failed to update start time for partition: %w", err)
+	}
+
 	logger.Info("replicating partition " + partition.PartitionId)
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
 		return fmt.Sprintf("syncing partition - %s: %d of %d total.", partition.PartitionId, idx, total)
@@ -698,8 +711,6 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 		if err != nil {
 			return err
 		}
-	} else {
-		logger.Info("no records to push for partition " + partition.PartitionId)
 	}
 
 	err = monitoring.UpdateEndTimeForPartition(ctx, a.CatalogPool, runUUID, partition)
@@ -718,7 +729,7 @@ func (a *FlowableActivity) ConsolidateQRepPartitions(ctx context.Context, config
 	defer connectors.CloseConnector(ctx, dstConn)
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
-		return "consolidating partitions for job - " + config.FlowJobName
+		return "consolidating partitions for job"
 	})
 	defer shutdown()
 
@@ -973,7 +984,7 @@ func (a *FlowableActivity) RenameTables(ctx context.Context, config *protos.Rena
 	defer connectors.CloseConnector(ctx, dstConn)
 
 	shutdown := utils.HeartbeatRoutine(ctx, func() string {
-		return "renaming tables for job - " + config.FlowJobName
+		return "renaming tables for job"
 	})
 	defer shutdown()
 

flow/activities/snapshot_activity.go

@@ -32,7 +32,7 @@ func (a *SnapshotActivity) CloseSlotKeepAlive(ctx context.Context, flowJobName s
 		connectors.CloseConnector(ctx, s.connector)
 		delete(a.SnapshotConnections, flowJobName)
 	}
-	a.Alerter.LogFlowEvent(ctx, flowJobName, "Ended Snapshot Flow Job - "+flowJobName)
+	a.Alerter.LogFlowEvent(ctx, flowJobName, "Ended Snapshot Flow Job")
 
 	return nil
 }
@@ -50,7 +50,7 @@ func (a *SnapshotActivity) SetupReplication(
 		return nil, nil
 	}
 
-	a.Alerter.LogFlowEvent(ctx, config.FlowJobName, "Started Snapshot Flow Job - "+config.FlowJobName)
+	a.Alerter.LogFlowEvent(ctx, config.FlowJobName, "Started Snapshot Flow Job")
 
 	conn, err := connectors.GetCDCPullConnector(ctx, config.PeerConnectionConfig)
 	if err != nil {
@@ -84,7 +84,7 @@ func (a *SnapshotActivity) SetupReplication(
 	var slotInfo connpostgres.SlotCreationResult
 	select {
 	case slotInfo = <-slotSignal.SlotCreated:
-		logger.Info("slot created", slotInfo.SlotName)
+		logger.Info("slot created", slog.String("SlotName", slotInfo.SlotName))
 	case err := <-replicationErr:
 		closeConnectionForError(err)
 		return nil, fmt.Errorf("failed to setup replication: %w", err)

flow/cmd/mirror_status.go

@@ -144,7 +144,8 @@ func (h *FlowRequestHandler) cloneTableSummary(
 
 	rows, err := h.pool.Query(ctx, q, "clone_"+flowJobName+"_%")
 	if err != nil {
-		slog.Error("unable to query initial load partition - "+flowJobName, slog.Any("error", err))
+		slog.Error("unable to query initial load partition",
+			slog.String(string(shared.FlowNameKey), flowJobName), slog.Any("error", err))
 		return nil, fmt.Errorf("unable to query initial load partition - %s: %w", flowJobName, err)
 	}
 

flow/connectors/bigquery/qrep.go

@@ -30,15 +30,6 @@ func (c *BigQueryConnector) SyncQRepRecords(
 		return 0, err
 	}
 
-	done, err := c.pgMetadata.IsQrepPartitionSynced(ctx, config.FlowJobName, partition.PartitionId)
-	if err != nil {
-		return 0, fmt.Errorf("failed to check if partition %s is synced: %w", partition.PartitionId, err)
-	}
-
-	if done {
-		c.logger.Info(fmt.Sprintf("Partition %s has already been synced", partition.PartitionId))
-		return 0, nil
-	}
 	c.logger.Info(fmt.Sprintf("QRep sync function called and partition existence checked for"+
 		" partition %s of destination table %s",
 		partition.PartitionId, destTable))
@@ -111,3 +102,9 @@ func (c *BigQueryConnector) SetupQRepMetadataTables(ctx context.Context, config
 
 	return nil
 }
+
+func (c *BigQueryConnector) IsQRepPartitionSynced(ctx context.Context,
+	req *protos.IsQRepPartitionSyncedInput,
+) (bool, error) {
+	return c.pgMetadata.IsQrepPartitionSynced(ctx, req.FlowJobName, req.PartitionId)
+}

flow/connectors/clickhouse/qrep.go

@@ -33,16 +33,6 @@ func (c *ClickhouseConnector) SyncQRepRecords(
 		slog.String("destinationTable", destTable),
 	)
 
-	done, err := c.isPartitionSynced(ctx, partition.PartitionId)
-	if err != nil {
-		return 0, fmt.Errorf("failed to check if partition %s is synced: %w", partition.PartitionId, err)
-	}
-
-	if done {
-		c.logger.Info("Partition has already been synced", flowLog)
-		return 0, nil
-	}
-
 	tblSchema, err := c.getTableSchema(destTable)
 	if err != nil {
 		return 0, fmt.Errorf("failed to get schema of table %s: %w", destTable, err)
@@ -96,9 +86,11 @@ func (c *ClickhouseConnector) getTableSchema(tableName string) ([]*sql.ColumnTyp
 	return columnTypes, nil
 }
 
-func (c *ClickhouseConnector) isPartitionSynced(ctx context.Context, partitionID string) (bool, error) {
+func (c *ClickhouseConnector) IsQRepPartitionSynced(ctx context.Context,
+	req *protos.IsQRepPartitionSyncedInput,
+) (bool, error) {
 	//nolint:gosec
-	queryString := fmt.Sprintf(`SELECT COUNT(*) FROM %s WHERE partitionID = '%s'`, qRepMetadataTableName, partitionID)
+	queryString := fmt.Sprintf(`SELECT COUNT(*) FROM %s WHERE partitionID = '%s'`, qRepMetadataTableName, req.PartitionId)
 
 	row := c.database.QueryRowContext(ctx, queryString)
 

flow/connectors/clickhouse/qrep_avro_sync.go

@@ -130,12 +130,12 @@ func (s *ClickhouseAvroSyncMethod) SyncQRepRecords(
 			continue
 		}
 
-		selector = append(selector, colName)
+		selector = append(selector, "`"+colName+"`")
 	}
 	selectorStr := strings.Join(selector, ",")
 	//nolint:gosec
-	query := fmt.Sprintf("INSERT INTO %s(%s) SELECT * FROM s3('%s','%s','%s', 'Avro')",
-		config.DestinationTableIdentifier, selectorStr, avroFileUrl,
+	query := fmt.Sprintf("INSERT INTO %s(%s) SELECT %s FROM s3('%s','%s','%s', 'Avro')",
+		config.DestinationTableIdentifier, selectorStr, selectorStr, avroFileUrl,
 		s.connector.creds.AccessKeyID, s.connector.creds.SecretAccessKey)
 
 	_, err = s.connector.database.ExecContext(ctx, query)

flow/connectors/core.go

@@ -146,6 +146,9 @@ type QRepPullConnector interface {
 type QRepSyncConnector interface {
 	Connector
 
+	// IsQRepPartitionSynced returns true if a partition has already been synced
+	IsQRepPartitionSynced(ctx context.Context, req *protos.IsQRepPartitionSyncedInput) (bool, error)
+
 	// SetupQRepMetadataTables sets up the metadata tables for QRep.
 	SetupQRepMetadataTables(ctx context.Context, config *protos.QRepConfig) error
 
@@ -261,6 +264,7 @@ var (
 	_ QRepSyncConnector = &connbigquery.BigQueryConnector{}
 	_ QRepSyncConnector = &connsnowflake.SnowflakeConnector{}
 	_ QRepSyncConnector = &connclickhouse.ClickhouseConnector{}
+	_ QRepSyncConnector = &conns3.S3Connector{}
 
 	_ QRepConsolidateConnector = &connsnowflake.SnowflakeConnector{}
 	_ QRepConsolidateConnector = &connclickhouse.ClickhouseConnector{}

flow/connectors/postgres/postgres_schema_delta_test.go

@@ -196,23 +196,25 @@ func (s PostgresSchemaDeltaTestSuite) TestAddDropWhitespaceColumnNames() {
 }
 
 func TestPostgresSchemaDeltaTestSuite(t *testing.T) {
-	e2eshared.RunSuite(t, SetupSuite, func(s PostgresSchemaDeltaTestSuite) {
-		teardownTx, err := s.connector.conn.Begin(context.Background())
-		require.NoError(s.t, err)
-		defer func() {
-			err := teardownTx.Rollback(context.Background())
-			if err != pgx.ErrTxClosed {
-				require.NoError(s.t, err)
-			}
-		}()
-		_, err = teardownTx.Exec(context.Background(), fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE",
-			s.schema))
-		require.NoError(s.t, err)
-		err = teardownTx.Commit(context.Background())
-		require.NoError(s.t, err)
-
-		require.NoError(s.t, s.connector.ConnectionActive(context.Background()))
-		require.NoError(s.t, s.connector.Close())
-		require.Error(s.t, s.connector.ConnectionActive(context.Background()))
-	})
+	e2eshared.RunSuite(t, SetupSuite)
+}
+
+func (s PostgresSchemaDeltaTestSuite) Teardown() {
+	teardownTx, err := s.connector.conn.Begin(context.Background())
+	require.NoError(s.t, err)
+	defer func() {
+		err := teardownTx.Rollback(context.Background())
+		if err != pgx.ErrTxClosed {
+			require.NoError(s.t, err)
+		}
+	}()
+	_, err = teardownTx.Exec(context.Background(), fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE",
+		s.schema))
+	require.NoError(s.t, err)
+	err = teardownTx.Commit(context.Background())
+	require.NoError(s.t, err)
+
+	require.NoError(s.t, s.connector.ConnectionActive(context.Background()))
+	require.NoError(s.t, s.connector.Close())
+	require.Error(s.t, s.connector.ConnectionActive(context.Background()))
 }

flow/connectors/postgres/qrep.go

@@ -458,15 +458,6 @@ func (c *PostgresConnector) SyncQRepRecords(
 		return 0, fmt.Errorf("table %s does not exist, used schema: %s", dstTable.Table, dstTable.Schema)
 	}
 
-	done, err := c.isPartitionSynced(ctx, partition.PartitionId)
-	if err != nil {
-		return 0, fmt.Errorf("failed to check if partition is synced: %w", err)
-	}
-
-	if done {
-		c.logger.Info(fmt.Sprintf("partition %s already synced", partition.PartitionId))
-		return 0, nil
-	}
 	c.logger.Info("SyncRecords called and initial checks complete.")
 
 	stagingTableSync := &QRepStagingTableSync{connector: c}
@@ -573,18 +564,20 @@ func BuildQuery(logger log.Logger, query string, flowJobName string) (string, er
 	return res, nil
 }
 
-// isPartitionSynced checks whether a specific partition is synced
-func (c *PostgresConnector) isPartitionSynced(ctx context.Context, partitionID string) (bool, error) {
+// IsQRepPartitionSynced checks whether a specific partition is synced
+func (c *PostgresConnector) IsQRepPartitionSynced(ctx context.Context,
+	req *protos.IsQRepPartitionSyncedInput,
+) (bool, error) {
 	// setup the query string
 	metadataTableIdentifier := pgx.Identifier{c.metadataSchema, qRepMetadataTableName}
 	queryString := fmt.Sprintf(
-		"SELECT COUNT(*)>0 FROM %s WHERE partitionID = $1;",
+		"SELECT COUNT(*)>0 FROM %s WHERE partitionID=$1;",
 		metadataTableIdentifier.Sanitize(),
 	)
 
 	// prepare and execute the query
 	var result bool
-	err := c.conn.QueryRow(ctx, queryString, partitionID).Scan(&result)
+	err := c.conn.QueryRow(ctx, queryString, req.PartitionId).Scan(&result)
 	if err != nil {
 		return false, fmt.Errorf("failed to execute query: %w", err)
 	}

flow/connectors/s3/qrep.go

@@ -81,3 +81,10 @@ func (c *S3Connector) SetupQRepMetadataTables(_ context.Context, config *protos.
 	c.logger.Info("QRep metadata setup not needed for S3.")
 	return nil
 }
+
+// S3 doesn't check if partition is already synced, but file with same name is overwritten
+func (c *S3Connector) IsQRepPartitionSynced(_ context.Context,
+	config *protos.IsQRepPartitionSyncedInput,
+) (bool, error) {
+	return false, nil
+}

flow/connectors/snowflake/qrep.go

@@ -35,16 +35,6 @@ func (c *SnowflakeConnector) SyncQRepRecords(
 	}
 	c.logger.Info("Called QRep sync function and obtained table schema", flowLog)
 
-	done, err := c.pgMetadata.IsQrepPartitionSynced(ctx, config.FlowJobName, partition.PartitionId)
-	if err != nil {
-		return 0, fmt.Errorf("failed to check if partition %s is synced: %w", partition.PartitionId, err)
-	}
-
-	if done {
-		c.logger.Info("Partition has already been synced", flowLog)
-		return 0, nil
-	}
-
 	avroSync := NewSnowflakeAvroSyncHandler(config, c)
 	return avroSync.SyncQRepRecords(ctx, config, partition, tblSchema, stream)
 }
@@ -282,3 +272,9 @@ func (c *SnowflakeConnector) dropStage(ctx context.Context, stagingPath string,
 func (c *SnowflakeConnector) getStageNameForJob(job string) string {
 	return fmt.Sprintf("%s.peerdb_stage_%s", c.rawSchema, job)
 }
+
+func (c *SnowflakeConnector) IsQRepPartitionSynced(ctx context.Context,
+	req *protos.IsQRepPartitionSyncedInput,
+) (bool, error) {
+	return c.pgMetadata.IsQrepPartitionSynced(ctx, req.FlowJobName, req.PartitionId)
+}