Merge #144225 #144240

144225: kvcoord: add request validation to txnWriteBuffer r=yuzefovich a=stevendanna

This adds early detection of unsupported request types and request
options. SendLocked now errors on any request that the txnWriteBuffer
doesn't explicitly support. We may consider flushing the buffer
instead to increase the set of transactions that can use this feature.

Fixes #143894

Release note: None

144240: sql/schemachanger: avoid falling back when adding multiple pks r=fqazi a=fqazi

Previously, the declarative schema changer would incorrectly fallback if the user was attempting to add multiple primary keys. This was a gap in the logic for automatically cleaning up default pks, which return a not implemented error if the rowid column was not found. To address this, this patch will always generates the correct error.

Fixes: #144235

Release note: None

Co-authored-by: Steven Danna <[email protected]>
Co-authored-by: Faizan Qazi <[email protected]>

Author: 3 people

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer.go

@@ -20,6 +20,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/storage/enginepb"
 	"github.com/cockroachdb/cockroach/pkg/storage/mvccencoding"
 	"github.com/cockroachdb/cockroach/pkg/storage/mvcceval"
+	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
@@ -46,8 +47,8 @@ var bufferedWritesMaxBufferSize = settings.RegisterIntSetting(
 
 // txnWriteBuffer is a txnInterceptor that buffers transactional writes until
 // commit time. Moreover, it also decomposes read-write KV operations (e.g.
-// CPuts, InitPuts) into separate (locking) read and write operations, buffering
-// the latter until commit time.
+// CPuts) into separate (locking) read and write operations, buffering the
+// latter until commit time.
 //
 // Buffering writes until commit time has four main benefits:
 //
@@ -253,7 +254,16 @@ func (twb *txnWriteBuffer) SendLocked(
 		return twb.flushBufferAndSendBatch(ctx, ba)
 	}
 
-	transformedBa, ts := twb.applyTransformations(ctx, ba)
+	if err := twb.validateBatch(ba); err != nil {
+		// We could choose to twb.flushBufferAndSendBatch
+		// here. For now, we return an error.
+		return nil, kvpb.NewError(err)
+	}
+
+	transformedBa, ts, pErr := twb.applyTransformations(ctx, ba)
+	if pErr != nil {
+		return nil, pErr
+	}
 
 	if len(transformedBa.Requests) == 0 {
 		// Lower layers (the DistSender and the KVServer) do not expect/handle empty
@@ -278,6 +288,80 @@ func (twb *txnWriteBuffer) SendLocked(
 	return twb.mergeResponseWithTransformations(ctx, ts, br)
 }
 
+// validateBatch returns an error if the batch is unsupported
+// by the txnWriteBuffer.
+func (twb *txnWriteBuffer) validateBatch(ba *kvpb.BatchRequest) error {
+	if ba.WriteOptions != nil {
+		// OriginTimestamp and OriginID are currently only used by Logical Data
+		// Replication (LDR). These options are unsupported at the moment as we
+		// don't store the inbound batch options in the buffer.
+		if ba.WriteOptions.OriginTimestamp.IsSet() {
+			return errors.AssertionFailedf("transaction write buffer does not support batches with OriginTimestamp set")
+		}
+		if ba.WriteOptions.OriginID != 0 {
+			return errors.AssertionFailedf("transaction write buffer does not support batches with OriginID set")
+		}
+	}
+	return twb.validateRequests(ba)
+}
+
+// validateRequests returns an error if any of the requests in the batch
+// are unsupported by the txnWriteBuffer.
+func (twb *txnWriteBuffer) validateRequests(ba *kvpb.BatchRequest) error {
+	for _, ru := range ba.Requests {
+		req := ru.GetInner()
+		switch t := req.(type) {
+		case *kvpb.ConditionalPutRequest:
+			// Our client side ConditionalPutRequest evaluation does not know how to
+			// handle the origin timestamp setting. Doing so would require sending a
+			// GetRequest with RawMVCCValues set and parsing the MVCCValueHeader.
+			if t.OriginTimestamp.IsSet() {
+				return unsupportedOptionError(t.Method(), "OriginTimestamp")
+			}
+		case *kvpb.PutRequest:
+		case *kvpb.DeleteRequest:
+		case *kvpb.GetRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+		case *kvpb.ScanRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+			if t.ScanFormat == kvpb.COL_BATCH_RESPONSE {
+				return unsupportedOptionError(t.Method(), "COL_BATCH_RESPONSE scan format")
+			}
+		case *kvpb.ReverseScanRequest:
+			// ReturnRawMVCCValues is unsupported because we don't know how to serve
+			// such reads from the write buffer currently.
+			if t.ReturnRawMVCCValues {
+				return unsupportedOptionError(t.Method(), "ReturnRawMVCCValue")
+			}
+			if t.ScanFormat == kvpb.COL_BATCH_RESPONSE {
+				return unsupportedOptionError(t.Method(), "COL_BATCH_RESPONSE scan format")
+			}
+		default:
+			// All other requests are unsupported. Note that we assume EndTxn and
+			// DeleteRange requests were handled explicitly before this method was
+			// called.
+			return unsupportedMethodError(t.Method())
+		}
+	}
+	return nil
+}
+
+func unsupportedMethodError(m kvpb.Method) error {
+	return errors.AssertionFailedf("transaction write buffer does not support %s requests", m)
+}
+
+func unsupportedOptionError(m kvpb.Method, option string) error {
+	return errors.AssertionFailedf("transaction write buffer does not support %s requests with %s", m, option)
+}
+
 // estimateSize returns a conservative estimate by which the buffer will grow in
 // size if the writes from the supplied batch request are buffered.
 func (twb *txnWriteBuffer) estimateSize(ba *kvpb.BatchRequest) int64 {
@@ -539,7 +623,7 @@ func (twb *txnWriteBuffer) closeLocked() {}
 // TODO(arul): Augment this comment as these expand.
 func (twb *txnWriteBuffer) applyTransformations(
 	ctx context.Context, ba *kvpb.BatchRequest,
-) (*kvpb.BatchRequest, transformations) {
+) (*kvpb.BatchRequest, transformations, *kvpb.Error) {
 	baRemote := ba.ShallowCopy()
 	// TODO(arul): We could improve performance here by pre-allocating
 	// baRemote.Requests to the correct size by counting the number of Puts/Dels
@@ -699,9 +783,6 @@ func (twb *txnWriteBuffer) applyTransformations(
 				// We've constructed a response that we'll stitch together with the
 				// result on the response path; eschew sending the request to the KV
 				// layer.
-				//
-				// TODO(arul): if the ReturnRawMVCCValues flag is set, we'll need to
-				// flush the buffer.
 				continue
 			}
 			// Wasn't served locally; send the request to the KV layer.
@@ -736,10 +817,10 @@ func (twb *txnWriteBuffer) applyTransformations(
 			baRemote.Requests = append(baRemote.Requests, ru)
 
 		default:
-			baRemote.Requests = append(baRemote.Requests, ru)
+			return nil, nil, kvpb.NewError(unsupportedMethodError(t.Method()))
 		}
 	}
-	return baRemote, ts
+	return baRemote, ts, nil
 }
 
 // seekItemForSpan returns a bufferedWrite appropriate for use with a
@@ -1103,9 +1184,7 @@ func (t transformation) toResp(
 		ru.MustSetInner(reverseScanResp)
 
 	default:
-		// This is only possible once we start decomposing read-write requests into
-		// separate bits.
-		panic("unimplemented")
+		return ru, kvpb.NewError(unsupportedMethodError(req.Method()))
 	}
 
 	return ru, nil
@@ -1725,7 +1804,7 @@ func (m *respMerger) toReverseScanResp(
 
 // assertTrue panics with a message if the supplied condition isn't true.
 func assertTrue(cond bool, msg string) {
-	if !cond {
+	if !cond && buildutil.CrdbTestBuild {
 		panic(msg)
 	}
 }

pkg/kv/kvclient/kvcoord/txn_interceptor_write_buffer_test.go

@@ -1851,3 +1851,139 @@ func TestTxnWriteBufferClearsBufferOnEpochBump(t *testing.T) {
 	require.Equal(t, 0, int(twb.bufferSize))
 	require.Equal(t, numCalled, mockSender.NumCalled())
 }
+
+// TestTxnWriteBufferBatchRequestValidation verifies that the txnWriteBuffer
+// rejects requests that it doesn't know how to support.
+func TestTxnWriteBufferBatchRequestValidation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	twb, mockSender := makeMockTxnWriteBuffer(cluster.MakeClusterSettings())
+
+	type testCase struct {
+		name string
+		ba   func() *kvpb.BatchRequest
+	}
+
+	txn := makeTxnProto()
+	txn.Sequence = 1
+	keyA, keyC := roachpb.Key("a"), roachpb.Key("c")
+
+	tests := []testCase{
+		{
+			name: "batch with OriginTimestamp",
+			ba: func() *kvpb.BatchRequest {
+				header := kvpb.Header{
+					Txn: &txn,
+					WriteOptions: &kvpb.WriteOptions{
+						OriginTimestamp: hlc.Timestamp{WallTime: 1},
+					},
+				}
+				return &kvpb.BatchRequest{Header: header}
+			},
+		},
+		{
+			name: "batch with OriginID",
+			ba: func() *kvpb.BatchRequest {
+				header := kvpb.Header{
+					Txn: &txn,
+					WriteOptions: &kvpb.WriteOptions{
+						OriginID: 1,
+					},
+				}
+				return &kvpb.BatchRequest{Header: header}
+			},
+		},
+		{
+			name: "batch with InitPut",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				b.Add(&kvpb.InitPutRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+					Value:         roachpb.Value{},
+				})
+				return b
+			},
+		},
+		{
+			name: "batch with Increment",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				b.Add(&kvpb.IncrementRequest{
+					RequestHeader: kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				})
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues Scan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ScanRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues ReverseScan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ReverseScanRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with ReturnRawMVCCValues Get",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.GetRequest{
+					ReturnRawMVCCValues: true,
+					RequestHeader:       kvpb.RequestHeader{Key: keyA, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with COL_BATCH_RESPONSE Scan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ScanRequest{
+					ScanFormat:    kvpb.COL_BATCH_RESPONSE,
+					RequestHeader: kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+		{
+			name: "batch with COL_BATCH_RESPONSE ReverseScan",
+			ba: func() *kvpb.BatchRequest {
+				b := &kvpb.BatchRequest{Header: kvpb.Header{Txn: &txn}}
+				r := &kvpb.ReverseScanRequest{
+					ScanFormat:    kvpb.COL_BATCH_RESPONSE,
+					RequestHeader: kvpb.RequestHeader{Key: keyA, EndKey: keyC, Sequence: txn.Sequence},
+				}
+				b.Add(r)
+				return b
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			numCalledBefore := mockSender.NumCalled()
+			_, pErr := twb.SendLocked(ctx, tc.ba())
+			require.NotNil(t, pErr)
+			require.Equal(t, numCalledBefore, mockSender.NumCalled())
+
+		})
+	}
+}

pkg/sql/logictest/testdata/logic_test/new_schema_changer

@@ -1540,3 +1540,23 @@ statement ok
 EXPLAIN (DDL) CREATE SCHEMA sc1
 
 subtest end
+
+
+subtest multiple_primary_key
+
+statement ok
+CREATE TABLE multiple_pk_attempt(n INT PRIMARY KEY, j INT)
+
+let $use_decl_sc
+SHOW use_declarative_schema_changer
+
+statement ok
+SET use_declarative_schema_changer = 'unsafe_always'
+
+statement error pgcode 42611 multiple primary keys for table \"multiple_pk_attempt\" are not allowed
+ALTER TABLE multiple_pk_attempt ADD PRIMARY KEY (j);
+
+statement ok
+SET use_declarative_schema_changer = $use_decl_sc
+
+subtest end

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_add_constraint.go

@@ -85,7 +85,12 @@ func alterTableAddPrimaryKey(
 	if getPrimaryIndexDefaultRowIDColumn(
 		b, tbl.TableID, oldPrimaryIndex.IndexID,
 	) == nil {
-		panic(scerrors.NotImplementedError(t))
+		// If the constraint already exists then nothing to do here.
+		if oldPrimaryIndex != nil && d.IfNotExists {
+			return
+		}
+		panic(pgerror.Newf(pgcode.InvalidColumnDefinition,
+			"multiple primary keys for table %q are not allowed", tn.Object()))
 	}
 	alterPrimaryKey(b, tn, tbl, stmt, alterPrimaryKeySpec{
 		n:             t,

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/alter_table_alter_primary_key.go

@@ -447,12 +447,18 @@ func fallBackIfRegionalByRowTable(b BuildCtx, t tree.NodeFormatter, tableID cati
 	}
 }
 
+// mustRetrieveCurrentPrimaryIndexElement retrieves the current primary index,
+// which must be public.
 func mustRetrieveCurrentPrimaryIndexElement(
 	b BuildCtx, tableID catid.DescID,
 ) (res *scpb.PrimaryIndex) {
 	scpb.ForEachPrimaryIndex(b.QueryByID(tableID), func(
 		current scpb.Status, target scpb.TargetStatus, e *scpb.PrimaryIndex,
 	) {
+		// TODO(fqazi): We don't support DROP CONSTRAINT PRIMARY KEY, so there is no
+		// risk of ever seeing a non-public PrimaryIndex element. In the future when
+		// we do support DROP CONSTRAINT PRIMARY KEY, we should adapt callers of
+		// this function to handle the absent primary index case.
 		if current == scpb.Status_PUBLIC {
 			res = e
 		}

pkg/sql/schemachanger/scbuild/testdata/unimplemented_alter_table

@@ -32,10 +32,6 @@ unimplemented
 ALTER TABLE defaultdb.foo EXPERIMENTAL_AUDIT SET READ WRITE
 ----
 
-unimplemented
-ALTER TABLE defaultdb.foo ADD PRIMARY KEY (l);
-----
-
 unimplemented
 ALTER TABLE defaultdb.foo PARTITION BY NOTHING
 ----

pkg/sql/schemachanger/scbuild/testdata/zone_cfg

@@ -46,10 +46,6 @@ ALTER TABLE defaultdb.foo_index_zone_cfg ADD COLUMN j INT
 - [[IndexColumn:{DescID: 105, ColumnID: 7, IndexID: 1}, PUBLIC], ABSENT]
   {columnId: 7, indexId: 1, kind: STORED, ordinalInKind: 5, tableId: 105}
 
-unimplemented
-ALTER TABLE defaultdb.foo_index_zone_cfg ADD PRIMARY KEY (i, n)
-----
-
 build
 ALTER TABLE defaultdb.foo_index_zone_cfg DROP COLUMN k
 ----