Merge #40281 #40426 #40428

40281: sql: Add range size to SHOW RANGES command. r=rohany a=rohany

Fixes #39767.

This is done by adding a new builtin function `crdb_internal.range_stats` that returns a json object containing stats about the requested range.

Release note (sql change): Add range size information to SHOW RANGES. Adds a builtin function crdb_internal.range_stats which returns a json object containing stats about the requested range.

40426: workload/cli: deprecate the --init flag on "workload run" r=nvanbenschoten a=danhhz

workload is now exposed to users and having two ways to initialize is
confusing. Previous users of the --init flag should use `workload init`
directly.

Release note (cli change): The --init flag on (the experimental)
"workload run" command is deprecated. Use "workload init" instead.

40428: cli: Change with load demo flag to be a persistent flag r=knz a=rohany

Previously, the command `cockroach demo movr --with-load` would
not work due to with-load not being a persistent flag. This
is undesired behavior as a workload should be able to be run
assuming that the chosen database has a workload defined.

Fixes #40427.

Release note: None

Co-authored-by: Rohan Yadav <[email protected]>
Co-authored-by: Daniel Harrison <[email protected]>

docs/generated/sql/functions.md

@@ -993,6 +993,8 @@ SELECT * FROM crdb_internal.check_consistency(true, ‘\x02’, ‘\x04’)</p>
 </span></td></tr>
 <tr><td><code>crdb_internal.pretty_key(raw_key: <a href="bytes.html">bytes</a>, skip_fields: <a href="int.html">int</a>) &rarr; <a href="string.html">string</a></code></td><td><span class="funcdesc"><p>This function is used only by CockroachDB’s developers for testing purposes.</p>
 </span></td></tr>
+<tr><td><code>crdb_internal.range_stats(key: <a href="bytes.html">bytes</a>) &rarr; jsonb</code></td><td><span class="funcdesc"><p>This function is used to retrieve range statistics information as a JSON object.</p>
+</span></td></tr>
 <tr><td><code>crdb_internal.round_decimal_values(val: <a href="decimal.html">decimal</a>, scale: <a href="int.html">int</a>) &rarr; <a href="decimal.html">decimal</a></code></td><td><span class="funcdesc"><p>This function is used internally to round decimal values during mutations.</p>
 </span></td></tr>
 <tr><td><code>crdb_internal.round_decimal_values(val: <a href="decimal.html">decimal</a>[], scale: <a href="int.html">int</a>) &rarr; <a href="decimal.html">decimal</a>[]</code></td><td><span class="funcdesc"><p>This function is used internally to round decimal array values during mutations.</p>

pkg/cli/flags.go

