From 19382ae24c75b94cf5391db8817dff68b26a1a90 Mon Sep 17 00:00:00 2001 From: Evan Wall Date: Thu, 20 Oct 2022 13:34:57 -0400 Subject: [PATCH] multitenant: move SPLIT AT tenant check from SQL to KV layer fixes #90435 Move the SPLIT AT tenant check from the SQL to KV layer to prepare for tenant permissions work. This check will eventually be replaced by a tenant permissions lookup to determine if the tenant can perform the requested KV operation (AdminSplit in this case). `Codec.ForSystemTenant()` checks still exist in the SQL and will be handled in separate PRs. Release note: None --- pkg/BUILD.bazel | 2 + pkg/ccl/backupccl/BUILD.bazel | 2 + .../backupccl/restore_data_processor_test.go | 15 +- .../backupccl/split_and_scatter_processor.go | 3 +- .../testdata/logic_test/crdb_internal_tenant | 2 +- .../testdata/logic_test/tenant_unsupported | 2 +- pkg/kv/batch.go | 6 +- pkg/kv/bulk/buffering_adder.go | 2 +- pkg/kv/bulk/sst_batcher.go | 4 +- pkg/kv/bulk/sst_batcher_test.go | 14 +- pkg/kv/db.go | 3 +- .../kvcoord/dist_sender_server_test.go | 66 ++++- pkg/kv/kvclient/kvcoord/split_test.go | 64 +++-- .../kvclient/kvcoord/txn_coord_sender_test.go | 28 +- pkg/kv/kvclient/kvcoord/txn_test.go | 7 +- pkg/kv/kvnemesis/applier.go | 2 +- pkg/kv/kvserver/client_replica_test.go | 4 +- pkg/kv/kvserver/consistency_queue_test.go | 7 +- .../intent_resolver_integration_test.go | 2 +- pkg/kv/kvserver/range_log_test.go | 21 +- pkg/kv/kvserver/replica_command_test.go | 14 +- pkg/kv/kvserver/replica_test.go | 7 +- pkg/kv/kvserver/ts_maintenance_queue_test.go | 7 +- pkg/kv/txn_external_test.go | 4 +- pkg/roachpb/api.proto | 9 + pkg/rpc/auth_tenant.go | 62 ++--- pkg/rpc/auth_test.go | 2 +- pkg/server/intent_test.go | 7 +- pkg/server/node_test.go | 14 +- pkg/sql/BUILD.bazel | 1 + pkg/sql/backfill.go | 3 +- pkg/sql/importer/BUILD.bazel | 1 + pkg/sql/importer/import_processor_planning.go | 3 +- pkg/sql/multitenant_admin_function_test.go | 263 ++++++++++-------- pkg/sql/oppurpose/BUILD.bazel | 12 + pkg/sql/oppurpose/op_purpose.go | 28 ++ pkg/sql/opt_exec_factory.go | 15 +- pkg/sql/schema_changer.go | 3 +- pkg/sql/schemachanger/scdeps/BUILD.bazel | 1 + pkg/sql/schemachanger/scdeps/exec_deps.go | 3 +- pkg/sql/split.go | 18 +- pkg/sql/tenant.go | 3 +- pkg/sql/unsplit_range_test.go | 2 +- 43 files changed, 499 insertions(+), 239 deletions(-) create mode 100644 pkg/sql/oppurpose/BUILD.bazel create mode 100644 pkg/sql/oppurpose/op_purpose.go diff --git a/pkg/BUILD.bazel b/pkg/BUILD.bazel index 99273d43f9f0..e718b0a8d47a 100644 --- a/pkg/BUILD.bazel +++ b/pkg/BUILD.bazel @@ -1545,6 +1545,7 @@ GO_TARGETS = [ "//pkg/sql/mutations:mutations", "//pkg/sql/oidext:oidext", "//pkg/sql/oidext:oidext_test", + "//pkg/sql/oppurpose:oppurpose", "//pkg/sql/opt/bench:bench", "//pkg/sql/opt/bench:bench_test", "//pkg/sql/opt/cat:cat", @@ -2710,6 +2711,7 @@ GET_X_DATA_TARGETS = [ "//pkg/sql/memsize:get_x_data", "//pkg/sql/mutations:get_x_data", "//pkg/sql/oidext:get_x_data", + "//pkg/sql/oppurpose:get_x_data", "//pkg/sql/opt:get_x_data", "//pkg/sql/opt/bench:get_x_data", "//pkg/sql/opt/cat:get_x_data", diff --git a/pkg/ccl/backupccl/BUILD.bazel b/pkg/ccl/backupccl/BUILD.bazel index 3be3c74ceceb..7e84f5d6a23c 100644 --- a/pkg/ccl/backupccl/BUILD.bazel +++ b/pkg/ccl/backupccl/BUILD.bazel @@ -96,6 +96,7 @@ go_library( "//pkg/sql/execinfra", "//pkg/sql/execinfrapb", "//pkg/sql/exprutil", + "//pkg/sql/oppurpose", "//pkg/sql/parser", "//pkg/sql/pgwire/pgcode", "//pkg/sql/pgwire/pgerror", @@ -245,6 +246,7 @@ go_test( "//pkg/sql/execinfra", "//pkg/sql/execinfrapb", "//pkg/sql/importer", + "//pkg/sql/oppurpose", "//pkg/sql/parser", "//pkg/sql/pgwire/pgerror", "//pkg/sql/randgen", diff --git a/pkg/ccl/backupccl/restore_data_processor_test.go b/pkg/ccl/backupccl/restore_data_processor_test.go index 306ba223b37b..806c8f2b8850 100644 --- a/pkg/ccl/backupccl/restore_data_processor_test.go +++ b/pkg/ccl/backupccl/restore_data_processor_test.go @@ -36,6 +36,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/execinfra" "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/sem/eval" "github.com/cockroachdb/cockroach/pkg/storage" "github.com/cockroachdb/cockroach/pkg/testutils" @@ -359,10 +360,20 @@ func runTestIngest(t *testing.T, init func(*cluster.Settings)) { t.Fatalf("failed to rewrite key: %s", reqMidKey2) } - if err := kvDB.AdminSplit(ctx, reqMidKey1, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + reqMidKey1, + hlc.MaxTimestamp, /* expirationTime */ + oppurpose.SplitBackup, + ); err != nil { t.Fatal(err) } - if err := kvDB.AdminSplit(ctx, reqMidKey2, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + reqMidKey2, + hlc.MaxTimestamp, /* expirationTime */ + oppurpose.SplitBackup, + ); err != nil { t.Fatal(err) } diff --git a/pkg/ccl/backupccl/split_and_scatter_processor.go b/pkg/ccl/backupccl/split_and_scatter_processor.go index f44e5c46f152..51908f6ecc37 100644 --- a/pkg/ccl/backupccl/split_and_scatter_processor.go +++ b/pkg/ccl/backupccl/split_and_scatter_processor.go @@ -21,6 +21,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/execinfra" "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" "github.com/cockroachdb/cockroach/pkg/sql/rowexec" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" @@ -96,7 +97,7 @@ func (s dbSplitAndScatterer) split( newSplitKey = splitAt } log.VEventf(ctx, 1, "presplitting new key %+v", newSplitKey) - if err := s.db.AdminSplit(ctx, newSplitKey, expirationTime); err != nil { + if err := s.db.AdminSplit(ctx, newSplitKey, expirationTime, oppurpose.SplitBackup); err != nil { return errors.Wrapf(err, "splitting key %s", newSplitKey) } diff --git a/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal_tenant b/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal_tenant index 6d14ccb90af7..89bbe7e8b60c 100644 --- a/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal_tenant +++ b/pkg/ccl/logictestccl/testdata/logic_test/crdb_internal_tenant @@ -390,7 +390,7 @@ test {} statement ok CREATE TABLE foo (a INT PRIMARY KEY, INDEX idx(a)); INSERT INTO foo VALUES(1) -statement error unsupported in multi-tenancy mode +statement error pq: rpc error: code = Unauthenticated desc = request \[1 AdmSplit\] not permitted ALTER TABLE foo SPLIT AT VALUES(2) # Make sure that the cluster id isn't unset. diff --git a/pkg/ccl/logictestccl/testdata/logic_test/tenant_unsupported b/pkg/ccl/logictestccl/testdata/logic_test/tenant_unsupported index 1eec331aca5f..2ba8cfdfdb0c 100644 --- a/pkg/ccl/logictestccl/testdata/logic_test/tenant_unsupported +++ b/pkg/ccl/logictestccl/testdata/logic_test/tenant_unsupported @@ -45,7 +45,7 @@ SELECT * FROM crdb_internal.kv_node_status # Cannot perform operations that issue Admin requests. -statement error operation is unsupported in multi-tenancy mode +statement error pq: rpc error: code = Unauthenticated desc = request \[1 AdmSplit\] not permitted ALTER TABLE kv SPLIT AT VALUES ('foo') statement error operation is unsupported in multi-tenancy mode diff --git a/pkg/kv/batch.go b/pkg/kv/batch.go index bd35a786f59d..890e4472146c 100644 --- a/pkg/kv/batch.go +++ b/pkg/kv/batch.go @@ -699,7 +699,10 @@ func (b *Batch) adminMerge(key interface{}) { // adminSplit is only exported on DB. It is here for symmetry with the // other operations. func (b *Batch) adminSplit( - splitKeyIn interface{}, expirationTime hlc.Timestamp, predicateKeys []roachpb.Key, + splitKeyIn interface{}, + expirationTime hlc.Timestamp, + class roachpb.AdminSplitRequest_Class, + predicateKeys []roachpb.Key, ) { splitKey, err := marshalKey(splitKeyIn) if err != nil { @@ -713,6 +716,7 @@ func (b *Batch) adminSplit( SplitKey: splitKey, ExpirationTime: expirationTime, PredicateKeys: predicateKeys, + Class: class, } b.appendReqs(req) b.initResult(1, 0, notRaw, nil) diff --git a/pkg/kv/bulk/buffering_adder.go b/pkg/kv/bulk/buffering_adder.go index e04a99213f02..82d073a8a040 100644 --- a/pkg/kv/bulk/buffering_adder.go +++ b/pkg/kv/bulk/buffering_adder.go @@ -403,7 +403,7 @@ func (b *BufferingAdder) createInitialSplits(ctx context.Context) error { } predicateKey := b.curBuf.Key(predicateAt) log.VEventf(ctx, 1, "pre-splitting span %d of %d at %s", i, b.initialSplits, splitKey) - if err := b.sink.db.AdminSplit(ctx, splitKey, expire, predicateKey); err != nil { + if err := b.sink.db.AdminSplit(ctx, splitKey, expire, roachpb.AdminSplitRequest_INGESTION, predicateKey); err != nil { // TODO(dt): a typed error would be nice here. if strings.Contains(err.Error(), "predicate") { log.VEventf(ctx, 1, "%s adder split at %s rejected, had previously split and no longer included %s", diff --git a/pkg/kv/bulk/sst_batcher.go b/pkg/kv/bulk/sst_batcher.go index c5e33a95bb5c..0a543fae6fef 100644 --- a/pkg/kv/bulk/sst_batcher.go +++ b/pkg/kv/bulk/sst_batcher.go @@ -521,7 +521,7 @@ func (b *SSTBatcher) doFlush(ctx context.Context, reason int) error { log.Warningf(ctx, "%s failed to generate split-above key: %v", b.name, err) } else { beforeSplit := timeutil.Now() - err := b.db.AdminSplit(ctx, splitAbove, expire) + err := b.db.AdminSplit(ctx, splitAbove, expire, roachpb.AdminSplitRequest_INGESTION) b.currentStats.SplitWait += timeutil.Since(beforeSplit) if err != nil { log.Warningf(ctx, "%s failed to split-above: %v", b.name, err) @@ -536,7 +536,7 @@ func (b *SSTBatcher) doFlush(ctx context.Context, reason int) error { log.Warningf(ctx, "%s failed to generate split key: %v", b.name, err) } else { beforeSplit := timeutil.Now() - err := b.db.AdminSplit(ctx, splitAt, expire) + err := b.db.AdminSplit(ctx, splitAt, expire, roachpb.AdminSplitRequest_INGESTION) b.currentStats.SplitWait += timeutil.Since(beforeSplit) if err != nil { log.Warningf(ctx, "%s failed to split: %v", b.name, err) diff --git a/pkg/kv/bulk/sst_batcher_test.go b/pkg/kv/bulk/sst_batcher_test.go index 68436511eb0e..249df6f299a7 100644 --- a/pkg/kv/bulk/sst_batcher_test.go +++ b/pkg/kv/bulk/sst_batcher_test.go @@ -277,7 +277,12 @@ func runTestImport(t *testing.T, batchSizeValue int64) { } t.Logf("splitting at %s", key(split1)) - require.NoError(t, kvDB.AdminSplit(ctx, key(split1), hlc.MaxTimestamp /* expirationTime */)) + require.NoError(t, kvDB.AdminSplit( + ctx, + key(split1), + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) // We want to make sure our range-aware batching knows about one of our // splits to exercise that code path, but we also want to make sure we @@ -293,7 +298,12 @@ func runTestImport(t *testing.T, batchSizeValue int64) { } t.Logf("splitting at %s", key(split2)) - require.NoError(t, kvDB.AdminSplit(ctx, key(split2), hlc.MaxTimestamp /* expirationTime */)) + require.NoError(t, kvDB.AdminSplit( + ctx, + key(split2), + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) ts := hlc.Timestamp{WallTime: 100} mem := mon.NewUnlimitedMonitor(ctx, "lots", mon.MemoryResource, nil, nil, 0, nil) diff --git a/pkg/kv/db.go b/pkg/kv/db.go index 1d50f56ad0d1..055f762d53c7 100644 --- a/pkg/kv/db.go +++ b/pkg/kv/db.go @@ -590,10 +590,11 @@ func (db *DB) AdminSplit( ctx context.Context, splitKey interface{}, expirationTime hlc.Timestamp, + class roachpb.AdminSplitRequest_Class, predicateKeys ...roachpb.Key, ) error { b := &Batch{} - b.adminSplit(splitKey, expirationTime, predicateKeys) + b.adminSplit(splitKey, expirationTime, class, predicateKeys) return getOneErr(db.Run(ctx, b), b) } diff --git a/pkg/kv/kvclient/kvcoord/dist_sender_server_test.go b/pkg/kv/kvclient/kvcoord/dist_sender_server_test.go index 9aa37aefc10f..37e9cc738fa7 100644 --- a/pkg/kv/kvclient/kvcoord/dist_sender_server_test.go +++ b/pkg/kv/kvclient/kvcoord/dist_sender_server_test.go @@ -119,8 +119,13 @@ func TestRangeLookupWithOpenTransaction(t *testing.T) { // client. func setupMultipleRanges(ctx context.Context, db *kv.DB, splitAt ...string) error { // Split the keyspace at the given keys. - for _, key := range splitAt { - if err := db.AdminSplit(ctx, key /* splitKey */, hlc.MaxTimestamp /* expirationTime */); err != nil { + for _, splitKey := range splitAt { + if err := db.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { return err } } @@ -1173,7 +1178,12 @@ func TestMultiRangeScanDeleteRange(t *testing.T) { defer s.Stopper().Stop(ctx) tds := db.NonTransactionalSender() - if err := db.AdminSplit(ctx, "m", hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := db.AdminSplit( + ctx, + "m", /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } writes := []roachpb.Key{roachpb.Key("a"), roachpb.Key("z")} @@ -1325,8 +1335,13 @@ func TestMultiRangeScanWithPagination(t *testing.T) { defer s.Stopper().Stop(ctx) tds := db.NonTransactionalSender() - for _, sk := range tc.splitKeys { - if err := db.AdminSplit(ctx, sk, hlc.MaxTimestamp /* expirationTime */); err != nil { + for _, splitKey := range tc.splitKeys { + if err := db.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } @@ -1472,8 +1487,13 @@ func TestParallelSender(t *testing.T) { // Split into multiple ranges. splitKeys := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j"} - for _, key := range splitKeys { - if err := db.AdminSplit(context.Background(), key, hlc.MaxTimestamp /* expirationTime */); err != nil { + for _, splitKey := range splitKeys { + if err := db.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } @@ -1516,9 +1536,14 @@ func initReverseScanTestEnv(s serverutils.TestServerInterface, t *testing.T) *kv // Set up multiple ranges: // ["", "b"),["b", "e") ,["e", "g") and ["g", "\xff\xff"). - for _, key := range []string{"b", "e", "g"} { + for _, splitKey := range []string{"b", "e", "g"} { // Split the keyspace at the given key. - if err := db.AdminSplit(context.Background(), key, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := db.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } @@ -1618,8 +1643,13 @@ func TestBatchPutWithConcurrentSplit(t *testing.T) { // Split first using the default client and scan to make sure that // the range descriptor cache reflects the split. - for _, key := range []string{"b", "f"} { - if err := db.AdminSplit(context.Background(), key, hlc.MaxTimestamp /* expirationTime */); err != nil { + for _, splitKey := range []string{"b", "f"} { + if err := db.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } @@ -1676,7 +1706,12 @@ func TestReverseScanWithSplitAndMerge(t *testing.T) { // Case 1: An encounter with a range split. // Split the range ["b", "e") at "c". - if err := db.AdminSplit(context.Background(), "c", hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := db.AdminSplit( + context.Background(), + "c", /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -3538,7 +3573,12 @@ func BenchmarkReturnOnRangeBoundary(b *testing.B) { for r := 0; r < Ranges; r++ { rangeKey := string(rune('a' + r)) - require.NoError(b, db.AdminSplit(ctx, rangeKey, hlc.MaxTimestamp)) + require.NoError(b, db.AdminSplit( + ctx, + rangeKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) for k := 0; k < KeysPerRange; k++ { key := fmt.Sprintf("%s%d", rangeKey, k) diff --git a/pkg/kv/kvclient/kvcoord/split_test.go b/pkg/kv/kvclient/kvcoord/split_test.go index b48a189dfa38..f4832f6c5390 100644 --- a/pkg/kv/kvclient/kvcoord/split_test.go +++ b/pkg/kv/kvclient/kvcoord/split_test.go @@ -109,7 +109,12 @@ func TestRangeSplitMeta(t *testing.T) { for _, splitRKey := range splitKeys { splitKey := roachpb.Key(splitRKey) log.Infof(ctx, "starting split at key %q...", splitKey) - if err := s.DB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } log.Infof(ctx, "split at key %q complete", splitKey) @@ -156,7 +161,12 @@ func TestRangeSplitsWithConcurrentTxns(t *testing.T) { <-txnChannel } log.Infof(ctx, "starting split at key %q...", splitKey) - if pErr := s.DB.AdminSplit(context.Background(), splitKey, hlc.MaxTimestamp /* expirationTime */); pErr != nil { + if pErr := s.DB.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); pErr != nil { t.Error(pErr) } log.Infof(ctx, "split at key %q complete", splitKey) @@ -251,11 +261,21 @@ func TestRangeSplitsWithSameKeyTwice(t *testing.T) { splitKey := roachpb.Key("aa") log.Infof(ctx, "starting split at key %q...", splitKey) - if err := s.DB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } log.Infof(ctx, "split at key %q first time complete", splitKey) - if err := s.DB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } @@ -280,7 +300,12 @@ func TestRangeSplitsStickyBit(t *testing.T) { descKey := keys.RangeDescriptorKey(splitKey) // Splitting range. - if err := s.DB.AdminSplit(ctx, splitKey.AsRawKey(), hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey.AsRawKey(), + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -300,7 +325,12 @@ func TestRangeSplitsStickyBit(t *testing.T) { } // Splitting range. - if err := s.DB.AdminSplit(ctx, splitKey.AsRawKey(), hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey.AsRawKey(), + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -330,27 +360,27 @@ func TestSplitPredicates(t *testing.T) { expire := hlc.MaxTimestamp // Setup a known-span range [c, g) for some simple single predicate checks. - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("b"), expire)) - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("g"), expire)) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("b"), expire, roachpb.AdminSplitRequest_INGESTION)) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("g"), expire, roachpb.AdminSplitRequest_INGESTION)) // c is below split key f, and is in [b, g). - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("f"), expire, roachpb.Key("c"))) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("f"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("c"))) // e is above split key d, and is in [b, f). - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("d"), expire, roachpb.Key("e"))) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("d"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("e"))) // b is above split key c, and is in [b, d) although just barely. - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("c"), expire, roachpb.Key("b"))) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("c"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("b"))) // Setup another known span [g, n) and test rejections with it. - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("n"), expire)) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("n"), expire, roachpb.AdminSplitRequest_INGESTION)) // Reject split at h that wanted b to be in range [g, n). - require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.Key("b"))) + require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("b"))) // Reject split at h that wanted i, j, and z to be in range [g, n). - require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.Key("i"), roachpb.Key("j"), roachpb.Key("z"))) + require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("i"), roachpb.Key("j"), roachpb.Key("z"))) // Reject split at h that wanted i, j, and n to be in range [g, n). - require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.Key("i"), roachpb.Key("j"), roachpb.Key("n"))) + require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("i"), roachpb.Key("j"), roachpb.Key("n"))) // Reject split at h that wanted i, n and j to be in range [g, n). - require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.Key("i"), roachpb.Key("n"), roachpb.Key("j"))) + require.Error(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("i"), roachpb.Key("n"), roachpb.Key("j"))) // Allow split at h that wanted i, k and j to be in range [g, n). - require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.Key("i"), roachpb.Key("k"), roachpb.Key("j"))) + require.NoError(t, s.DB.AdminSplit(ctx, roachpb.Key("h"), expire, roachpb.AdminSplitRequest_INGESTION, roachpb.Key("i"), roachpb.Key("k"), roachpb.Key("j"))) } diff --git a/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go b/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go index b7aa70c26b1a..24cd33850b6e 100644 --- a/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go +++ b/pkg/kv/kvclient/kvcoord/txn_coord_sender_test.go @@ -163,7 +163,12 @@ func TestTxnCoordSenderHeartbeat(t *testing.T) { keyA := roachpb.Key("a") keyC := roachpb.Key("c") splitKey := roachpb.Key("b") - if err := s.DB.AdminSplit(ctx, splitKey /* splitKey */, hlc.MaxTimestamp /* expirationTimestamp */); err != nil { + if err := s.DB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -302,7 +307,12 @@ func TestDB_PrepareForRetryAfterHeartbeatFailure(t *testing.T) { keyA := roachpb.Key("a") keyC := roachpb.Key("c") splitKey := roachpb.Key("b") - require.NoError(t, s.DB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp)) + require.NoError(t, s.DB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) require.NoError(t, txn.Put(ctx, keyA, "1")) { @@ -945,7 +955,12 @@ func TestWTOBitTerminatedOnErrorResponses(t *testing.T) { defer s.Stopper().Stop(ctx) // Split to ensure batch requests get split through the distsender. - require.NoError(t, db.AdminSplit(ctx, keyB /* splitKey */, hlc.MaxTimestamp /* expirationTimestamp */)) + require.NoError(t, db.AdminSplit( + ctx, + keyB, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) // Write a key that the txn-al CPut below doesn't expect, failing the batch // request. @@ -2675,7 +2690,12 @@ func TestPutsInStagingTxn(t *testing.T) { }) defer s.Stopper().Stop(ctx) - require.NoError(t, db.AdminSplit(ctx, keyB /* splitKey */, hlc.MaxTimestamp /* expirationTimestamp */)) + require.NoError(t, db.AdminSplit( + ctx, + keyB, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) txn := db.NewTxn(ctx, "test") diff --git a/pkg/kv/kvclient/kvcoord/txn_test.go b/pkg/kv/kvclient/kvcoord/txn_test.go index 994c45abda60..35dfe4338afe 100644 --- a/pkg/kv/kvclient/kvcoord/txn_test.go +++ b/pkg/kv/kvclient/kvcoord/txn_test.go @@ -429,7 +429,12 @@ func TestTxnRepeatGetWithRangeSplit(t *testing.T) { } s.Manual.Advance(time.Second) // Split range by keyB. - if err := s.DB.AdminSplit(context.Background(), splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := s.DB.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } // Wait till split complete. diff --git a/pkg/kv/kvnemesis/applier.go b/pkg/kv/kvnemesis/applier.go index 7f000c1ff841..a6aae4df7642 100644 --- a/pkg/kv/kvnemesis/applier.go +++ b/pkg/kv/kvnemesis/applier.go @@ -86,7 +86,7 @@ func applyOp(ctx context.Context, env *Env, db *kv.DB, op *Operation) { *DeleteRangeOperation: applyClientOp(ctx, db, op, false /* inTxn */) case *SplitOperation: - err := db.AdminSplit(ctx, o.Key, hlc.MaxTimestamp) + err := db.AdminSplit(ctx, o.Key, hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION) o.Result = resultInit(ctx, err) case *MergeOperation: err := db.AdminMerge(ctx, o.Key) diff --git a/pkg/kv/kvserver/client_replica_test.go b/pkg/kv/kvserver/client_replica_test.go index d6e9a9325c29..13a9353e8baa 100644 --- a/pkg/kv/kvserver/client_replica_test.go +++ b/pkg/kv/kvserver/client_replica_test.go @@ -1020,7 +1020,7 @@ func TestTxnReadWithinUncertaintyIntervalAfterRangeMerge(t *testing.T) { require.NoError(t, tc.Server(0).DB().AdminMerge(ctx, keyA)) if alsoSplit { - require.NoError(t, tc.Server(0).DB().AdminSplit(ctx, keyC, hlc.MaxTimestamp)) + require.NoError(t, tc.Server(0).DB().AdminSplit(ctx, keyC, hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION)) } // Try and read the transaction from the context of a new transaction. This @@ -4007,7 +4007,7 @@ func TestChangeReplicasLeaveAtomicRacesWithMerge(t *testing.T) { err = db.AdminMerge(ctx, lhs) require.NoError(t, err) if resplit { - require.NoError(t, db.AdminSplit(ctx, rhs, hlc.Timestamp{WallTime: math.MaxInt64})) + require.NoError(t, db.AdminSplit(ctx, rhs, hlc.Timestamp{WallTime: math.MaxInt64}, roachpb.AdminSplitRequest_INGESTION)) err = tc.WaitForSplitAndInitialization(rhs) require.NoError(t, err) } diff --git a/pkg/kv/kvserver/consistency_queue_test.go b/pkg/kv/kvserver/consistency_queue_test.go index 8ecd2e96ddb4..549dca4306ca 100644 --- a/pkg/kv/kvserver/consistency_queue_test.go +++ b/pkg/kv/kvserver/consistency_queue_test.go @@ -480,7 +480,12 @@ func testConsistencyQueueRecomputeStatsImpl(t *testing.T, hadEstimates bool) { // Split off a range so that we get away from the timeseries writes, which // pollute the stats with ContainsEstimates=true. Note that the split clears // the right hand side (which is what we operate on) from that flag. - require.NoError(t, db0.AdminSplit(ctx, key, hlc.MaxTimestamp /* expirationTime */)) + require.NoError(t, db0.AdminSplit( + ctx, + key, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) delta := computeDelta(db0) diff --git a/pkg/kv/kvserver/intent_resolver_integration_test.go b/pkg/kv/kvserver/intent_resolver_integration_test.go index 9470fe925808..94d56e4fd385 100644 --- a/pkg/kv/kvserver/intent_resolver_integration_test.go +++ b/pkg/kv/kvserver/intent_resolver_integration_test.go @@ -318,7 +318,7 @@ func TestReliableIntentCleanup(t *testing.T) { // Split off 16 ranges by the first hex digit (4 bits) after prefix: // key\x00\x00 key\x00\x10 key\x00\x20 key\x00\x30 ... for i := 0; i < 16; i++ { - require.NoError(t, db.AdminSplit(ctx, append(prefix, byte(i<<4)), hlc.MaxTimestamp)) + require.NoError(t, db.AdminSplit(ctx, append(prefix, byte(i<<4)), hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION)) } require.NoError(t, tc.WaitForFullReplication()) diff --git a/pkg/kv/kvserver/range_log_test.go b/pkg/kv/kvserver/range_log_test.go index 3c327cae0c96..5882c0b96c2e 100644 --- a/pkg/kv/kvserver/range_log_test.go +++ b/pkg/kv/kvserver/range_log_test.go @@ -56,7 +56,12 @@ func TestLogSplits(t *testing.T) { initialSplits := countSplits() // Generate an explicit split event. - if err := kvDB.AdminSplit(ctx, "splitkey", hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + "splitkey", + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -178,10 +183,20 @@ func TestLogMerges(t *testing.T) { } // Create two ranges, then merge them. - if err := kvDB.AdminSplit(ctx, "a", hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + "a", /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } - if err := kvDB.AdminSplit(ctx, "b", hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + "b", /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } if err := kvDB.AdminMerge(ctx, "a"); err != nil { diff --git a/pkg/kv/kvserver/replica_command_test.go b/pkg/kv/kvserver/replica_command_test.go index 7b6e9e772059..847072763334 100644 --- a/pkg/kv/kvserver/replica_command_test.go +++ b/pkg/kv/kvserver/replica_command_test.go @@ -42,7 +42,12 @@ func TestRangeDescriptorUpdateProtoChangedAcrossVersions(t *testing.T) { defer s.Stopper().Stop(ctx) bKey := roachpb.Key("b") - if err := kvDB.AdminSplit(ctx, bKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + bKey, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -87,7 +92,12 @@ func TestRangeDescriptorUpdateProtoChangedAcrossVersions(t *testing.T) { // merges and replica changes, but they all go through updateRangeDescriptor // so it's unnecessary. cKey := roachpb.Key("c") - if err := kvDB.AdminSplit(ctx, cKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + cKey, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } } diff --git a/pkg/kv/kvserver/replica_test.go b/pkg/kv/kvserver/replica_test.go index b7df08e20c61..b3efcc171869 100644 --- a/pkg/kv/kvserver/replica_test.go +++ b/pkg/kv/kvserver/replica_test.go @@ -9611,7 +9611,12 @@ func TestErrorInRaftApplicationClearsIntents(t *testing.T) { defer s.Stopper().Stop(context.Background()) splitKey := roachpb.Key("b") - if err := kvDB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + ctx, + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } diff --git a/pkg/kv/kvserver/ts_maintenance_queue_test.go b/pkg/kv/kvserver/ts_maintenance_queue_test.go index 00b4beb90155..fbc50eb3c18d 100644 --- a/pkg/kv/kvserver/ts_maintenance_queue_test.go +++ b/pkg/kv/kvserver/ts_maintenance_queue_test.go @@ -279,7 +279,12 @@ func TestTimeSeriesMaintenanceQueueServer(t *testing.T) { // Force a range split in between near past and far past. This guarantees // that the pruning operation will issue a DeleteRange which spans ranges. - if err := db.AdminSplit(context.Background(), splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := db.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } diff --git a/pkg/kv/txn_external_test.go b/pkg/kv/txn_external_test.go index 8824aedc224a..9bd6af17f4fc 100644 --- a/pkg/kv/txn_external_test.go +++ b/pkg/kv/txn_external_test.go @@ -530,8 +530,8 @@ func TestRevScanAndGet(t *testing.T) { defer tc.Stopper().Stop(ctx) db := tc.Servers[0].DB() - require.NoError(t, db.AdminSplit(ctx, "b", hlc.MaxTimestamp)) - require.NoError(t, db.AdminSplit(ctx, "h", hlc.MaxTimestamp)) + require.NoError(t, db.AdminSplit(ctx, "b", hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION)) + require.NoError(t, db.AdminSplit(ctx, "h", hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION)) // Setup: // Ranges: [keyMin-------b) [b--------h) [h------keyMax) diff --git a/pkg/roachpb/api.proto b/pkg/roachpb/api.proto index e537ec4bc2c9..2a274cb1a1b6 100644 --- a/pkg/roachpb/api.proto +++ b/pkg/roachpb/api.proto @@ -812,6 +812,13 @@ message EndTxnResponse { // of hlc.Timestamp{} (I.E. The zero timestamp so they are always eligible for // automatic merging). message AdminSplitRequest { + + enum Class { + ARBITRARY = 0; + INGESTION = 1; + ORGANIZATION = 2; + } + RequestHeader header = 1 [(gogoproto.nullable) = false, (gogoproto.embed) = true]; bytes split_key = 2 [(gogoproto.casttype) = "Key"]; reserved 3; @@ -821,6 +828,8 @@ message AdminSplitRequest { // cause the split to be rejected. This can be used by a caller to effectively // send a "conditional split" request, i.e. a split if not already split. repeated bytes predicate_keys = 5 [(gogoproto.casttype) = "Key"]; + + Class class = 6; } // An AdminSplitResponse is the return value from the AdminSplit() diff --git a/pkg/rpc/auth_tenant.go b/pkg/rpc/auth_tenant.go index e6f1e5c8a4d9..e8ec04ad2a97 100644 --- a/pkg/rpc/auth_tenant.go +++ b/pkg/rpc/auth_tenant.go @@ -129,10 +129,10 @@ func (a tenantAuthorizer) authorize( // authBatch authorizes the provided tenant to invoke the Batch RPC with the // provided args. func (a tenantAuthorizer) authBatch(tenID roachpb.TenantID, args *roachpb.BatchRequest) error { - // Consult reqMethodAllowlist to determine whether each request in the batch + // Consult reqAllowed to determine whether each request in the batch // is permitted. If not, reject the entire batch. for _, ru := range args.Requests { - if !reqAllowed(ru.GetInner()) { + if !reqAllowed(ru.GetInner(), tenID) { return authErrorf("request [%s] not permitted", args.Summary()) } } @@ -149,35 +149,35 @@ func (a tenantAuthorizer) authBatch(tenID roachpb.TenantID, args *roachpb.BatchR return nil } -var reqMethodAllowlist = [...]bool{ - roachpb.Get: true, - roachpb.Put: true, - roachpb.ConditionalPut: true, - roachpb.Increment: true, - roachpb.Delete: true, - roachpb.DeleteRange: true, - roachpb.ClearRange: true, - roachpb.RevertRange: true, - roachpb.Scan: true, - roachpb.ReverseScan: true, - roachpb.EndTxn: true, - roachpb.AdminSplit: true, - roachpb.HeartbeatTxn: true, - roachpb.QueryTxn: true, - roachpb.QueryIntent: true, - roachpb.QueryLocks: true, - roachpb.InitPut: true, - roachpb.Export: true, - roachpb.AdminScatter: true, - roachpb.AddSSTable: true, - roachpb.Refresh: true, - roachpb.RefreshRange: true, - roachpb.IsSpanEmpty: true, -} - -func reqAllowed(r roachpb.Request) bool { - m := int(r.Method()) - return m < len(reqMethodAllowlist) && reqMethodAllowlist[m] +func reqAllowed(r roachpb.Request, tenID roachpb.TenantID) bool { + switch t := r.(type) { + case *roachpb.GetRequest, + *roachpb.PutRequest, + *roachpb.ConditionalPutRequest, + *roachpb.IncrementRequest, + *roachpb.DeleteRequest, + *roachpb.DeleteRangeRequest, + *roachpb.ClearRangeRequest, + *roachpb.RevertRangeRequest, + *roachpb.ScanRequest, + *roachpb.ReverseScanRequest, + *roachpb.EndTxnRequest, + *roachpb.HeartbeatTxnRequest, + *roachpb.QueryTxnRequest, + *roachpb.QueryIntentRequest, + *roachpb.QueryLocksRequest, + *roachpb.InitPutRequest, + *roachpb.ExportRequest, + *roachpb.AdminScatterRequest, + *roachpb.AddSSTableRequest, + *roachpb.RefreshRequest, + *roachpb.RefreshRangeRequest, + *roachpb.IsSpanEmptyRequest: + return true + case *roachpb.AdminSplitRequest: + return t.Class != roachpb.AdminSplitRequest_ARBITRARY || tenID == roachpb.SystemTenantID + } + return false } // authRangeLookup authorizes the provided tenant to invoke the RangeLookup RPC diff --git a/pkg/rpc/auth_test.go b/pkg/rpc/auth_test.go index 42bc04b10d03..bab6fcf926d9 100644 --- a/pkg/rpc/auth_test.go +++ b/pkg/rpc/auth_test.go @@ -179,7 +179,7 @@ func TestTenantAuthRequest(t *testing.T) { makeAdminSplitReq := func(key string) roachpb.Request { s := makeSpan(key) h := roachpb.RequestHeaderFromSpan(s) - return &roachpb.AdminSplitRequest{RequestHeader: h, SplitKey: s.Key} + return &roachpb.AdminSplitRequest{RequestHeader: h, SplitKey: s.Key, Class: roachpb.AdminSplitRequest_INGESTION} } makeAdminScatterReq := func(key string) roachpb.Request { s := makeSpan(key) diff --git a/pkg/server/intent_test.go b/pkg/server/intent_test.go index e25471844261..715ffe250d8d 100644 --- a/pkg/server/intent_test.go +++ b/pkg/server/intent_test.go @@ -117,7 +117,12 @@ func TestIntentResolution(t *testing.T) { Knobs: base.TestingKnobs{Store: &storeKnobs}}) defer s.Stopper().Stop(context.Background()) // Split the Range. This should not have any asynchronous intents. - if err := kvDB.AdminSplit(context.Background(), splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := kvDB.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } diff --git a/pkg/server/node_test.go b/pkg/server/node_test.go index 90c9eea65191..eae793f65be0 100644 --- a/pkg/server/node_test.go +++ b/pkg/server/node_test.go @@ -545,7 +545,12 @@ func TestNodeStatusWritten(t *testing.T) { // ======================================== // Split the range. - if err := ts.db.AdminSplit(context.Background(), splitKey, hlc.MaxTimestamp /* expirationTime */); err != nil { + if err := ts.db.AdminSplit( + context.Background(), + splitKey, + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + ); err != nil { t.Fatal(err) } @@ -719,7 +724,12 @@ func TestGetTenantWeights(t *testing.T) { // relevant store(s). const otherTenantID = 5 prefix := keys.MakeTenantPrefix(roachpb.MakeTenantID(otherTenantID)) - require.NoError(t, s.DB().AdminSplit(ctx, prefix, hlc.MaxTimestamp)) + require.NoError(t, s.DB().AdminSplit( + ctx, + prefix, /* splitKey */ + hlc.MaxTimestamp, /* expirationTime */ + roachpb.AdminSplitRequest_INGESTION, + )) // The range can have replicas on multiple stores, so wait for the split to // be applied everywhere. stores := s.GetStores().(*kvserver.Stores) diff --git a/pkg/sql/BUILD.bazel b/pkg/sql/BUILD.bazel index e82848c50e35..0285b63531cf 100644 --- a/pkg/sql/BUILD.bazel +++ b/pkg/sql/BUILD.bazel @@ -373,6 +373,7 @@ go_library( "//pkg/sql/lexbase", "//pkg/sql/memsize", "//pkg/sql/mutations", + "//pkg/sql/oppurpose", "//pkg/sql/opt", "//pkg/sql/opt/cat", "//pkg/sql/opt/constraint", diff --git a/pkg/sql/backfill.go b/pkg/sql/backfill.go index 597904989a5a..acd0b251dc58 100644 --- a/pkg/sql/backfill.go +++ b/pkg/sql/backfill.go @@ -37,6 +37,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc" "github.com/cockroachdb/cockroach/pkg/sql/execinfra" "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/row" @@ -2266,7 +2267,7 @@ func (sc *SchemaChanger) backfillIndexes( if sc.execCfg.Codec.ForSystemTenant() { expirationTime := sc.db.Clock().Now().Add(time.Hour.Nanoseconds(), 0) for _, span := range addingSpans { - if err := sc.db.AdminSplit(ctx, span.Key, expirationTime); err != nil { + if err := sc.db.AdminSplit(ctx, span.Key, expirationTime, oppurpose.SplitImport); err != nil { return err } } diff --git a/pkg/sql/importer/BUILD.bazel b/pkg/sql/importer/BUILD.bazel index 9492e0529ce3..47a04b5e749c 100644 --- a/pkg/sql/importer/BUILD.bazel +++ b/pkg/sql/importer/BUILD.bazel @@ -72,6 +72,7 @@ go_library( "//pkg/sql/faketreeeval", "//pkg/sql/gcjob", "//pkg/sql/lexbase", + "//pkg/sql/oppurpose", "//pkg/sql/opt/memo", "//pkg/sql/parser", "//pkg/sql/pgwire/pgcode", diff --git a/pkg/sql/importer/import_processor_planning.go b/pkg/sql/importer/import_processor_planning.go index 67fe23e6320f..7a9d61414d00 100644 --- a/pkg/sql/importer/import_processor_planning.go +++ b/pkg/sql/importer/import_processor_planning.go @@ -27,6 +27,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc" "github.com/cockroachdb/cockroach/pkg/sql/execinfrapb" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/physicalplan" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" "github.com/cockroachdb/cockroach/pkg/sql/types" @@ -352,7 +353,7 @@ func presplitTableBoundaries( // TODO(ajwerner): Consider passing in the wrapped descriptors. tblDesc := tabledesc.NewBuilder(tbl.Desc).BuildImmutableTable() for _, span := range tblDesc.AllIndexSpans(cfg.Codec) { - if err := cfg.DB.AdminSplit(ctx, span.Key, expirationTime); err != nil { + if err := cfg.DB.AdminSplit(ctx, span.Key, expirationTime, oppurpose.SplitImport); err != nil { return err } } diff --git a/pkg/sql/multitenant_admin_function_test.go b/pkg/sql/multitenant_admin_function_test.go index 181d945dce05..fc3feddebac2 100644 --- a/pkg/sql/multitenant_admin_function_test.go +++ b/pkg/sql/multitenant_admin_function_test.go @@ -12,44 +12,50 @@ package sql import ( "context" + gosql "database/sql" "testing" "github.com/cockroachdb/cockroach/pkg/base" "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" - "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" "github.com/cockroachdb/cockroach/pkg/util/errorutil" "github.com/cockroachdb/cockroach/pkg/util/leaktest" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/stretchr/testify/require" ) -func createSQLRunners( - t *testing.T, -) (systemTenant *sqlutils.SQLRunner, secondaryTenant *sqlutils.SQLRunner, cleanup func()) { - createTestClusterAndServer := func() (serverutils.TestClusterInterface, serverutils.TestServerInterface) { - systemTenantTestCluster := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{}) - return systemTenantTestCluster, systemTenantTestCluster.Server(0) +const createTable = "CREATE TABLE t(i int PRIMARY KEY);" + +func createTestServer(t *testing.T) (serverutils.TestServerInterface, func()) { + testCluster := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{}) + return testCluster.Server(0), func() { + testCluster.Stopper().Stop(context.Background()) } +} - systemTenantTestCluster, systemTenantTestServer := createTestClusterAndServer() - systemTenantDB := serverutils.OpenDBConn( +func createSystemTenantDB(t *testing.T, testServer serverutils.TestServerInterface) *gosql.DB { + return serverutils.OpenDBConn( t, - systemTenantTestServer.ServingSQLAddr(), + testServer.ServingSQLAddr(), "", /* useDatabase */ false, /* insecure */ - systemTenantTestServer.Stopper(), + testServer.Stopper(), ) +} - secondaryTenantTestCluster, secondaryTenantTestServer := createTestClusterAndServer() - _, secondaryTenantDB := serverutils.StartTenant( - t, secondaryTenantTestServer, base.TestTenantArgs{ +func createSecondaryTenantDB( + t *testing.T, testServer serverutils.TestServerInterface, allowSplitAndScatter bool, +) *gosql.DB { + _, db := serverutils.StartTenant( + t, testServer, base.TestTenantArgs{ + TestingKnobs: base.TestingKnobs{ + TenantTestingKnobs: &TenantTestingKnobs{ + AllowSplitAndScatter: allowSplitAndScatter, + }, + }, TenantID: serverutils.TestTenantID(), }, ) - return sqlutils.MakeSQLRunner(systemTenantDB), sqlutils.MakeSQLRunner(secondaryTenantDB), func() { - systemTenantTestCluster.Stopper().Stop(context.Background()) - secondaryTenantTestCluster.Stopper().Stop(context.Background()) - } + return db } func TestMultiTenantAdminFunction(t *testing.T) { @@ -57,153 +63,166 @@ func TestMultiTenantAdminFunction(t *testing.T) { defer log.Scope(t).Close(t) testCases := []struct { - desc string - query string + desc string + setup string + setups []string + query string + errorMessage string + allowSplitAndScatter bool }{ { - desc: "ALTER RANGE x RELOCATE LEASE", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER RANGE (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]) RELOCATE LEASE TO 1; -`, + desc: "ALTER RANGE x RELOCATE LEASE", + query: "ALTER RANGE (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]) RELOCATE LEASE TO 1;", }, { - desc: "ALTER RANGE RELOCATE LEASE", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER RANGE RELOCATE LEASE TO 1 FOR (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]); -`, + desc: "ALTER RANGE RELOCATE LEASE", + query: "ALTER RANGE RELOCATE LEASE TO 1 FOR (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]);", }, { - desc: "ALTER RANGE x RELOCATE VOTERS", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER RANGE (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]) RELOCATE VOTERS FROM 1 TO 1; -`, + desc: "ALTER RANGE x RELOCATE VOTERS", + query: "ALTER RANGE (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]) RELOCATE VOTERS FROM 1 TO 1;", }, { - desc: "ALTER RANGE RELOCATE VOTERS", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER RANGE RELOCATE VOTERS FROM 1 TO 1 FOR (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]); -`, + desc: "ALTER RANGE RELOCATE VOTERS", + query: "ALTER RANGE RELOCATE VOTERS FROM 1 TO 1 FOR (SELECT min(range_id) FROM [SHOW RANGES FROM TABLE t]);", }, { - desc: "ALTER TABLE x EXPERIMENTAL_RELOCATE LEASE", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t EXPERIMENTAL_RELOCATE LEASE SELECT 1, 1; -`, + desc: "ALTER TABLE x EXPERIMENTAL_RELOCATE LEASE", + query: "ALTER TABLE t EXPERIMENTAL_RELOCATE LEASE SELECT 1, 1;", }, { - desc: "ALTER TABLE x EXPERIMENTAL_RELOCATE VOTERS", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t EXPERIMENTAL_RELOCATE VOTERS SELECT ARRAY[1], 1; -`, + desc: "ALTER TABLE x EXPERIMENTAL_RELOCATE VOTERS", + query: "ALTER TABLE t EXPERIMENTAL_RELOCATE VOTERS SELECT ARRAY[1], 1;", }, { - desc: "ALTER TABLE x SPLIT AT", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t SPLIT AT VALUES (1); -`, + desc: "ALTER TABLE x SPLIT AT", + query: "ALTER TABLE t SPLIT AT VALUES (1);", + errorMessage: "request [1 AdmSplit] not permitted", }, { - desc: "ALTER INDEX x SPLIT AT", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -CREATE INDEX idx on t(i); -ALTER INDEX t@idx SPLIT AT VALUES (1); -`, + desc: "ALTER INDEX x SPLIT AT", + setup: "CREATE INDEX idx on t(i);", + query: "ALTER INDEX t@idx SPLIT AT VALUES (1);", + errorMessage: "request [1 AdmSplit] not permitted", }, { - // todo(ewall): Support 'ALTER TABLE x SPLIT AT' first. - desc: "ALTER TABLE x UNSPLIT AT", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t SPLIT AT VALUES (1); -ALTER TABLE t UNSPLIT AT VALUES (1); -`, + desc: "ALTER TABLE x UNSPLIT AT", + setup: "ALTER TABLE t SPLIT AT VALUES (1);", + query: "ALTER TABLE t UNSPLIT AT VALUES (1);", + allowSplitAndScatter: true, }, { - // todo(ewall): Support 'ALTER INDEX x SPLIT AT' first. desc: "ALTER INDEX x UNSPLIT AT", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -CREATE INDEX idx on t(i); -ALTER INDEX t@idx SPLIT AT VALUES (1); -ALTER INDEX t@idx UNSPLIT AT VALUES (1); -`, + // Multiple setup statements required because of https://github.com/cockroachdb/cockroach/issues/90535. + setups: []string{ + "CREATE INDEX idx on t(i);", + "ALTER INDEX t@idx SPLIT AT VALUES (1);", + }, + query: "ALTER INDEX t@idx UNSPLIT AT VALUES (1);", + allowSplitAndScatter: true, }, { - // todo(ewall): Support 'ALTER TABLE x SPLIT AT' first. - desc: "TRUNCATE TABLE x", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t SPLIT AT VALUES (1); -TRUNCATE TABLE t; -`, + desc: "ALTER TABLE x UNSPLIT ALL", + query: "ALTER TABLE t UNSPLIT ALL;", }, { - desc: "ALTER TABLE x UNSPLIT ALL", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t UNSPLIT ALL; -`, + desc: "ALTER INDEX x UNSPLIT ALL", + setup: "CREATE INDEX idx on t(i);", + query: "ALTER INDEX t@idx UNSPLIT ALL;", }, { - desc: "ALTER INDEX x UNSPLIT ALL", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -CREATE INDEX idx on t(i); -ALTER INDEX t@idx UNSPLIT ALL; -`, + desc: "ALTER TABLE x SCATTER", + query: "ALTER TABLE t SCATTER;", }, { - desc: "ALTER TABLE x SCATTER", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t SCATTER; -`, + desc: "ALTER INDEX x SCATTER", + setup: "CREATE INDEX idx on t(i);", + query: "ALTER INDEX t@idx SCATTER;", }, { - desc: "ALTER INDEX x SCATTER", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -CREATE INDEX idx on t(i); -ALTER INDEX t@idx SCATTER; -`, - }, - { - desc: "CONFIGURE ZONE", - query: ` -CREATE TABLE t(i int PRIMARY KEY); -ALTER TABLE t CONFIGURE ZONE DISCARD; -`, + desc: "CONFIGURE ZONE", + query: "ALTER TABLE t CONFIGURE ZONE DISCARD;", }, } + ctx := context.Background() + for _, tc := range testCases { t.Run(tc.desc, func(t *testing.T) { - systemTenant, secondaryTenant, cleanup := createSQLRunners(t) - defer cleanup() - ctx := context.Background() - query := tc.query + execQueries := func(db *gosql.DB) error { + setups := tc.setups + setup := tc.setup + if setup != "" { + setups = []string{setup} + } + setups = append([]string{createTable}, setups...) + for _, setup := range setups { + _, err := db.ExecContext(ctx, setup) + require.NoError(t, err, setup) + } + _, err := db.ExecContext(ctx, tc.query) + return err + } // Test system tenant. - { - _, err := systemTenant.DB.ExecContext(ctx, query) + func() { + testServer, cleanup := createTestServer(t) + defer cleanup() + db := createSystemTenantDB(t, testServer) + err := execQueries(db) require.NoError(t, err) - } + }() // Test secondary tenant. - { - _, err := secondaryTenant.DB.ExecContext(ctx, query) + func() { + testServer, cleanup := createTestServer(t) + defer cleanup() + db := createSecondaryTenantDB(t, testServer, tc.allowSplitAndScatter) + err := execQueries(db) require.Error(t, err) - require.Contains(t, err.Error(), errorutil.UnsupportedWithMultiTenancyMessage) - } + errorMessage := tc.errorMessage + if errorMessage == "" { + errorMessage = errorutil.UnsupportedWithMultiTenancyMessage + } + require.Contains(t, err.Error(), errorMessage) + }() }) } } + +func TestTruncateTable(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx := context.Background() + execQueries := func(db *gosql.DB) error { + _, err := db.ExecContext(ctx, createTable) + require.NoError(t, err) + _, err = db.ExecContext(ctx, "ALTER TABLE t SPLIT AT VALUES (1);") + require.NoError(t, err) + _, err = db.ExecContext(ctx, "TRUNCATE TABLE t;") + require.NoError(t, err) + _, err = db.QueryContext(ctx, "SHOW RANGES FROM TABLE t;") + // TODO(ewall): Verify splits are not retained after `SHOW RANGES` is supported by secondary tenants. + return err + } + + // Test system tenant. + func() { + testServer, cleanup := createTestServer(t) + defer cleanup() + db := createSystemTenantDB(t, testServer) + err := execQueries(db) + require.NoError(t, err) + }() + + // Test secondary tenant. + func() { + testServer, cleanup := createTestServer(t) + defer cleanup() + db := createSecondaryTenantDB(t, testServer, true /* allowSplitAndScatter */) + err := execQueries(db) + require.Error(t, err) + }() +} diff --git a/pkg/sql/oppurpose/BUILD.bazel b/pkg/sql/oppurpose/BUILD.bazel new file mode 100644 index 000000000000..dfc4ad6e828e --- /dev/null +++ b/pkg/sql/oppurpose/BUILD.bazel @@ -0,0 +1,12 @@ +load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data") +load("@io_bazel_rules_go//go:def.bzl", "go_library") + +go_library( + name = "oppurpose", + srcs = ["op_purpose.go"], + importpath = "github.com/cockroachdb/cockroach/pkg/sql/oppurpose", + visibility = ["//visibility:public"], + deps = ["//pkg/roachpb"], +) + +get_x_data(name = "get_x_data") diff --git a/pkg/sql/oppurpose/op_purpose.go b/pkg/sql/oppurpose/op_purpose.go new file mode 100644 index 000000000000..acb1025d2fe9 --- /dev/null +++ b/pkg/sql/oppurpose/op_purpose.go @@ -0,0 +1,28 @@ +// Copyright 2022 The Cockroach Authors. +// +// Use of this software is governed by the Business Source License +// included in the file licenses/BSL.txt. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0, included in the file +// licenses/APL.txt. + +package oppurpose + +import "github.com/cockroachdb/cockroach/pkg/roachpb" + +const ( + // SplitBackup is a split caused by a backup. + SplitBackup = roachpb.AdminSplitRequest_INGESTION + // SplitImport is a split caused by an import. + SplitImport = roachpb.AdminSplitRequest_INGESTION + // SplitSchema is a split caused by a schema change. + SplitSchema = roachpb.AdminSplitRequest_INGESTION + // SplitCreateTenant is a split caused by tenant creation. + SplitCreateTenant = roachpb.AdminSplitRequest_INGESTION + // SplitManual is a split caused by a manual action. + SplitManual = roachpb.AdminSplitRequest_ARBITRARY + // SplitManualTest is a split caused by a manual action test. + SplitManualTest = roachpb.AdminSplitRequest_INGESTION +) diff --git a/pkg/sql/opt_exec_factory.go b/pkg/sql/opt_exec_factory.go index ae9e79501fd3..a7dcab0734ee 100644 --- a/pkg/sql/opt_exec_factory.go +++ b/pkg/sql/opt_exec_factory.go @@ -1944,21 +1944,18 @@ func (ef *execFactory) ConstructAlterTableSplit( return nil, err } - knobs := ef.planner.ExecCfg().TenantTestingKnobs - if !(knobs != nil && knobs.AllowSplitAndScatter) && !ef.planner.ExecCfg().Codec.ForSystemTenant() { - return nil, errorutil.UnsupportedWithMultiTenancy(54254) - } - expirationTime, err := parseExpirationTime(ef.ctx, ef.planner.EvalContext(), expiration) if err != nil { return nil, err } + knobs := ef.planner.ExecCfg().TenantTestingKnobs return &splitNode{ - tableDesc: index.Table().(*optTable).desc, - index: index.(*optIndex).idx, - rows: input.(planNode), - expirationTime: expirationTime, + tableDesc: index.Table().(*optTable).desc, + index: index.(*optIndex).idx, + rows: input.(planNode), + expirationTime: expirationTime, + testAllowSplitAndScatter: knobs != nil && knobs.AllowSplitAndScatter, }, nil } diff --git a/pkg/sql/schema_changer.go b/pkg/sql/schema_changer.go index 8251bc12c639..451a9fac3120 100644 --- a/pkg/sql/schema_changer.go +++ b/pkg/sql/schema_changer.go @@ -44,6 +44,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/faketreeeval" "github.com/cockroachdb/cockroach/pkg/sql/flowinfra" "github.com/cockroachdb/cockroach/pkg/sql/gcjob/gcjobnotifier" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/parser" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" @@ -3173,7 +3174,7 @@ func (sc *SchemaChanger) preSplitHashShardedIndexRanges(ctx context.Context) err func splitAndScatter( ctx context.Context, db *kv.DB, key roachpb.Key, expirationTime hlc.Timestamp, ) error { - if err := db.AdminSplit(ctx, key, expirationTime); err != nil { + if err := db.AdminSplit(ctx, key, expirationTime, oppurpose.SplitSchema); err != nil { return err } _, err := db.AdminScatter(ctx, key, 0 /* maxSize */) diff --git a/pkg/sql/schemachanger/scdeps/BUILD.bazel b/pkg/sql/schemachanger/scdeps/BUILD.bazel index dd00365fc9ff..c7b3df1470a6 100644 --- a/pkg/sql/schemachanger/scdeps/BUILD.bazel +++ b/pkg/sql/schemachanger/scdeps/BUILD.bazel @@ -31,6 +31,7 @@ go_library( "//pkg/sql/catalog/descs", "//pkg/sql/catalog/resolver", "//pkg/sql/descmetadata", + "//pkg/sql/oppurpose", "//pkg/sql/rowenc", "//pkg/sql/schemachanger/scbuild", "//pkg/sql/schemachanger/scexec", diff --git a/pkg/sql/schemachanger/scdeps/exec_deps.go b/pkg/sql/schemachanger/scdeps/exec_deps.go index b578e3290d27..79e9396924b5 100644 --- a/pkg/sql/schemachanger/scdeps/exec_deps.go +++ b/pkg/sql/schemachanger/scdeps/exec_deps.go @@ -29,6 +29,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descs" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec" "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scexec/scmutationexec" "github.com/cockroachdb/cockroach/pkg/sql/sem/tree" @@ -323,7 +324,7 @@ func (d *txnDeps) MaybeSplitIndexSpans( span := table.IndexSpan(d.codec, indexToBackfill.GetID()) const backfillSplitExpiration = time.Hour expirationTime := d.txn.DB().Clock().Now().Add(backfillSplitExpiration.Nanoseconds(), 0) - return d.txn.DB().AdminSplit(ctx, span.Key, expirationTime) + return d.txn.DB().AdminSplit(ctx, span.Key, expirationTime, oppurpose.SplitSchema) } // GetResumeSpans implements the scexec.BackfillerTracker interface. diff --git a/pkg/sql/split.go b/pkg/sql/split.go index 6a7cd8db60df..87f584e73a6b 100644 --- a/pkg/sql/split.go +++ b/pkg/sql/split.go @@ -16,6 +16,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/keys" "github.com/cockroachdb/cockroach/pkg/sql/catalog" "github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" @@ -29,11 +30,12 @@ import ( type splitNode struct { optColumnsSlot - tableDesc catalog.TableDescriptor - index catalog.Index - rows planNode - run splitRun - expirationTime hlc.Timestamp + tableDesc catalog.TableDescriptor + index catalog.Index + rows planNode + run splitRun + expirationTime hlc.Timestamp + testAllowSplitAndScatter bool } // splitRun contains the run-time state of splitNode during local execution. @@ -60,7 +62,11 @@ func (n *splitNode) Next(params runParams) (bool, error) { return false, err } - if err := params.ExecCfg().DB.AdminSplit(params.ctx, rowKey, n.expirationTime); err != nil { + purpose := oppurpose.SplitManual + if n.testAllowSplitAndScatter { + purpose = oppurpose.SplitManualTest + } + if err := params.ExecCfg().DB.AdminSplit(params.ctx, rowKey, n.expirationTime, purpose); err != nil { return false, err } diff --git a/pkg/sql/tenant.go b/pkg/sql/tenant.go index 909709610d0a..a6d2f57ccdd4 100644 --- a/pkg/sql/tenant.go +++ b/pkg/sql/tenant.go @@ -29,6 +29,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/catalog/bootstrap" "github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb" "github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema" + "github.com/cockroachdb/cockroach/pkg/sql/oppurpose" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" "github.com/cockroachdb/cockroach/pkg/sql/rowenc" @@ -306,7 +307,7 @@ func (p *planner) CreateTenant(ctx context.Context, tenID uint64, name string) e // the span configs infrastructure, in `system.span_configurations`. expTime := p.ExecCfg().Clock.Now().Add(time.Hour.Nanoseconds(), 0) for _, key := range splits { - if err := p.ExecCfg().DB.AdminSplit(ctx, key, expTime); err != nil { + if err := p.ExecCfg().DB.AdminSplit(ctx, key, expTime, oppurpose.SplitCreateTenant); err != nil { return err } } diff --git a/pkg/sql/unsplit_range_test.go b/pkg/sql/unsplit_range_test.go index a1fb5267cd5c..f7bd9b0f1ef4 100644 --- a/pkg/sql/unsplit_range_test.go +++ b/pkg/sql/unsplit_range_test.go @@ -124,7 +124,7 @@ func splitLastRangeInSpan( if err != nil { t.Fatal(err) } - if err := kvDB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp); err != nil { + if err := kvDB.AdminSplit(ctx, splitKey, hlc.MaxTimestamp, roachpb.AdminSplitRequest_INGESTION); err != nil { t.Fatal(err) }