logical: support array primary keys in the crud writer

This adds a second point select row reader to the LDR crud writer. This
is a work around for #32552.

This change also includes a small fix to the batch row reader. It now
initializes the array parameter with the canonical type. This is
required because the arguments come from the cdc event decoder and the
event decoder always returns canonical datums.

Release note: none
Fixes: #148303

Author: jeffswenson

pkg/crosscluster/logical/replication_statements.go

@@ -475,6 +475,129 @@ func newBulkSelectStatement(
 	return stmt, paramTypes, nil
 }
 
+// newPointSelectStatement returns a statement that can be used to query
+// a single row by primary key. Unlike newBulkSelectStatement which handles
+// multiple rows with arrays, this generates a simple SELECT statement with
+// individual parameters for each primary key column.
+//
+// The statement will have one parameter for each primary key column, where
+// each parameter is the value for that column. The columns are expected in
+// column ID order.
+//
+// For example, given a table with primary key columns (id, secondary_id) and
+// additional columns (value1, value2), the generated statement would be
+// equivalent to:
+//
+//	SELECT
+//		replication_target.crdb_internal_origin_timestamp,
+//		replication_target.crdb_internal_mvcc_timestamp,
+//		replication_target.id, replication_target.secondary_id,
+//		replication_target.value1, replication_target.value2
+//	FROM [table_id AS replication_target]
+//	WHERE replication_target.id = $1 AND replication_target.secondary_id = $2
+func newPointSelectStatement(
+	table catalog.TableDescriptor,
+) (statements.Statement[tree.Statement], []*types.T, error) {
+	cols := getColumnSchema(table)
+	primaryKeyColumns := make([]columnSchema, 0, len(cols))
+	for _, col := range cols {
+		if col.isPrimaryKey {
+			primaryKeyColumns = append(primaryKeyColumns, col)
+		}
+	}
+
+	// Create parameter types for primary key values
+	paramTypes := make([]*types.T, 0, len(primaryKeyColumns))
+	for _, pkCol := range primaryKeyColumns {
+		paramTypes = append(paramTypes, pkCol.columnType)
+	}
+
+	// Create the table reference for `replication_target`
+	targetName, err := tree.NewUnresolvedObjectName(1, [3]string{"replication_target"}, tree.NoAnnotation)
+	if err != nil {
+		return statements.Statement[tree.Statement]{}, nil, err
+	}
+
+	// Build the SELECT clause columns: timestamps first, then all table columns
+	selectColumns := make(tree.SelectExprs, 0, 2+len(cols))
+
+	// Add `replication_target.crdb_internal_origin_timestamp`
+	selectColumns = append(selectColumns, tree.SelectExpr{
+		Expr: &tree.ColumnItem{
+			ColumnName: "crdb_internal_origin_timestamp",
+			TableName:  targetName,
+		},
+	})
+	// Add `replication_target.crdb_internal_mvcc_timestamp`
+	selectColumns = append(selectColumns, tree.SelectExpr{
+		Expr: &tree.ColumnItem{
+			ColumnName: "crdb_internal_mvcc_timestamp",
+			TableName:  targetName,
+		},
+	})
+
+	// Add all table columns: `replication_target.column_name`
+	for _, col := range cols {
+		selectColumns = append(selectColumns, tree.SelectExpr{
+			Expr: &tree.ColumnItem{
+				ColumnName: tree.Name(col.column.GetName()),
+				TableName:  targetName,
+			},
+		})
+	}
+
+	// Build the WHERE clause: `replication_target.pk_col1 = $1 AND replication_target.pk_col2 = $2`
+	var whereClause tree.Expr
+	for i, pkCol := range primaryKeyColumns {
+		placeholder, err := newTypedPlaceholder(i+1, pkCol.column)
+		if err != nil {
+			return statements.Statement[tree.Statement]{}, nil, err
+		}
+
+		eqExpr := &tree.ComparisonExpr{
+			// Use EQ operator to compare primary key columns because primary key
+			// columns are guaranteed to be non-NULL.
+			Operator: treecmp.MakeComparisonOperator(treecmp.EQ),
+			Left: &tree.ColumnItem{
+				TableName:  targetName,
+				ColumnName: tree.Name(pkCol.column.GetName()),
+			},
+			Right: placeholder,
+		}
+
+		if i == 0 {
+			whereClause = eqExpr
+		} else {
+			whereClause = &tree.AndExpr{
+				Left:  whereClause,
+				Right: eqExpr,
+			}
+		}
+	}
+
+	// Construct the complete SELECT statement
+	selectStmt := &tree.Select{
+		Select: &tree.SelectClause{
+			Exprs: selectColumns,
+			From: tree.From{
+				Tables: tree.TableExprs{
+					&tree.TableRef{
+						TableID: int64(table.GetID()),
+						As:      tree.AliasClause{Alias: "replication_target"},
+					},
+				},
+			},
+			Where: &tree.Where{Type: tree.AstWhere, Expr: whereClause},
+		},
+	}
+
+	stmt, err := toParsedStatement(selectStmt)
+	if err != nil {
+		return statements.Statement[tree.Statement]{}, nil, err
+	}
+	return stmt, paramTypes, nil
+}
+
 func toParsedStatement(stmt tree.Statement) (statements.Statement[tree.Statement], error) {
 	// User Serialize instead of String to ensure the type casts use fully
 	// qualified names.

pkg/crosscluster/logical/replication_statements_test.go

@@ -130,6 +130,18 @@ func TestReplicationStatements(t *testing.T) {
 
 				prepareStatement(t, sqlDB, types, stmt)
 
+				return stmt.SQL
+			case "show-point-select":
+				var tableName string
+				d.ScanArgs(t, "table", &tableName)
+
+				desc := getTableDesc(tableName)
+
+				stmt, types, err := newPointSelectStatement(desc)
+				require.NoError(t, err)
+
+				prepareStatement(t, sqlDB, types, stmt)
+
 				return stmt.SQL
 			default:
 				return "unknown command: " + d.Cmd

pkg/crosscluster/logical/sql_row_reader.go

@@ -11,18 +11,18 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/errors"
 )
 
-type sqlRowReader struct {
-	session isql.Session
-
-	selectStatement isql.PreparedStatement
-
-	// keyColumnIndices is the index of the datums that are part of the primary key.
-	keyColumnIndices []int
-	columns          []columnSchema
+type sqlRowReader interface {
+	// ReadRows reads the rows from the table using the provided transaction. A row
+	// will only be present in the result set if it exists. The index of the row in
+	// the input is the key to the output map.
+	//
+	// E.g. result[i] and rows[i] are the same row.
+	ReadRows(ctx context.Context, rows []tree.Datums) (map[int]priorRow, error)
 }
 
 // priorRow is a row returned by the SQL reader. It contains the rows local
@@ -42,7 +42,37 @@ type priorRow struct {
 
 func newSQLRowReader(
 	ctx context.Context, table catalog.TableDescriptor, session isql.Session,
-) (*sqlRowReader, error) {
+) (sqlRowReader, error) {
+	hasArrayPrimaryKey := false
+	for _, col := range getColumnSchema(table) {
+		if col.isPrimaryKey && col.columnType.Family() == types.ArrayFamily {
+			hasArrayPrimaryKey = true
+			break
+		}
+	}
+	if hasArrayPrimaryKey {
+		// TODO(#32552): delete point row reader once CockroachDB supports nested
+		// array types. We can't use the bulk reader because it passes all of the
+		// primary key values in an array, which results in an array of arrays when
+		// a primary key column is an array.
+		return newPointRowReader(ctx, table, session)
+	}
+	return newBulkRowReader(ctx, table, session)
+}
+
+type bulkRowReader struct {
+	session isql.Session
+
+	selectStatement isql.PreparedStatement
+
+	// keyColumnIndices is the index of the datums that are part of the primary key.
+	keyColumnIndices []int
+	columns          []columnSchema
+}
+
+func newBulkRowReader(
+	ctx context.Context, table catalog.TableDescriptor, session isql.Session,
+) (*bulkRowReader, error) {
 	cols := getColumnSchema(table)
 	keyColumns := make([]int, 0, len(cols))
 	for i, col := range cols {
@@ -60,20 +90,17 @@ func newSQLRowReader(
 		return nil, err
 	}
 
-	return &sqlRowReader{
+	return &bulkRowReader{
 		session:          session,
 		selectStatement:  selectStatement,
 		keyColumnIndices: keyColumns,
 		columns:          cols,
 	}, nil
 }
 
-// ReadRows reads the rows from the table using the provided transaction. A row
-// will only be present in the result set if it exists. The index of the row in
-// the input is the key to the output map.
-//
-// E.g. result[i] and rows[i] are the same row.
-func (r *sqlRowReader) ReadRows(ctx context.Context, rows []tree.Datums) (map[int]priorRow, error) {
+func (r *bulkRowReader) ReadRows(
+	ctx context.Context, rows []tree.Datums,
+) (map[int]priorRow, error) {
 	// TODO(jeffswenson): optimize allocations. It may require a change to the
 	// API. For now, this probably isn't a performance bottleneck because:
 	// 1. Many of the allocations are one per batch instead of one per row.
@@ -86,7 +113,7 @@ func (r *sqlRowReader) ReadRows(ctx context.Context, rows []tree.Datums) (map[in
 
 	params := make([]tree.Datum, 0, len(r.keyColumnIndices))
 	for _, index := range r.keyColumnIndices {
-		array := tree.NewDArray(r.columns[index].column.GetType())
+		array := tree.NewDArray(r.columns[index].columnType)
 		for _, row := range rows {
 			if err := array.Append(row[index]); err != nil {
 				return nil, err
@@ -145,3 +172,105 @@ func (r *sqlRowReader) ReadRows(ctx context.Context, rows []tree.Datums) (map[in
 
 	return result, nil
 }
+
+type pointReadRowReader struct {
+	session isql.Session
+
+	selectStatement isql.PreparedStatement
+
+	// keyColumnIndices is the index of the datums that are part of the primary key.
+	keyColumnIndices []int
+	columns          []columnSchema
+}
+
+func newPointRowReader(
+	ctx context.Context, table catalog.TableDescriptor, session isql.Session,
+) (*pointReadRowReader, error) {
+	cols := getColumnSchema(table)
+	keyColumns := make([]int, 0, len(cols))
+	for i, col := range cols {
+		if col.isPrimaryKey {
+			keyColumns = append(keyColumns, i)
+		}
+	}
+
+	selectStatementRaw, types, err := newPointSelectStatement(table)
+	if err != nil {
+		return nil, err
+	}
+	selectStatement, err := session.Prepare(ctx, "replication-read-point", selectStatementRaw, types)
+	if err != nil {
+		return nil, err
+	}
+
+	return &pointReadRowReader{
+		session:          session,
+		selectStatement:  selectStatement,
+		keyColumnIndices: keyColumns,
+		columns:          cols,
+	}, nil
+}
+
+func (p *pointReadRowReader) ReadRows(
+	ctx context.Context, rows []tree.Datums,
+) (map[int]priorRow, error) {
+	if len(rows) == 0 {
+		return nil, nil
+	}
+
+	result := make(map[int]priorRow, len(rows))
+
+	for i, row := range rows {
+		params := make([]tree.Datum, 0, len(p.keyColumnIndices))
+		for _, keyIndex := range p.keyColumnIndices {
+			params = append(params, row[keyIndex])
+		}
+
+		queryRows, err := p.session.QueryPrepared(ctx, p.selectStatement, params)
+		if err != nil {
+			return nil, err
+		}
+
+		if len(queryRows) > 1 {
+			return nil, errors.AssertionFailedf("expected at most 1 row, got %d", len(queryRows))
+		}
+		if len(queryRows) == 0 {
+			continue
+		}
+
+		resultRow := queryRows[0]
+		// The columns are:
+		// 0. The origin timestamp.
+		// 1. The mvcc timestamp.
+		// 2+. The table columns.
+		const prefixColumns = 2
+		if len(resultRow) != len(p.columns)+prefixColumns {
+			return nil, errors.AssertionFailedf("expected %d columns, got %d", len(p.columns)+prefixColumns, len(resultRow))
+		}
+
+		isLocal := false
+		timestamp := resultRow[0]
+		if timestamp == tree.DNull {
+			timestamp = resultRow[1]
+			isLocal = true
+		}
+
+		decimal, ok := timestamp.(*tree.DDecimal)
+		if !ok {
+			return nil, errors.AssertionFailedf("expected column 0 or 1 to be origin timestamp")
+		}
+
+		logicalTimestamp, err := hlc.DecimalToHLC(&decimal.Decimal)
+		if err != nil {
+			return nil, err
+		}
+
+		result[i] = priorRow{
+			row:              resultRow[prefixColumns:],
+			logicalTimestamp: logicalTimestamp,
+			isLocal:          isLocal,
+		}
+	}
+
+	return result, nil
+}