Skip to content

Commit

Permalink
bootstrap: add string representation of initial values
Browse files Browse the repository at this point in the history
This commit adds InitialValuesToString and InitialValuesFromString which
convert KV datums and KV split keys from a MetadataSchema to a string
representation and vice-versa. The tenant prefix is omitted from the
string representation.

This functionality exists to allow bootstrapping clusters using the
values and splits as generated by an earlier cockroach binary.

Informs cockroachdb#94773.

Release note: None
  • Loading branch information
Marius Posta committed Jan 20, 2023
1 parent b38da1d commit 713b512
Show file tree
Hide file tree
Showing 5 changed files with 456 additions and 1 deletion.
2 changes: 2 additions & 0 deletions pkg/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ ALL_TESTS = [
"//pkg/spanconfig:spanconfig_test",
"//pkg/sql/backfill:backfill_test",
"//pkg/sql/cacheutil:cacheutil_test",
"//pkg/sql/catalog/bootstrap:bootstrap_test",
"//pkg/sql/catalog/catalogkeys:catalogkeys_test",
"//pkg/sql/catalog/catenumpb:catenumpb_disallowed_imports_test",
"//pkg/sql/catalog/catformat:catformat_test",
Expand Down Expand Up @@ -1416,6 +1417,7 @@ GO_TARGETS = [
"//pkg/sql/cacheutil:cacheutil",
"//pkg/sql/cacheutil:cacheutil_test",
"//pkg/sql/catalog/bootstrap:bootstrap",
"//pkg/sql/catalog/bootstrap:bootstrap_test",
"//pkg/sql/catalog/catalogkeys:catalogkeys",
"//pkg/sql/catalog/catalogkeys:catalogkeys_test",
"//pkg/sql/catalog/catenumpb:catenumpb",
Expand Down
19 changes: 18 additions & 1 deletion pkg/sql/catalog/bootstrap/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
load("//build/bazelutil/unused_checker:unused.bzl", "get_x_data")
load("@io_bazel_rules_go//go:def.bzl", "go_library")
load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")

go_library(
name = "bootstrap",
Expand Down Expand Up @@ -31,4 +31,21 @@ go_library(
],
)

go_test(
name = "bootstrap_test",
srcs = ["bootstrap_test.go"],
args = ["-test.timeout=295s"],
data = glob(["testdata/**"]),
embed = [":bootstrap"],
deps = [
"//pkg/config/zonepb",
"//pkg/keys",
"//pkg/roachpb",
"//pkg/testutils/datapathutils",
"//pkg/util/leaktest",
"@com_github_cockroachdb_datadriven//:datadriven",
"@com_github_stretchr_testify//require",
],
)

get_x_data(name = "get_x_data")
105 changes: 105 additions & 0 deletions pkg/sql/catalog/bootstrap/bootstrap_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,105 @@
// 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 bootstrap

import (
"crypto/sha256"
"encoding/hex"
"testing"

"github.com/cockroachdb/cockroach/pkg/config/zonepb"
"github.com/cockroachdb/cockroach/pkg/keys"
"github.com/cockroachdb/cockroach/pkg/roachpb"
"github.com/cockroachdb/cockroach/pkg/testutils/datapathutils"
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
"github.com/cockroachdb/datadriven"
"github.com/stretchr/testify/require"
)

func TestInitialValuesToString(t *testing.T) {
defer leaktest.AfterTest(t)()
datadriven.Walk(t, datapathutils.TestDataPath(t), func(t *testing.T, path string) {
datadriven.RunTest(t, path, func(t *testing.T, d *datadriven.TestData) string {
codec := keys.SystemSQLCodec
switch d.Cmd {
case "system":
break

case "tenant":
const dummyTenantID = 12345
codec = keys.MakeSQLCodec(roachpb.MustMakeTenantID(dummyTenantID))

default:
t.Fatalf("unexpected command %q", d.Cmd)
}
var expectedHash string
d.ScanArgs(t, "hash", &expectedHash)
ms := MakeMetadataSchema(codec, zonepb.DefaultZoneConfigRef(), zonepb.DefaultSystemZoneConfigRef())
result := InitialValuesToString(ms)
h := sha256.Sum256([]byte(result))
if actualHash := hex.EncodeToString(h[:]); expectedHash != actualHash {
t.Errorf(`Unexpected hash value %s for %s.
If you're seeing this error message, this means that the bootstrapped system
schema has changed. Assuming that this is expected:
- If this occurred during development on the main branch, rewrite the expected
test output and the hash value and move on.
- If this occurred during development of a patch for a release branch, make
very sure that the underlying change really is expected and is backward-
compatible and is absolutely necessary. If that's the case, then there are
hardcoded literals in the main development branch as well as any subsequent
release branches that need to be updated also.`, actualHash, d.Cmd)
}
return result
})
})
}

func TestRoundTripInitialValuesStringRepresentation(t *testing.T) {
t.Run("system", func(t *testing.T) {
roundTripInitialValuesStringRepresentation(t, 0 /* tenantID */)
})
t.Run("tenant", func(t *testing.T) {
const dummyTenantID = 54321
roundTripInitialValuesStringRepresentation(t, dummyTenantID)
})
t.Run("tenants", func(t *testing.T) {
const dummyTenantID1, dummyTenantID2 = 54321, 12345
require.Equal(t,
InitialValuesToString(makeMetadataSchema(dummyTenantID1)),
InitialValuesToString(makeMetadataSchema(dummyTenantID2)),
)
})
}

func roundTripInitialValuesStringRepresentation(t *testing.T, tenantID uint64) {
ms := makeMetadataSchema(tenantID)
expectedKVs, expectedSplits := ms.GetInitialValues()
actualKVs, actualSplits, err := InitialValuesFromString(ms.codec, InitialValuesToString(ms))
require.NoError(t, err)
require.Len(t, actualKVs, len(expectedKVs))
require.Len(t, actualSplits, len(expectedSplits))
for i, actualKV := range actualKVs {
expectedKV := expectedKVs[i]
require.EqualValues(t, expectedKV, actualKV)
}
for i, actualSplit := range actualSplits {
expectedSplit := expectedSplits[i]
require.EqualValues(t, expectedSplit, actualSplit)
}
}

func makeMetadataSchema(tenantID uint64) MetadataSchema {
codec := keys.SystemSQLCodec
if tenantID > 0 {
codec = keys.MakeSQLCodec(roachpb.MustMakeTenantID(tenantID))
}
return MakeMetadataSchema(codec, zonepb.DefaultZoneConfigRef(), zonepb.DefaultSystemZoneConfigRef())
}
91 changes: 91 additions & 0 deletions pkg/sql/catalog/bootstrap/metadata.go
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,12 @@
package bootstrap

import (
"bytes"
"context"
"encoding/hex"
"encoding/json"
"sort"
"strings"

"github.com/cockroachdb/cockroach/pkg/config/zonepb"
"github.com/cockroachdb/cockroach/pkg/keys"
Expand Down Expand Up @@ -251,6 +255,93 @@ func (ms MetadataSchema) GetInitialValues() ([]roachpb.KeyValue, []roachpb.RKey)
return ret, splits
}

type initialValueStrings struct {
Key string `json:"key"`
Value string `json:"value,omitempty"`
}

// InitialValuesToString returns a string representation of the return values
// of MetadataSchema.GetInitialValues. The tenant prefix is stripped.
func InitialValuesToString(ms MetadataSchema) string {
kvs, splits := ms.GetInitialValues()
// Collect the records.
type record struct {
k, v []byte
}
records := make([]record, 0, len(kvs)+len(splits))
for _, kv := range kvs {
records = append(records, record{k: kv.Key, v: kv.Value.TagAndDataBytes()})
}
for _, s := range splits {
records = append(records, record{k: s})
}
// Strip the tenant prefix if there is one.
p := []byte(ms.codec.TenantPrefix())
for i, r := range records {
if !bytes.Equal(p, r.k[:len(p)]) {
panic("unexpected prefix")
}
records[i].k = r.k[len(p):]
}
// Build the string representation.
s := make([]initialValueStrings, len(records))
for i, r := range records {
s[i].Key = hex.EncodeToString(r.k)
s[i].Value = hex.EncodeToString(r.v)
}
// Sort the records by key.
sort.Slice(s, func(i, j int) bool {
return s[i].Key < s[j].Key
})
// Build the string representation.
var sb strings.Builder
sb.WriteRune('[')
for i, r := range s {
if i > 0 {
sb.WriteRune(',')
}
b, err := json.Marshal(r)
if err != nil {
panic(err)
}
sb.Write(b)
sb.WriteRune('\n')
}
sb.WriteRune(']')
return sb.String()
}

// InitialValuesFromString is the reciprocal to InitialValuesToString and
// appends the tenant prefix from the given codec.
func InitialValuesFromString(
codec keys.SQLCodec, str string,
) (kvs []roachpb.KeyValue, splits []roachpb.RKey, _ error) {
p := codec.TenantPrefix()
var s []initialValueStrings
if err := json.Unmarshal([]byte(str), &s); err != nil {
return nil, nil, err
}
for i, r := range s {
k, err := hex.DecodeString(r.Key)
if err != nil {
return nil, nil, errors.Errorf("failed to decode hex key %s for record #%d", r.Key, i+1)
}
v, err := hex.DecodeString(r.Value)
if err != nil {
return nil, nil, errors.Errorf("failed to decode hex value %s fo record #%d", r.Value, i+1)
}
k = append(p[:len(p):len(p)], k...)
if len(v) == 0 {
splits = append(splits, k)
} else {
kv := roachpb.KeyValue{Key: k}
kv.Value.SetTagAndData(v)
kvs = append(kvs, kv)
}
}
return kvs, splits, nil
}

// DescriptorIDs returns the descriptor IDs present in the metadata schema in
// sorted order.
func (ms MetadataSchema) DescriptorIDs() descpb.IDs {
Expand Down
Loading

0 comments on commit 713b512

Please sign in to comment.