pg qrep: don't assume all columns nullable (#3515)

following along from #3504, that PR likely unnecessary since I think hamba infers as much

Author: serprex

flow/connectors/postgres/qrep_query_executor.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"maps"
+	"slices"
 
 	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgconn"
@@ -58,43 +60,68 @@ func (qe *QRepQueryExecutor) ExecuteQuery(ctx context.Context, query string, arg
 	return rows, nil
 }
 
-func (qe *QRepQueryExecutor) executeQueryInTx(ctx context.Context, tx pgx.Tx, cursorName string, fetchSize int) (pgx.Rows, error) {
-	qe.logger.Info("Executing query in transaction")
-	q := fmt.Sprintf("FETCH %d FROM %s", fetchSize, cursorName)
+// FieldDescriptionsToSchema converts a slice of pgconn.FieldDescription to a QRecordSchema.
+func (qe *QRepQueryExecutor) cursorToSchema(
+	ctx context.Context,
+	tx pgx.Tx,
+	cursorName string,
+) (types.QRecordSchema, error) {
+	type attId struct {
+		relid uint32
+		num   uint16
+	}
 
-	rows, err := tx.Query(ctx, q)
+	rows, err := tx.Query(ctx, "FETCH 0 FROM "+cursorName)
 	if err != nil {
-		qe.logger.Error("[pg_query_executor] failed to execute query in tx", slog.Any("error", err))
-		return nil, err
+		return types.QRecordSchema{}, fmt.Errorf("failed to fetch 0 for field descriptions: %w", err)
 	}
-
-	return rows, nil
-}
-
-// FieldDescriptionsToSchema converts a slice of pgconn.FieldDescription to a QRecordSchema.
-func (qe *QRepQueryExecutor) fieldDescriptionsToSchema(fds []pgconn.FieldDescription) types.QRecordSchema {
+	fds := rows.FieldDescriptions()
+	tableOIDset := make(map[uint32]struct{})
+	nullPointers := make(map[attId]*bool, len(fds))
 	qfields := make([]types.QField, len(fds))
 	for i, fd := range fds {
+		tableOIDset[fd.TableOID] = struct{}{}
 		ctype := qe.postgresOIDToQValueKind(fd.DataTypeOID, qe.customTypeMapping, qe.version)
-		// there isn't a way to know if a column is nullable or not
 		if ctype == types.QValueKindNumeric || ctype == types.QValueKindArrayNumeric {
 			precision, scale := datatypes.ParseNumericTypmod(fd.TypeModifier)
 			qfields[i] = types.QField{
 				Name:      fd.Name,
 				Type:      ctype,
-				Nullable:  true,
+				Nullable:  false,
 				Precision: precision,
 				Scale:     scale,
 			}
 		} else {
 			qfields[i] = types.QField{
 				Name:     fd.Name,
 				Type:     ctype,
-				Nullable: true,
+				Nullable: false,
 			}
 		}
+		nullPointers[attId{
+			relid: fd.TableOID,
+			num:   fd.TableAttributeNumber,
+		}] = &qfields[i].Nullable
 	}
-	return types.NewQRecordSchema(qfields)
+	rows.Close()
+	tableOIDs := slices.Collect(maps.Keys(tableOIDset))
+
+	rows, err = tx.Query(ctx, "SELECT a.attrelid,a.attnum FROM pg_attribute a WHERE a.attrelid = ANY($1) AND NOT a.attnotnull", tableOIDs)
+	if err != nil {
+		return types.QRecordSchema{}, fmt.Errorf("failed to query schema for field descriptions: %w", err)
+	}
+
+	var att attId
+	if _, err := pgx.ForEachRow(rows, []any{&att.relid, &att.num}, func() error {
+		if nullPointer, ok := nullPointers[att]; ok {
+			*nullPointer = true
+		}
+		return nil
+	}); err != nil {
+		return types.QRecordSchema{}, fmt.Errorf("failed to process schema for field descriptions: %w", err)
+	}
+
+	return types.NewQRecordSchema(qfields), nil
 }
 
 func (qe *QRepQueryExecutor) processRowsStream(
@@ -152,20 +179,17 @@ func (qe *QRepQueryExecutor) processFetchedRows(
 	fetchSize int,
 	stream *model.QRecordStream,
 ) (int64, int64, error) {
-	rows, err := qe.executeQueryInTx(ctx, tx, cursorName, fetchSize)
+	qe.logger.Info("[pg_query_executor] fetching from cursor", slog.String("cursor", cursorName))
+
+	rows, err := tx.Query(ctx, fmt.Sprintf("FETCH %d FROM %s", fetchSize, cursorName))
 	if err != nil {
-		qe.logger.Error("[pg_query_executor] failed to execute query in tx",
+		qe.logger.Error("[pg_query_executor] failed to fetch cursor in tx",
 			slog.Any("error", err), slog.String("query", query))
 		return 0, 0, fmt.Errorf("[pg_query_executor] failed to execute query in tx: %w", err)
 	}
 	defer rows.Close()
 
 	fieldDescriptions := rows.FieldDescriptions()
-	if !stream.IsSchemaSet() {
-		schema := qe.fieldDescriptionsToSchema(fieldDescriptions)
-		stream.SetSchema(schema)
-	}
-
 	numRows, numBytes, err := qe.processRowsStream(ctx, cursorName, stream, rows, fieldDescriptions)
 	if err != nil {
 		qe.logger.Error("[pg_query_executor] failed to process rows", slog.Any("error", err))

flow/connectors/postgres/sink_q.go

@@ -45,7 +45,6 @@ func (stream RecordStreamSink) ExecuteQueryWithTx(
 	randomUint := rand.Uint64()
 
 	cursorName := fmt.Sprintf("peerdb_cursor_%d", randomUint)
-	fetchSize := shared.FetchAndChannelSize
 	cursorQuery := fmt.Sprintf("DECLARE %s CURSOR FOR %s", cursorName, query)
 
 	if _, err := tx.Exec(ctx, cursorQuery, args...); err != nil {
@@ -60,10 +59,18 @@ func (stream RecordStreamSink) ExecuteQueryWithTx(
 		slog.String("query", query),
 		slog.Int("channelLen", len(stream.Records)))
 
+	if !stream.IsSchemaSet() {
+		schema, err := qe.cursorToSchema(ctx, tx, cursorName)
+		if err != nil {
+			return 0, 0, err
+		}
+		stream.SetSchema(schema)
+	}
+
 	var totalNumRows int64
 	var totalNumBytes int64
 	for {
-		numRows, numBytes, err := qe.processFetchedRows(ctx, query, tx, cursorName, fetchSize, stream.QRecordStream)
+		numRows, numBytes, err := qe.processFetchedRows(ctx, query, tx, cursorName, shared.FetchAndChannelSize, stream.QRecordStream)
 		if err != nil {
 			qe.logger.Error("[pg_query_executor] failed to process fetched rows", slog.Any("error", err))
 			return totalNumRows, totalNumBytes, err

flow/connectors/snowflake/snowflake.go

@@ -35,7 +35,6 @@ const (
 		_PEERDB_RECORD_TYPE INTEGER NOT NULL, _PEERDB_MATCH_DATA STRING,_PEERDB_BATCH_ID INT,
 		_PEERDB_UNCHANGED_TOAST_COLUMNS STRING)`
 	createDummyTableSQL               = "CREATE TABLE IF NOT EXISTS %s.%s(_PEERDB_DUMMY_COL STRING)"
-	rawTableMultiValueInsertSQL       = "INSERT INTO %s.%s VALUES%s"
 	createNormalizedTableSQL          = "CREATE TABLE IF NOT EXISTS %s(%s)"
 	createOrReplaceNormalizedTableSQL = "CREATE OR REPLACE TABLE %s(%s)"
 	toVariantColumnName               = "VAR_COLS"

flow/e2e/snowflake_test.go

@@ -90,8 +90,8 @@ func (s PeerFlowE2ETestSuiteSF) Test_Invalid_Numeric() {
 	_, err := s.Conn().Exec(s.t.Context(), fmt.Sprintf(`
 		CREATE TABLE IF NOT EXISTS %s (
 			id INT PRIMARY KEY,
-			num1 NUMERIC(100, 50) NOT NULL,
-			num2 NUMERIC(100, 50) NOT NULL
+			num1 NUMERIC(100, 50),
+			num2 NUMERIC(100, 50)
 		);
 	`, srcTableName))
 	require.NoError(s.t, err)
@@ -115,7 +115,7 @@ func (s PeerFlowE2ETestSuiteSF) Test_Invalid_Numeric() {
 	env := ExecutePeerflow(s.t, tc, flowConnConfig)
 	SetupCDCFlowStatusQuery(s.t, env, flowConnConfig)
 
-	EnvWaitFor(s.t, env, 3*time.Minute, "normalize shapes", func() bool {
+	EnvWaitFor(s.t, env, 3*time.Minute, "init", func() bool {
 		records, err := s.sfHelper.ExecuteAndProcessQuery(s.t.Context(), "select num1, num2 from "+dstTableName+" where id = 1")
 		if err != nil || len(records.Records) == 0 {
 			return false
@@ -130,7 +130,7 @@ func (s PeerFlowE2ETestSuiteSF) Test_Invalid_Numeric() {
 		"9999999999999999")
 	EnvNoError(s.t, env, err)
 
-	EnvWaitFor(s.t, env, 3*time.Minute, "normalize shapes", func() bool {
+	EnvWaitFor(s.t, env, 3*time.Minute, "cdc", func() bool {
 		records, err := s.sfHelper.ExecuteAndProcessQuery(s.t.Context(), "select num1, num2 from "+dstTableName+" where id = 2")
 		if err != nil || len(records.Records) == 0 {
 			return false