@@ -584,11 +584,11 @@ func init() {
 	// We add this command as a persistent flag so you can do stuff like
 	// ./cockroach demo movr --nodes=3.
 	IntFlag(demoFlags, &demoCtx.nodes, cliflags.DemoNodes, 1)
+	BoolFlag(demoFlags, &demoCtx.runWorkload, cliflags.RunDemoWorkload, false)
+	VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
 	// The --empty flag is only valid for the top level demo command,
 	// so we use the regular flag set.
 	BoolFlag(demoCmd.Flags(), &demoCtx.useEmptyDatabase, cliflags.UseEmptyDatabase, false)
-	BoolFlag(demoCmd.Flags(), &demoCtx.runWorkload, cliflags.RunDemoWorkload, false)
-	VarFlag(demoFlags, &demoCtx.localities, cliflags.DemoNodeLocality)
 
 	// sqlfmt command.
 	fmtFlags := sqlfmtCmd.Flags()

pkg/sql/crdb_internal.go

@@ -1869,7 +1869,9 @@ CREATE VIEW crdb_internal.ranges AS SELECT
 	replicas,
 	learner_replicas,
 	split_enforced_until,
-	crdb_internal.lease_holder(start_key) AS lease_holder
+	crdb_internal.lease_holder(start_key) AS lease_holder,
+	(crdb_internal.range_stats(start_key)->>'key_bytes')::INT +
+	(crdb_internal.range_stats(start_key)->>'val_bytes')::INT AS range_size
 FROM crdb_internal.ranges_no_leases
 `,
 	resultColumns: sqlbase.ResultColumns{
@@ -1885,6 +1887,7 @@ FROM crdb_internal.ranges_no_leases
 		{Name: "learner_replicas", Typ: types.Int2Vector},
 		{Name: "split_enforced_until", Typ: types.Timestamp},
 		{Name: "lease_holder", Typ: types.Int},
+		{Name: "range_size", Typ: types.Int},
 	},
 }
 
@@ -1906,7 +1909,7 @@ CREATE TABLE crdb_internal.ranges_no_leases (
   index_name           STRING NOT NULL,
 	replicas             INT[] NOT NULL,
 	learner_replicas     INT[] NOT NULL,
-  split_enforced_until TIMESTAMP
+	split_enforced_until TIMESTAMP
 )
 `,
 	generator: func(ctx context.Context, p *planner, _ *DatabaseDescriptor) (virtualTableGenerator, error) {

pkg/sql/delegate/show_ranges.go

@@ -41,6 +41,7 @@ func (d *delegator) delegateShowRanges(n *tree.ShowRanges) (tree.Statement, erro
 				ELSE crdb_internal.pretty_key(r.end_key, 2)
 			END AS end_key,
 			range_id,
+			range_size / 1000000 as range_size_mb,
 			replicas,
 			lease_holder,
 			locality
@@ -71,6 +72,7 @@ SELECT
   CASE WHEN r.start_key <= x'%s' THEN NULL ELSE crdb_internal.pretty_key(r.start_key, 2) END AS start_key,
   CASE WHEN r.end_key >= x'%s' THEN NULL ELSE crdb_internal.pretty_key(r.end_key, 2) END AS end_key,
   range_id,
+  range_size / 1000000 as range_size_mb,
   replicas,
   lease_holder,
   locality

pkg/sql/logictest/testdata/logic_test/crdb_internal

@@ -213,10 +213,10 @@ SELECT * FROM crdb_internal.zones WHERE false
 ----
 zone_id  target  range_name  database_name  table_name  index_name  partition_name  config_yaml  config_sql  config_protobuf
 
-query ITTTTTTTTTTT colnames
+query ITTTTTTTTTTTI colnames
 SELECT * FROM crdb_internal.ranges WHERE range_id < 0
 ----
-range_id  start_key  start_pretty  end_key  end_pretty  database_name  table_name  index_name  replicas  learner_replicas  split_enforced_until  lease_holder
+range_id  start_key  start_pretty  end_key  end_pretty  database_name  table_name  index_name  replicas  learner_replicas  split_enforced_until  lease_holder range_size
 
 query ITTTTTTTTTT colnames
 SELECT * FROM crdb_internal.ranges_no_leases WHERE range_id < 0

pkg/sql/sem/builtins/builtins.go

@@ -16,6 +16,7 @@ import (
 	"crypto/sha1"
 	"crypto/sha256"
 	"crypto/sha512"
+	gojson "encoding/json"
 	"fmt"
 	"hash"
 	"hash/crc32"
@@ -3109,6 +3110,42 @@ may increase either contention or retry errors, or both.`,
 		},
 	),
 
+	// Return statistics about a range.
+	"crdb_internal.range_stats": makeBuiltin(
+		tree.FunctionProperties{
+			Category: categorySystemInfo,
+		},
+		tree.Overload{
+			Types: tree.ArgTypes{
+				{"key", types.Bytes},
+			},
+			ReturnType: tree.FixedReturnType(types.Jsonb),
+			Fn: func(ctx *tree.EvalContext, args tree.Datums) (tree.Datum, error) {
+				key := []byte(tree.MustBeDBytes(args[0]))
+				b := &client.Batch{}
+				b.AddRawRequest(&roachpb.RangeStatsRequest{
+					RequestHeader: roachpb.RequestHeader{
+						Key: key,
+					},
+				})
+				if err := ctx.Txn.Run(ctx.Context, b); err != nil {
+					return nil, pgerror.Newf(pgcode.InvalidParameterValue, "message: %s", err)
+				}
+				resp := b.RawResponse().Responses[0].GetInner().(*roachpb.RangeStatsResponse).MVCCStats
+				jsonStr, err := gojson.Marshal(&resp)
+				if err != nil {
+					return nil, err
+				}
+				jsonDatum, err := tree.ParseDJSON(string(jsonStr))
+				if err != nil {
+					return nil, err
+				}
+				return jsonDatum, nil
+			},
+			Info: "This function is used to retrieve range statistics information as a JSON object.",
+		},
+	),
+
 	"crdb_internal.set_vmodule": makeBuiltin(
 		tree.FunctionProperties{
 			Category: categorySystemInfo,

pkg/sql/show_ranges_test.go

@@ -35,8 +35,8 @@ func TestShowRangesWithLocality(t *testing.T) {
 	sqlDB.Exec(t, `CREATE TABLE t (x INT PRIMARY KEY)`)
 	sqlDB.Exec(t, `ALTER TABLE t SPLIT AT SELECT i FROM generate_series(0, 20) AS g(i)`)
 
-	const nodeColIdx = 4
-	const localityColIdx = 5
+	const nodeColIdx = 5
+	const localityColIdx = 6
 
 	result := sqlDB.QueryStr(t, `SHOW RANGES FROM TABLE t`)
 	for _, row := range result {

pkg/workload/cli/run.go

@@ -45,7 +45,7 @@ var maxRate = runFlags.Float64(
 	"max-rate", 0, "Maximum frequency of operations (reads/writes). If 0, no limit.")
 var maxOps = runFlags.Uint64("max-ops", 0, "Maximum number of operations to run")
 var duration = runFlags.Duration("duration", 0, "The duration to run. If 0, run forever.")
-var doInit = runFlags.Bool("init", false, "Automatically run init")
+var doInit = runFlags.Bool("init", false, "Automatically run init. DEPRECATED: Use workload init instead.")
 var ramp = runFlags.Duration("ramp", 0*time.Second, "The duration over which to ramp up load.")
 
 var initFlags = pflag.NewFlagSet(`init`, pflag.ContinueOnError)
@@ -326,6 +326,8 @@ func runRun(gen workload.Generator, urls []string, dbName string) error {
 		return err
 	}
 	if *doInit || *drop {
+		log.Info(ctx, `DEPRECATION: `+
+			`the --init flag on "workload run" will no longer be supported after 19.2`)
 		for {
 			err = runInitImpl(ctx, gen, initDB, dbName)
 			if err == nil {