sql: support scalar expressions (& placeholders) for relocate store IDs

This commit makes it possible to use arbitrary scalar
expressions (incl subqueries, placeholders etc) in the FROM/TO clauses
of ALTER RANGE ... RELOCATE statements.

For example:

```
ALTER RANGE 123 RELOCATE FROM $1 TO $2
ALTER RANGE 123 RELOCATE FROM $1 TO $1+10
```

Release note (sql change): The experimental ALTER RANGE...RELOCATE
syntax now accepts arbitrary scalar expressions as the source and target
store IDs.

docs/generated/sql/bnf/alter_range_relocate_stmt.bnf

@@ -1,5 +1,5 @@
 alter_range_relocate_stmt ::=
-	'ALTER' 'RANGE' relocate_kw 'LEASE' 'TO' iconst64 'FOR' select_stmt
-	| 'ALTER' 'RANGE' iconst64 relocate_kw 'LEASE' 'TO' iconst64
-	| 'ALTER' 'RANGE' relocate_kw relocate_subject_nonlease 'FROM' iconst64 'TO' iconst64 'FOR' select_stmt
-	| 'ALTER' 'RANGE' iconst64 relocate_kw relocate_subject_nonlease 'FROM' iconst64 'TO' iconst64
+	'ALTER' 'RANGE' relocate_kw 'LEASE' 'TO' a_expr 'FOR' select_stmt
+	| 'ALTER' 'RANGE' iconst64 relocate_kw 'LEASE' 'TO' a_expr
+	| 'ALTER' 'RANGE' relocate_kw relocate_subject_nonlease 'FROM' a_expr 'TO' a_expr 'FOR' select_stmt
+	| 'ALTER' 'RANGE' iconst64 relocate_kw relocate_subject_nonlease 'FROM' a_expr 'TO' a_expr

docs/generated/sql/bnf/stmt_block.bnf

@@ -1846,10 +1846,10 @@ alter_zone_range_stmt ::=
 	'ALTER' 'RANGE' zone_name set_zone_config
 
 alter_range_relocate_stmt ::=
-	'ALTER' 'RANGE' relocate_kw 'LEASE' 'TO' iconst64 'FOR' select_stmt
-	| 'ALTER' 'RANGE' iconst64 relocate_kw 'LEASE' 'TO' iconst64
-	| 'ALTER' 'RANGE' relocate_kw relocate_subject_nonlease 'FROM' iconst64 'TO' iconst64 'FOR' select_stmt
-	| 'ALTER' 'RANGE' iconst64 relocate_kw relocate_subject_nonlease 'FROM' iconst64 'TO' iconst64
+	'ALTER' 'RANGE' relocate_kw 'LEASE' 'TO' a_expr 'FOR' select_stmt
+	| 'ALTER' 'RANGE' iconst64 relocate_kw 'LEASE' 'TO' a_expr
+	| 'ALTER' 'RANGE' relocate_kw relocate_subject_nonlease 'FROM' a_expr 'TO' a_expr 'FOR' select_stmt
+	| 'ALTER' 'RANGE' iconst64 relocate_kw relocate_subject_nonlease 'FROM' a_expr 'TO' a_expr
 
 alter_zone_partition_stmt ::=
 	'ALTER' 'PARTITION' partition_name 'OF' 'TABLE' table_name set_zone_config

pkg/kv/kvserver/client_lease_test.go

