From 4afb6309133b63e1421dc0133f7dfbecc6afe751 Mon Sep 17 00:00:00 2001
From: Jackson Owens <jackson@cockroachlabs.com>
Date: Thu, 24 Oct 2024 16:31:32 -0400
Subject: [PATCH] cmd/pebble: populate common KeySchemas

Update the cmd/pebble CLI tool to populate opts.KeySchemas with commonly used
KeySchemas. This was motivated by a desire to use the cmd/pebble CLI tool with
a sstable constructed by the metamorphic tests.
---
 cmd/pebble/main.go | 17 ++++++++++++++++-
 tool/tool.go       |  3 +++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/cmd/pebble/main.go b/cmd/pebble/main.go
index 688a868775..122e89e19c 100644
--- a/cmd/pebble/main.go
+++ b/cmd/pebble/main.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/cockroachdb/pebble/internal/crdbtest"
 	"github.com/cockroachdb/pebble/internal/testkeys"
+	"github.com/cockroachdb/pebble/sstable/colblk"
 	"github.com/cockroachdb/pebble/tool"
 	"github.com/spf13/cobra"
 )
@@ -31,6 +32,16 @@ var (
 	secondaryCacheSize int64
 )
 
+// Define the default testkeys key schema. Feeding this schema to the tool
+// ensures the cmd/pebble cli tool natively understands the sstables constructed
+// by our test cases.
+//
+// TODO(jackson): Ideally when a sstable.Reader finds a key schema that's a
+// DefaultKeySchema not already in the KeySchemas constructed with a Comparer
+// that it knows of, it would automatically construct the appropriate KeySchema.
+// Or at least the cli tool should.
+var testKeysSchema = colblk.DefaultKeySchema(testkeys.Comparer, 16)
+
 func main() {
 	log.SetFlags(0)
 
@@ -58,7 +69,11 @@ func main() {
 	}
 	rootCmd.AddCommand(benchCmd)
 
-	t := tool.New(tool.Comparers(&crdbtest.Comparer, testkeys.Comparer), tool.Mergers(fauxMVCCMerger))
+	t := tool.New(
+		tool.Comparers(&crdbtest.Comparer, testkeys.Comparer),
+		tool.Mergers(fauxMVCCMerger),
+		tool.KeySchemas(&crdbtest.KeySchema, &testKeysSchema),
+	)
 	rootCmd.AddCommand(t.Commands...)
 
 	for _, cmd := range []*cobra.Command{replayCmd, scanCmd, syncCmd, tombstoneCmd, writeBenchCmd, ycsbCmd} {
diff --git a/tool/tool.go b/tool/tool.go
index d83f5ab0ed..b42f195282 100644
--- a/tool/tool.go
+++ b/tool/tool.go
@@ -60,6 +60,9 @@ func Comparers(cmps ...*Comparer) Option {
 // introspection tools.
 func KeySchemas(schemas ...*colblk.KeySchema) Option {
 	return func(t *T) {
+		if t.opts.KeySchemas == nil {
+			t.opts.KeySchemas = make(map[string]*colblk.KeySchema)
+		}
 		for _, s := range schemas {
 			t.opts.KeySchemas[s.Name] = s
 		}