Merge pull request #51612 from yuzefovich/upsert-19.2

release-19.2: sql: fix pagination in UPSERT

pkg/sql/logictest/testdata/logic_test/upsert

@@ -985,3 +985,24 @@ query III
 SELECT * from table38627
 ----
 1  1  5
+
+# Regression test for UPSERT batching logic (#51391).
+statement ok
+SET CLUSTER SETTING kv.raft.command.max_size='4MiB';
+CREATE TABLE src (s STRING);
+CREATE TABLE dest (s STRING);
+INSERT INTO src
+SELECT
+	'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'
+FROM
+	generate_series(1, 50000)
+
+# This statement produces a raft command of about 6.6 MiB in size, so if the
+# batching logic is incorrect, we'll encounter "command is too large" error.
+statement ok
+UPSERT INTO dest (s) (SELECT s FROM src)
+
+statement ok
+RESET CLUSTER SETTING kv.raft.command.max_size;
+DROP TABLE src;
+DROP TABLE dest

pkg/sql/tablewriter_upsert_opt.go

@@ -92,7 +92,7 @@ func (tu *optTableUpserter) init(txn *client.Txn, evalCtx *tree.EvalContext) err
 		tu.rowsUpserted = rowcontainer.NewRowContainer(
 			evalCtx.Mon.MakeBoundAccount(),
 			sqlbase.ColTypeInfoFromColDescs(tu.returnCols),
-			tu.insertRows.Len(),
+			0, /* rowCapacity */
 		)
 	}
 
@@ -160,6 +160,11 @@ func (tu *optTableUpserter) atBatchEnd(ctx context.Context, traceKV bool) error
 	return nil
 }
 
+// curBatchSize is part of the extendedTableWriter interface. Note that we need
+// to override tableUpserterBase.curBatchSize because the optimizer-driven
+// UPSERTs do not store the insert rows in insertRows row container.
+func (tu *optTableUpserter) curBatchSize() int { return tu.batchSize }
+
 // insertNonConflictingRow inserts the given source row into the table when
 // there was no conflict. If the RETURNING clause was specified, then the
 // inserted row is stored in the rowsUpserted collection.