@@ -969,11 +969,11 @@ func TestAlterRangeRelocate(t *testing.T) {
 	// We start with having the range under test on (1,2,3).
 	db := tc.ServerConn(0)
 	// Move 2 -> 4.
-	_, err := db.Exec("ALTER RANGE " + rhsDesc.RangeID.String() + " RELOCATE FROM 2 TO 4")
+	_, err := db.Exec("ALTER RANGE "+rhsDesc.RangeID.String()+" RELOCATE FROM $1 TO $2", 2, 4)
 	require.NoError(t, err)
 	require.NoError(t, tc.WaitForVoters(rhsDesc.StartKey.AsRawKey(), tc.Targets(0, 2, 3)...))
 	// Move lease 1 -> 4.
-	_, err = db.Exec("ALTER RANGE " + rhsDesc.RangeID.String() + " RELOCATE LEASE TO 4")
+	_, err = db.Exec("ALTER RANGE "+rhsDesc.RangeID.String()+" RELOCATE LEASE TO $1", 4)
 	require.NoError(t, err)
 	testutils.SucceedsSoon(t, func() error {
 		repl := tc.GetFirstStoreFromServer(t, 3).LookupReplica(rhsDesc.StartKey)
@@ -988,7 +988,7 @@ func TestAlterRangeRelocate(t *testing.T) {
 	})
 
 	// Move lease 3 -> 5.
-	_, err = db.Exec("ALTER RANGE RELOCATE FROM 3 TO 5 FOR (SELECT range_id from crdb_internal.ranges where range_id = " + rhsDesc.RangeID.String() + ")")
+	_, err = db.Exec("ALTER RANGE RELOCATE FROM $1 TO $2 FOR (SELECT range_id from crdb_internal.ranges where range_id = $3)", 3, 5, rhsDesc.RangeID)
 	require.NoError(t, err)
 	require.NoError(t, tc.WaitForVoters(rhsDesc.StartKey.AsRawKey(), tc.Targets(0, 3, 4)...))
 }

pkg/sql/distsql_spec_exec_factory.go

@@ -1076,7 +1076,10 @@ func (e *distSQLSpecExecFactory) ConstructAlterTableRelocate(
 }
 
 func (e *distSQLSpecExecFactory) ConstructAlterRangeRelocate(
-	input exec.Node, relocateSubject tree.RelocateSubject, toStoreID int64, fromStoreID int64,
+	input exec.Node,
+	relocateSubject tree.RelocateSubject,
+	toStoreID tree.TypedExpr,
+	fromStoreID tree.TypedExpr,
 ) (exec.Node, error) {
 	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: alter range relocate")
 }

pkg/sql/opt/exec/execbuilder/statement.go

@@ -235,11 +235,20 @@ func (b *Builder) buildAlterRangeRelocate(relocate *memo.AlterRangeRelocateExpr)
 	if err != nil {
 		return execPlan{}, err
 	}
+	scalarCtx := buildScalarCtx{}
+	toStoreID, err := b.buildScalar(&scalarCtx, relocate.ToStoreID)
+	if err != nil {
+		return execPlan{}, err
+	}
+	fromStoreID, err := b.buildScalar(&scalarCtx, relocate.FromStoreID)
+	if err != nil {
+		return execPlan{}, err
+	}
 	node, err := b.factory.ConstructAlterRangeRelocate(
 		input.root,
 		relocate.SubjectReplicas,
-		relocate.ToStoreID,
-		relocate.FromStoreID,
+		toStoreID,
+		fromStoreID,
 	)
 	if err != nil {
 		return execPlan{}, err

pkg/sql/opt/exec/explain/emit.go

@@ -26,7 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/humanizeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
-	"github.com/dustin/go-humanize"
+	humanize "github.com/dustin/go-humanize"
 )
 
 // Emit produces the EXPLAIN output against the given OutputBuilder. The
@@ -828,6 +828,11 @@ func (e *emitter) emitNodeAttributes(n *Node) error {
 			ob.Attrf("deduplicate", "")
 		}
 
+	case alterRangeRelocateOp:
+		a := n.args.(*alterRangeRelocateArgs)
+		ob.Expr("to", a.toStoreID, nil /* columns */)
+		ob.Expr("from", a.fromStoreID, nil /* columns */)
+
 	case simpleProjectOp,
 		serializingProjectOp,
 		ordinalityOp,
@@ -848,7 +853,6 @@ func (e *emitter) emitNodeAttributes(n *Node) error {
 		alterTableUnsplitOp,
 		alterTableUnsplitAllOp,
 		alterTableRelocateOp,
-		alterRangeRelocateOp,
 		controlJobsOp,
 		controlSchedulesOp,
 		cancelQueriesOp,

pkg/sql/opt/exec/factory.opt

@@ -736,6 +736,6 @@ define Export {
 define AlterRangeRelocate {
     input exec.Node
     subjectReplicas tree.RelocateSubject
-    toStoreID int64
-    fromStoreID int64
+    toStoreID tree.TypedExpr
+    fromStoreID tree.TypedExpr
 }

pkg/sql/opt/ops/statement.opt

@@ -297,14 +297,18 @@ define CreateStatisticsPrivate {
 define AlterRangeRelocate {
     # The input expression provides range IDs as integers.
     Input RelExpr
+    # The origin store ID.
+    ToStoreID ScalarExpr
+    # The destination store ID.
+    FromStoreID ScalarExpr
+
     _ AlterRangeRelocatePrivate
 }
 
 [Private]
 define AlterRangeRelocatePrivate {
+    # The subject indicates which replicas will be relocated.
     SubjectReplicas RelocateSubject
-    ToStoreID int64
-    FromStoreID int64
 
     # Columns stores the column IDs for the statement result columns.
     Columns ColList

pkg/sql/opt/optbuilder/alter_range.go

@@ -27,6 +27,13 @@ func (b *Builder) buildAlterRangeRelocate(
 		panic(err)
 	}
 
+	toStoreID := b.resolveAndBuildScalar(
+		relocate.ToStoreID, types.Int, exprKindStoreID, tree.RejectSpecial, inScope,
+	)
+	fromStoreID := b.resolveAndBuildScalar(
+		relocate.FromStoreID, types.Int, exprKindStoreID, tree.RejectSpecial, inScope,
+	)
+
 	// Disable optimizer caching, as we do for other ALTER statements.
 	b.DisableMemoReuse = true
 
@@ -47,10 +54,10 @@ func (b *Builder) buildAlterRangeRelocate(
 
 	outScope.expr = b.factory.ConstructAlterRangeRelocate(
 		inputScope.expr,
+		toStoreID,
+		fromStoreID,
 		&memo.AlterRangeRelocatePrivate{
 			SubjectReplicas: relocate.SubjectReplicas,
-			ToStoreID:       relocate.ToStoreID,
-			FromStoreID:     relocate.FromStoreID,
 			Columns:         colsToColList(outScope.cols),
 			Props:           physical.MinRequired,
 		},

pkg/sql/opt/optbuilder/scope.go

@@ -126,6 +126,7 @@ const (
 	exprKindOrderBy
 	exprKindReturning
 	exprKindSelect
+	exprKindStoreID
 	exprKindValues
 	exprKindWhere
 	exprKindWindowFrameStart
@@ -146,6 +147,7 @@ var exprKindName = [...]string{
 	exprKindOrderBy:           "ORDER BY",
 	exprKindReturning:         "RETURNING",
 	exprKindSelect:            "SELECT",
+	exprKindStoreID:           "RELOCATE STORE ID",
 	exprKindValues:            "VALUES",
 	exprKindWhere:             "WHERE",
 	exprKindWindowFrameStart:  "WINDOW FRAME START",

pkg/sql/opt/optbuilder/testdata/alter_range

@@ -4,74 +4,86 @@ CREATE TABLE abc (a INT PRIMARY KEY, b INT, c STRING, INDEX b (b), UNIQUE INDEX
 
 # Tests for ALTER RANGE RELOCATE.
 build
-ALTER RANGE 1 RELOCATE FROM 1 TO 2
+ALTER RANGE 1 RELOCATE FROM 1+2 TO 3+4
 ----
-alter-range-relocate &{VOTERS 2 1 [3 4 5] []}
+alter-range-relocate &{VOTERS [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── values
-      ├── columns: column1:6!null
-      └── (1,)
+ ├── values
+ │    ├── columns: column1:6!null
+ │    └── (1,)
+ ├── 3 + 4
+ └── 1 + 2
 
 build
-ALTER RANGE RELOCATE FROM 1 TO 2 FOR SELECT a from abc
+ALTER RANGE RELOCATE FROM 1+2 TO 3+4 FOR SELECT a FROM abc
 ----
-alter-range-relocate &{VOTERS 2 1 [3 4 5] []}
+alter-range-relocate &{VOTERS [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── project
-      ├── columns: a:6!null
-      └── scan abc
-           └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── project
+ │    ├── columns: a:6!null
+ │    └── scan abc
+ │         └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── 3 + 4
+ └── 1 + 2
 
 build
-ALTER RANGE a RELOCATE FROM 1 TO 2
+ALTER RANGE a RELOCATE FROM 1+2 TO 3+4
 ----
 error (42601): at or near "relocate": syntax error
 
 build
-ALTER RANGE RELOCATE FROM 1 TO 2 FOR SELECT c from abc
+ALTER RANGE RELOCATE FROM 1+2 TO 3+4 FOR SELECT c FROM abc
 ----
 error (42601): RELOCATE VOTERS data column 1 (range ids) must be of type int, not type string
 
 build
-ALTER RANGE 1 RELOCATE NONVOTERS FROM 1 TO 2
+ALTER RANGE 1 RELOCATE NONVOTERS FROM 1+2 TO 3+4
 ----
-alter-range-relocate &{NONVOTERS 2 1 [3 4 5] []}
+alter-range-relocate &{NONVOTERS [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── values
-      ├── columns: column1:6!null
-      └── (1,)
+ ├── values
+ │    ├── columns: column1:6!null
+ │    └── (1,)
+ ├── 3 + 4
+ └── 1 + 2
 
 build
-ALTER RANGE RELOCATE NONVOTERS FROM 1 TO 2 FOR SELECT a from abc
+ALTER RANGE RELOCATE NONVOTERS FROM 1+2 TO 3+4 FOR SELECT a FROM abc
 ----
-alter-range-relocate &{NONVOTERS 2 1 [3 4 5] []}
+alter-range-relocate &{NONVOTERS [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── project
-      ├── columns: a:6!null
-      └── scan abc
-           └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── project
+ │    ├── columns: a:6!null
+ │    └── scan abc
+ │         └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── 3 + 4
+ └── 1 + 2
 
 # Tests for ALTER RANGE RELOCATE LEASE.
 build
-ALTER RANGE 1 RELOCATE LEASE TO 2
+ALTER RANGE 1 RELOCATE LEASE TO 1+2
 ----
-alter-range-relocate &{LEASE 2 0 [3 4 5] []}
+alter-range-relocate &{LEASE [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── values
-      ├── columns: column1:6!null
-      └── (1,)
+ ├── values
+ │    ├── columns: column1:6!null
+ │    └── (1,)
+ ├── 1 + 2
+ └── NULL::INT8
 
 build
-ALTER RANGE RELOCATE LEASE TO 2 FOR SELECT a from abc
+ALTER RANGE RELOCATE LEASE TO 1+2 FOR SELECT a FROM abc
 ----
-alter-range-relocate &{LEASE 2 0 [3 4 5] []}
+alter-range-relocate &{LEASE [3 4 5] []}
  ├── columns: range_id:3 pretty:4 result:5
- └── project
-      ├── columns: a:6!null
-      └── scan abc
-           └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── project
+ │    ├── columns: a:6!null
+ │    └── scan abc
+ │         └── columns: a:6!null b:7 c:8 crdb_internal_mvcc_timestamp:9 tableoid:10
+ ├── 1 + 2
+ └── NULL::INT8
 
 build
-ALTER RANGE RELOCATE LEASE TO 2 FOR SELECT c from abc
+ALTER RANGE RELOCATE LEASE TO 1+2 FOR SELECT c FROM abc
 ----
 error (42601): RELOCATE LEASE data column 1 (range ids) must be of type int, not type string

pkg/sql/opt_exec_factory.go

@@ -1971,7 +1971,10 @@ func (ef *execFactory) ConstructAlterTableRelocate(
 
 // ConstructAlterRangeRelocate is part of the exec.Factory interface.
 func (ef *execFactory) ConstructAlterRangeRelocate(
-	input exec.Node, relocateSubject tree.RelocateSubject, toStoreID int64, fromStoreID int64,
+	input exec.Node,
+	relocateSubject tree.RelocateSubject,
+	toStoreID tree.TypedExpr,
+	fromStoreID tree.TypedExpr,
 ) (exec.Node, error) {
 	if !ef.planner.ExecCfg().Codec.ForSystemTenant() {
 		return nil, errorutil.UnsupportedWithMultiTenancy(54250)
@@ -1980,8 +1983,8 @@ func (ef *execFactory) ConstructAlterRangeRelocate(
 	return &relocateRange{
 		rows:            input.(planNode),
 		subjectReplicas: relocateSubject,
-		toStoreID:       roachpb.StoreID(toStoreID),
-		fromStoreID:     roachpb.StoreID(fromStoreID),
+		toStoreID:       toStoreID,
+		fromStoreID:     fromStoreID,
 	}, nil
 }