From 6317bd81b0f44e4b198c5b8bfb51d2707e011f1f Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Tue, 13 Jun 2023 15:30:43 -0700 Subject: [PATCH 1/2] server: use server's codec instead of TestingSQLCodec in some tests Epic: None Release note: None --- pkg/server/BUILD.bazel | 1 - pkg/server/decommission_test.go | 7 +++---- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/pkg/server/BUILD.bazel b/pkg/server/BUILD.bazel index 882b99a4c6b1..f75c50a82230 100644 --- a/pkg/server/BUILD.bazel +++ b/pkg/server/BUILD.bazel @@ -521,7 +521,6 @@ go_test( "//pkg/testutils", "//pkg/testutils/datapathutils", "//pkg/testutils/diagutils", - "//pkg/testutils/keysutils", "//pkg/testutils/serverutils", "//pkg/testutils/skip", "//pkg/testutils/sqlutils", diff --git a/pkg/server/decommission_test.go b/pkg/server/decommission_test.go index 2d39fd7101bd..c4bc293a58ee 100644 --- a/pkg/server/decommission_test.go +++ b/pkg/server/decommission_test.go @@ -21,7 +21,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/security/username" "github.com/cockroachdb/cockroach/pkg/testutils" - "github.com/cockroachdb/cockroach/pkg/testutils/keysutils" "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" "github.com/cockroachdb/cockroach/pkg/testutils/skip" "github.com/cockroachdb/cockroach/pkg/util/leaktest" @@ -129,8 +128,8 @@ func TestDecommissionPreCheckEvaluation(t *testing.T) { require.NoError(t, err) tblBID, err := firstSvr.admin.queryTableID(ctx, username.RootUserName(), "test", "tblB") require.NoError(t, err) - startKeyTblA := keysutils.TestingSQLCodec.TablePrefix(uint32(tblAID)) - startKeyTblB := keysutils.TestingSQLCodec.TablePrefix(uint32(tblBID)) + startKeyTblA := firstSvr.Codec().TablePrefix(uint32(tblAID)) + startKeyTblB := firstSvr.Codec().TablePrefix(uint32(tblBID)) // Split off ranges for tblA and tblB. _, rDescA, err := firstSvr.SplitRange(startKeyTblA) @@ -234,7 +233,7 @@ func TestDecommissionPreCheckOddToEven(t *testing.T) { runQueries(alterQueries...) tblAID, err := firstSvr.admin.queryTableID(ctx, username.RootUserName(), "test", "tblA") require.NoError(t, err) - startKeyTblA := keysutils.TestingSQLCodec.TablePrefix(uint32(tblAID)) + startKeyTblA := firstSvr.Codec().TablePrefix(uint32(tblAID)) // Split off range for tblA. _, rDescA, err := firstSvr.SplitRange(startKeyTblA) From 2485688b2b79f1ac02116c7e81987132f4647d6b Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Tue, 13 Jun 2023 19:22:33 -0700 Subject: [PATCH 2/2] keysutils: extend MakePrettyScannerForNamedTables to support tenants This commit extends `MakePrettyScannerForNamedTables` to support keys from the secondary tenants' key space, thus, removing the last usage of `TestingSQLCodec`. The way this is achieved is similar to how we overwrite the tableKeyParse function for `/Table` prefix in the dictionary - in this commit, we overwrite the tenantKeyParse function for `/Tenant` prefix. Release note: None --- pkg/cli/flags_util.go | 2 +- pkg/keys/printer.go | 51 +++++++------ pkg/keys/printer_test.go | 2 +- .../loqrecovery/loqrecoverypb/recovery.go | 2 +- .../kvserver/loqrecovery/recovery_env_test.go | 2 +- .../reports/constraint_stats_report_test.go | 2 +- pkg/kv/kvserver/reports/reporter_test.go | 3 +- pkg/testutils/keysutils/pretty_scanner.go | 36 ++++++---- .../keysutils/pretty_scanner_test.go | 72 ++++++++++--------- pkg/util/keysutil/keys.go | 41 ++++++++--- 10 files changed, 128 insertions(+), 85 deletions(-) diff --git a/pkg/cli/flags_util.go b/pkg/cli/flags_util.go index f189f31904c9..6092971d32d1 100644 --- a/pkg/cli/flags_util.go +++ b/pkg/cli/flags_util.go @@ -168,7 +168,7 @@ func (k *mvccKey) Set(value string) error { } *k = mvccKey(storage.MakeMVCCMetadataKey(roachpb.Key(unquoted))) case human: - scanner := keysutil.MakePrettyScanner(nil /* tableParser */) + scanner := keysutil.MakePrettyScanner(nil /* tableParser */, nil /* tenantParser */) key, err := scanner.Scan(keyStr) if err != nil { return err diff --git a/pkg/keys/printer.go b/pkg/keys/printer.go index e22edc7792c4..4a320eefa14e 100644 --- a/pkg/keys/printer.go +++ b/pkg/keys/printer.go @@ -233,26 +233,37 @@ func localStoreKeyParse(input string) (remainder string, output roachpb.Key) { const strTable = "/Table/" -func tenantKeyParse(input string) (remainder string, output roachpb.Key) { - input = mustShiftSlash(input) - slashPos := strings.Index(input, "/") - if slashPos < 0 { - slashPos = len(input) - } - remainder = input[slashPos:] // `/something/else` -> `/else` - tenantIDStr := input[:slashPos] - tenantID, err := strconv.ParseUint(tenantIDStr, 10, 64) - if err != nil { - panic(&ErrUglifyUnsupported{err}) - } - output = MakeTenantPrefix(roachpb.MustMakeTenantID(tenantID)) - if strings.HasPrefix(remainder, strTable) { - var indexKey roachpb.Key - remainder = remainder[len(strTable)-1:] - remainder, indexKey = tableKeyParse(remainder) - output = append(output, indexKey...) +// GetTenantKeyParseFn returns a function that parses the relevant prefix of the +// tenant data into a roachpb.Key, returning the remainder and the key +// corresponding to the consumed prefix of 'input'. It is expected that the +// '/Tenant' prefix has already been removed (i.e. the input is assumed to be of +// the form '//...'). If the input is of the form +// '//Table//...', then passed-in tableKeyParseFn function is +// invoked on the '//...' part. +func GetTenantKeyParseFn( + tableKeyParseFn func(string) (string, roachpb.Key), +) func(input string) (remainder string, output roachpb.Key) { + return func(input string) (remainder string, output roachpb.Key) { + input = mustShiftSlash(input) + slashPos := strings.Index(input, "/") + if slashPos < 0 { + slashPos = len(input) + } + remainder = input[slashPos:] // `/something/else` -> `/else` + tenantIDStr := input[:slashPos] + tenantID, err := strconv.ParseUint(tenantIDStr, 10, 64) + if err != nil { + panic(&ErrUglifyUnsupported{err}) + } + output = MakeTenantPrefix(roachpb.MustMakeTenantID(tenantID)) + if strings.HasPrefix(remainder, strTable) { + var indexKey roachpb.Key + remainder = remainder[len(strTable)-1:] + remainder, indexKey = tableKeyParseFn(remainder) + output = append(output, indexKey...) + } + return remainder, output } - return remainder, output } func tableKeyParse(input string) (remainder string, output roachpb.Key) { @@ -830,7 +841,7 @@ func init() { {Name: "", prefix: nil, ppFunc: decodeKeyPrint, PSFunc: tableKeyParse, sfFunc: formatTableKey}, }}, {Name: "/Tenant", start: TenantTableDataMin, end: TenantTableDataMax, Entries: []DictEntry{ - {Name: "", prefix: nil, ppFunc: tenantKeyPrint, PSFunc: tenantKeyParse, sfFunc: formatTenantKey}, + {Name: "", prefix: nil, ppFunc: tenantKeyPrint, PSFunc: GetTenantKeyParseFn(tableKeyParse), sfFunc: formatTenantKey}, }}, } } diff --git a/pkg/keys/printer_test.go b/pkg/keys/printer_test.go index c73183ac9dd8..3ca6eee686e4 100644 --- a/pkg/keys/printer_test.go +++ b/pkg/keys/printer_test.go @@ -483,7 +483,7 @@ exp: %s t.Errorf("%d: from string expected %s, got %s", i, exp, test.key.String()) } - scanner := keysutil.MakePrettyScanner(nil /* tableParser */) + scanner := keysutil.MakePrettyScanner(nil /* tableParser */, nil /* tenantParser */) parsed, err := scanner.Scan(keyInfo) if err != nil { if !errors.HasType(err, (*keys.ErrUglifyUnsupported)(nil)) { diff --git a/pkg/kv/kvserver/loqrecovery/loqrecoverypb/recovery.go b/pkg/kv/kvserver/loqrecovery/loqrecoverypb/recovery.go index f7fa827c586d..5f6e18085530 100644 --- a/pkg/kv/kvserver/loqrecovery/loqrecoverypb/recovery.go +++ b/pkg/kv/kvserver/loqrecovery/loqrecoverypb/recovery.go @@ -37,7 +37,7 @@ func (r *RecoveryKey) UnmarshalYAML(fn func(interface{}) error) error { if err := fn(&pretty); err != nil { return err } - scanner := keysutil.MakePrettyScanner(nil /* tableParser */) + scanner := keysutil.MakePrettyScanner(nil /* tableParser */, nil /* tenantParser */) key, err := scanner.Scan(pretty) if err != nil { return errors.Wrapf(err, "failed to parse key %s", pretty) diff --git a/pkg/kv/kvserver/loqrecovery/recovery_env_test.go b/pkg/kv/kvserver/loqrecovery/recovery_env_test.go index ff44a23cf3c7..a94ba9902c8a 100644 --- a/pkg/kv/kvserver/loqrecovery/recovery_env_test.go +++ b/pkg/kv/kvserver/loqrecovery/recovery_env_test.go @@ -504,7 +504,7 @@ func raftLogFromPendingDescriptorUpdate( } func parsePrettyKey(t *testing.T, pretty string) roachpb.RKey { - scanner := keysutil.MakePrettyScanner(nil /* tableParser */) + scanner := keysutil.MakePrettyScanner(nil /* tableParser */, nil /* tenantParser */) key, err := scanner.Scan(pretty) if err != nil { t.Fatalf("failed to parse key %s: %v", pretty, err) diff --git a/pkg/kv/kvserver/reports/constraint_stats_report_test.go b/pkg/kv/kvserver/reports/constraint_stats_report_test.go index 0fc3f51266e5..17d300b15219 100644 --- a/pkg/kv/kvserver/reports/constraint_stats_report_test.go +++ b/pkg/kv/kvserver/reports/constraint_stats_report_test.go @@ -985,7 +985,7 @@ func compileTestCase(tc baseReportTestCase) (compiledTestCase, error) { allStores = append(allStores, sds...) } - keyScanner := keysutils.MakePrettyScannerForNamedTables(tableToID, idxToID) + keyScanner := keysutils.MakePrettyScannerForNamedTables(roachpb.SystemTenantID, tableToID, idxToID) ranges, err := processSplits(keyScanner, tc.splits, allStores) if err != nil { return compiledTestCase{}, err diff --git a/pkg/kv/kvserver/reports/reporter_test.go b/pkg/kv/kvserver/reports/reporter_test.go index e22efd7c8cfd..65dfb3676e07 100644 --- a/pkg/kv/kvserver/reports/reporter_test.go +++ b/pkg/kv/kvserver/reports/reporter_test.go @@ -565,7 +565,8 @@ func TestZoneChecker(t *testing.T) { splits[i].key = ranges[i].split } keyScanner := keysutils.MakePrettyScannerForNamedTables( - map[string]int{"t1": t1ID} /* tableNameToID */, nil /* idxNameToID */) + roachpb.SystemTenantID, map[string]int{"t1": t1ID} /* tableNameToID */, nil, /* idxNameToID */ + ) rngs, err := processSplits(keyScanner, splits, nil /* stores */) require.NoError(t, err) diff --git a/pkg/testutils/keysutils/pretty_scanner.go b/pkg/testutils/keysutils/pretty_scanner.go index 59f3f67ccce1..b22b7297068a 100644 --- a/pkg/testutils/keysutils/pretty_scanner.go +++ b/pkg/testutils/keysutils/pretty_scanner.go @@ -22,18 +22,25 @@ import ( "github.com/cockroachdb/cockroach/pkg/util/keysutil" ) -// MakePrettyScannerForNamedTables create a PrettyScanner that, beside what the +// MakePrettyScannerForNamedTables creates a PrettyScanner that, beside what the // PrettyScanner is generally able to decode, can also decode keys of the form // "///1/2/3/..." using supplied maps from names to ids. -// -// TODO(nvanbenschoten): support tenant SQL keys. +// If the passed-in tenantID is not of the system tenant, then it handles keys +// of the form "/Tenant//Table/
//1/2/3/...". func MakePrettyScannerForNamedTables( - tableNameToID map[string]int, idxNameToID map[string]int, + tenantID roachpb.TenantID, tableNameToID map[string]int, idxNameToID map[string]int, ) keysutil.PrettyScanner { - return keysutil.MakePrettyScanner(func(input string) (string, roachpb.Key) { - remainder, k := parseTableKeysAsAscendingInts(input, tableNameToID, idxNameToID) - return remainder, k - }) + var tableParser, tenantParser keys.KeyParserFunc + if tenantID == roachpb.SystemTenantID { + tableParser = func(input string) (string, roachpb.Key) { + return parseTableKeysAsAscendingInts(input, tableNameToID, idxNameToID) + } + } else { + tenantParser = keys.GetTenantKeyParseFn(func(input string) (string, roachpb.Key) { + return parseTableKeysAsAscendingInts(input, tableNameToID, idxNameToID) + }) + } + return keysutil.MakePrettyScanner(tableParser, tenantParser) } // parseTableKeysAsAscendingInts takes a pretty-printed key segment like @@ -49,7 +56,8 @@ func MakePrettyScannerForNamedTables( // // Notice that the input is not expected to have the "/Table" prefix (which is // generated by pretty-printing keys). That prefix is assumed to have been -// consumed before this function is invoked (by the PrettyScanner). +// consumed before this function is invoked (by the PrettyScanner or +// GetTenantKeyParseFn). // // The "/..." part is returned as the remainder (the first return value). func parseTableKeysAsAscendingInts( @@ -67,7 +75,10 @@ func parseTableKeysAsAscendingInts( if !ok { panic(fmt.Sprintf("unknown table: %s", tableName)) } - output := TestingSQLCodec.TablePrefix(uint32(tableID)) + // We assume that the tenant prefix (if there was any) was already removed + // and included into the output by the caller, so we use the system codec + // here. + output := keys.SystemSQLCodec.TablePrefix(uint32(tableID)) if remainder == "" { return "", output } @@ -167,8 +178,3 @@ func parseAscendingIntIndexKey(input string) (string, roachpb.Key) { } return remainder, key } - -// TestingSQLCodec is a SQL key codec. It is equivalent to keys.SystemSQLCodec, but -// should be used when it is unclear which tenant should be referenced by the -// surrounding context. -var TestingSQLCodec = keys.MakeSQLCodec(roachpb.SystemTenantID) diff --git a/pkg/testutils/keysutils/pretty_scanner_test.go b/pkg/testutils/keysutils/pretty_scanner_test.go index 72ea01b91cec..a30a33b6bbe3 100644 --- a/pkg/testutils/keysutils/pretty_scanner_test.go +++ b/pkg/testutils/keysutils/pretty_scanner_test.go @@ -23,25 +23,25 @@ import ( func TestPrettyScanner(t *testing.T) { tests := []struct { prettyKey string - expKey func() roachpb.Key + expKey func(roachpb.TenantID) roachpb.Key expRemainder string }{ { prettyKey: "/Table/t1", - expKey: func() roachpb.Key { - return keys.SystemSQLCodec.TablePrefix(50) + expKey: func(tenantID roachpb.TenantID) roachpb.Key { + return keys.MakeSQLCodec(tenantID).TablePrefix(50) }, }, { prettyKey: "/Table/t1/pk", - expKey: func() roachpb.Key { - return keys.SystemSQLCodec.IndexPrefix(50, 1) + expKey: func(tenantID roachpb.TenantID) roachpb.Key { + return keys.MakeSQLCodec(tenantID).IndexPrefix(50, 1) }, }, { prettyKey: "/Table/t1/pk/1/2/3", - expKey: func() roachpb.Key { - k := keys.SystemSQLCodec.IndexPrefix(50, 1) + expKey: func(tenantID roachpb.TenantID) roachpb.Key { + k := keys.MakeSQLCodec(tenantID).IndexPrefix(50, 1) k = encoding.EncodeVarintAscending(k, 1) k = encoding.EncodeVarintAscending(k, 2) k = encoding.EncodeVarintAscending(k, 3) @@ -55,8 +55,8 @@ func TestPrettyScanner(t *testing.T) { }, { prettyKey: "/Table/t1/idx1/1/2/3", - expKey: func() roachpb.Key { - k := keys.SystemSQLCodec.IndexPrefix(50, 5) + expKey: func(tenantID roachpb.TenantID) roachpb.Key { + k := keys.MakeSQLCodec(tenantID).IndexPrefix(50, 5) k = encoding.EncodeVarintAscending(k, 1) k = encoding.EncodeVarintAscending(k, 2) k = encoding.EncodeVarintAscending(k, 3) @@ -67,32 +67,38 @@ func TestPrettyScanner(t *testing.T) { tableToID := map[string]int{"t1": 50} idxToID := map[string]int{"t1.idx1": 5} - scanner := MakePrettyScannerForNamedTables(tableToID, idxToID) - for _, test := range tests { - t.Run(test.prettyKey, func(t *testing.T) { - k, err := scanner.Scan(test.prettyKey) - if err != nil { - if test.expRemainder != "" { - if testutils.IsError(err, fmt.Sprintf("can't parse\"%s\"", test.expRemainder)) { - t.Fatalf("expected remainder: %s, got err: %s", test.expRemainder, err) + for _, tenantID := range []roachpb.TenantID{roachpb.SystemTenantID, roachpb.MustMakeTenantID(42)} { + scanner := MakePrettyScannerForNamedTables(tenantID, tableToID, idxToID) + for _, test := range tests { + prettyKey := test.prettyKey + if tenantID != roachpb.SystemTenantID { + prettyKey = fmt.Sprintf("/Tenant/%s%s", tenantID, prettyKey) + } + t.Run(prettyKey, func(t *testing.T) { + k, err := scanner.Scan(prettyKey) + if err != nil { + if test.expRemainder != "" { + if testutils.IsError(err, fmt.Sprintf("can't parse\"%s\"", test.expRemainder)) { + t.Fatalf("expected remainder: %s, got err: %s", test.expRemainder, err) + } + } else { + t.Fatal(err) } - } else { - t.Fatal(err) } - } - if test.expRemainder != "" && err == nil { - t.Fatalf("expected a remainder but got none: %s", test.expRemainder) - } - if test.expKey == nil { - if k != nil { - t.Fatalf("unexpected key returned: %s", k) + if test.expRemainder != "" && err == nil { + t.Fatalf("expected a remainder but got none: %s", test.expRemainder) } - return - } - expKey := test.expKey() - if !k.Equal(expKey) { - t.Fatalf("expected: %+v, got %+v", []byte(expKey), []byte(k)) - } - }) + if test.expKey == nil { + if k != nil { + t.Fatalf("unexpected key returned: %s", k) + } + return + } + expKey := test.expKey(tenantID) + if !k.Equal(expKey) { + t.Fatalf("expected: %+v, got %+v", []byte(expKey), []byte(k)) + } + }) + } } } diff --git a/pkg/util/keysutil/keys.go b/pkg/util/keysutil/keys.go index ac886a0ccdd4..41e677303845 100644 --- a/pkg/util/keysutil/keys.go +++ b/pkg/util/keysutil/keys.go @@ -42,23 +42,34 @@ type PrettyScanner struct { // pretty-printed keys from the table part of the keys space (i.e. inputs // starting with "/Table"). The supplied function needs to parse the part that // comes after "/Table". -func MakePrettyScanner(tableParser keys.KeyParserFunc) PrettyScanner { +// +// If tenantParser is not nil, it will replace the default function for scanning +// pretty-printed keys from the tenant part of the keys space (i.e. inputs +// starting with "/Tenant"). The supplied function needs to parse the part that +// comes after "/Tenant". +// +// At most one function can be non-nil. +func MakePrettyScanner(tableParser, tenantParser keys.KeyParserFunc) PrettyScanner { dict := keys.KeyDict - if tableParser != nil { - dict = customizeKeyComprehension(dict, tableParser) + if tableParser != nil && tenantParser != nil { + panic("both tableParser and tenantParser are non-nil") + } + if tableParser != nil || tenantParser != nil { + dict = customizeKeyComprehension(dict, tableParser, tenantParser) } return PrettyScanner{ keyComprehension: dict, // If we specified a custom parser, forget about the roundtrip. - validateRoundTrip: tableParser == nil, + validateRoundTrip: tableParser == nil && tenantParser == nil, } } // customizeKeyComprehension takes as input a KeyComprehensionTable and // overwrites the "pretty scanner" function for the tables key space (i.e. for -// keys starting with "/Table"). The modified table is returned. +// keys starting with "/Table") or for the tenant keys space (i.e. for keys +// starting with "/Tenant"). The modified table is returned. func customizeKeyComprehension( - table keys.KeyComprehensionTable, tableParser keys.KeyParserFunc, + table keys.KeyComprehensionTable, tableParser, tenantParser keys.KeyParserFunc, ) keys.KeyComprehensionTable { // Make a deep copy of the table. cpy := make(keys.KeyComprehensionTable, len(table)) @@ -69,20 +80,28 @@ func customizeKeyComprehension( } table = cpy - // Find the part of the table that deals with parsing table data. - // We'll perform surgery on it to apply `tableParser`. + // Find the part of the table that deals with parsing table / tenant data. + // We'll perform surgery on it to apply `tableParser` or 'tenantParser'. for i := range table { region := &table[i] - if region.Name == "/Table" { + if region.Name == "/Table" && tableParser != nil { if len(region.Entries) != 1 { - panic(fmt.Sprintf("expected a single entry under \"/Table\", got: %d", len(region.Entries))) + panic(fmt.Sprintf(`expected a single entry under "/Table", got: %d`, len(region.Entries))) } subRegion := ®ion.Entries[0] subRegion.PSFunc = tableParser return table } + if region.Name == "/Tenant" && tenantParser != nil { + if len(region.Entries) != 1 { + panic(fmt.Sprintf(`expected a single entry under "/Tenant", got: %d`, len(region.Entries))) + } + subRegion := ®ion.Entries[0] + subRegion.PSFunc = tenantParser + return table + } } - panic("failed to find required \"/Table\" entry") + panic(`failed to find required "/Table" and / or "/Tenant" entry`) } // Scan is a partial right inverse to PrettyPrint: it takes a key formatted for