diff --git a/flow/cmd/mirror_status.go b/flow/cmd/mirror_status.go index 6ac815279..2635ee726 100644 --- a/flow/cmd/mirror_status.go +++ b/flow/cmd/mirror_status.go @@ -331,9 +331,8 @@ func (h *FlowRequestHandler) InitialLoadSummary( defer rows.Close() - cloneStatuses := []*protos.CloneTableSummary{} - for rows.Next() { - if err := rows.Scan( + cloneStatuses, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*protos.CloneTableSummary, error) { + if err := row.Scan( &flowName, &destinationTable, &sourceTable, @@ -345,7 +344,7 @@ func (h *FlowRequestHandler) InitialLoadSummary( &numRowsSynced, &avgTimePerPartitionMs, ); err != nil { - return nil, NewInternalApiError(fmt.Errorf("unable to scan initial load partition - %s: %w", parentMirrorName, err)) + return nil, err } res := &protos.CloneTableSummary{ @@ -392,7 +391,10 @@ func (h *FlowRequestHandler) InitialLoadSummary( res.AvgTimePerPartitionMs = int64(avgTimePerPartitionMs.Float64) } - cloneStatuses = append(cloneStatuses, res) + return res, nil + }) + if err != nil { + return nil, NewInternalApiError(fmt.Errorf("unable to scan initial load partition - %s: %w", parentMirrorName, err)) } return &protos.InitialLoadSummaryResponse{ TableSummaries: cloneStatuses, @@ -430,18 +432,15 @@ func (h *FlowRequestHandler) getPartitionStatuses( defer rows.Close() - res := []*protos.PartitionStatus{} var partitionId pgtype.Text var startTime pgtype.Timestamp var endTime pgtype.Timestamp var numRowsInPartition pgtype.Int8 var numRowsSynced pgtype.Int8 - for rows.Next() { - if err := rows.Scan(&partitionId, &startTime, &endTime, &numRowsInPartition, &numRowsSynced); err != nil { - slog.ErrorContext(ctx, "unable to scan qrep partition", - slog.String("flow", flowJobName), slog.Any("error", err)) - return nil, fmt.Errorf("unable to scan qrep partition - %s: %w", flowJobName, err) + res, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (*protos.PartitionStatus, error) { + if err := row.Scan(&partitionId, &startTime, &endTime, &numRowsInPartition, &numRowsSynced); err != nil { + return nil, err } partitionStatus := &protos.PartitionStatus{} @@ -462,7 +461,12 @@ func (h *FlowRequestHandler) getPartitionStatuses( partitionStatus.RowsSynced = numRowsSynced.Int64 } - res = append(res, partitionStatus) + return partitionStatus, nil + }) + if err != nil { + slog.ErrorContext(ctx, "unable to scan qrep partition", + slog.String("flow", flowJobName), slog.Any("error", err)) + return nil, fmt.Errorf("unable to scan qrep partition - %s: %w", flowJobName, err) } return res, nil diff --git a/flow/connectors/postgres/client.go b/flow/connectors/postgres/client.go index a06212d4f..d286d22bc 100644 --- a/flow/connectors/postgres/client.go +++ b/flow/connectors/postgres/client.go @@ -404,8 +404,7 @@ func getSlotInfo( return nil, fmt.Errorf("failed to read information for slots: %w", err) } defer rows.Close() - var slotInfoRows []*protos.SlotInfo - for rows.Next() { + return pgx.CollectRows(rows, func(row pgx.CollectableRow) (*protos.SlotInfo, error) { var slotName pgtype.Text var redoLSN pgtype.Text var restartLSN pgtype.Text @@ -427,7 +426,7 @@ func getSlotInfo( var spillCount *int64 var spillBytes *int64 - err := rows.Scan( + err := row.Scan( &slotName, &redoLSN, &restartLSN, @@ -453,7 +452,7 @@ func getSlotInfo( return nil, err } - slotInfoRows = append(slotInfoRows, &protos.SlotInfo{ + return &protos.SlotInfo{ SlotName: slotName.String, RedoLSN: redoLSN.String, RestartLSN: restartLSN.String, @@ -474,9 +473,8 @@ func getSlotInfo( SpillTxns: spillTxns, SpillCount: spillCount, SpillBytes: spillBytes, - }) - } - return slotInfoRows, nil + }, nil + }) } // GetSlotInfo gets the information about the replication slot size and LSNs. diff --git a/flow/connectors/postgres/postgres_source.go b/flow/connectors/postgres/postgres_source.go index 91b8d9f8e..c7c628852 100644 --- a/flow/connectors/postgres/postgres_source.go +++ b/flow/connectors/postgres/postgres_source.go @@ -433,13 +433,9 @@ func (c *PostgresConnector) GetSelectedColumns( return nil, fmt.Errorf("error getting selected columns for table %s: %w", sourceTable, err) } - columns := make([]string, 0) - for rows.Next() { - var columnName string - if err := rows.Scan(&columnName); err != nil { - return nil, fmt.Errorf("error scanning column while getting selected columns: %w", err) - } - columns = append(columns, columnName) + columns, err := pgx.CollectRows(rows, pgx.RowTo[string]) + if err != nil { + return nil, fmt.Errorf("error scanning columns while getting selected columns for table %s: %w", sourceTable, err) } return columns, nil