sql: enables distributed distsql queries for multi-tenant

This change allows SQL queries to be distributed in multi-tenant
environments. The distribution algorithm randomly assigns spans to SQL
instances, but if only one instance is used the spans are assigned
instead to the gateway instance. Distribution does not take locality
into account, which will be implemented in a future PR.

This change also supports running execbuilder tests with the
3node-tenant configuration, which is under CCL. These tests can be run
in the following manner:

```
make test PKG=./pkg/ccl/logictestccl TESTS=TestTenantExecBuild
./dev test pkg/ccl/logictestccl -f=TestTenantExecBuild
```

Fixes: cockroachdb#80680

Release note: None

pkg/ccl/logictestccl/BUILD.bazel

@@ -15,8 +15,9 @@ go_test(
         "main_test.go",
     ],
     data = glob(["testdata/**"]) + [
-        "//pkg/sql/logictest:testdata",
         "//c-deps:libgeos",
+        "//pkg/sql/logictest:testdata",
+        "//pkg/sql/opt/exec/execbuilder:testdata",
     ],
     embed = [":logictestccl"],
     deps = [
@@ -29,6 +30,7 @@ go_test(
         "//pkg/sql/logictest",
         "//pkg/testutils",
         "//pkg/testutils/serverutils",
+        "//pkg/testutils/skip",
         "//pkg/testutils/testcluster",
         "//pkg/util/leaktest",
         "//pkg/util/randutil",

pkg/ccl/logictestccl/logic_test.go

@@ -16,6 +16,7 @@ import (
 	_ "github.com/cockroachdb/cockroach/pkg/ccl"
 	"github.com/cockroachdb/cockroach/pkg/sql/logictest"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 )
 
@@ -54,3 +55,24 @@ func TestTenantSQLLiteLogic(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	logictest.RunSQLLiteLogicTest(t, "3node-tenant")
 }
+
+// TestTenantExecBuild runs execbuilder test files under the 3node-tenant
+// configuration, which constructs a secondary tenant and runs the test within
+// that secondary tenant's sandbox.
+func TestTenantExecBuild(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	skip.UnderDeadlock(t, "times out and/or hangs")
+
+	testdataDir := "../../sql/opt/exec/execbuilder/testdata/"
+	if bazel.BuiltWithBazel() {
+		runfile, err := bazel.Runfile("pkg/sql/opt/exec/execbuilder/testdata/")
+		if err != nil {
+			t.Fatal(err)
+		}
+		testdataDir = runfile
+	}
+
+	logictest.RunLogicTestWithDefaultConfig(
+		t, logictest.TestServerArgs{DisableWorkmemRandomization: true}, "3node-tenant", true, /* runCCLConfigs */
+		filepath.Join(testdataDir, "[^.]*"))
+}

pkg/keys/keys.go

@@ -835,10 +835,26 @@ func GetRowPrefixLength(key roachpb.Key) (int, error) {
 	}
 	sqlN := len(sqlKey)
 
+	// Check that the prefix contains a valid TableID.
 	if encoding.PeekType(sqlKey) != encoding.Int {
 		// Not a table key, so the row prefix is the entire key.
 		return n, nil
 	}
+	tableIDLen, err := encoding.GetUvarintLen(sqlKey)
+	if err != nil {
+		return 0, err
+	}
+
+	// Check whether the prefix contains a valid IndexID after the TableID. Not
+	// all keys contain an index ID.
+	if encoding.PeekType(sqlKey[tableIDLen:]) == encoding.Int {
+		indexIDLen, err := encoding.GetUvarintLen(sqlKey[tableIDLen:])
+		// If the IndexID is the last part of the key, the entire key is the prefix.
+		if err == nil && tableIDLen+indexIDLen == sqlN {
+			return n, nil
+		}
+	}
+
 	// The column family ID length is encoded as a varint and we take advantage
 	// of the fact that the column family ID itself will be encoded in 0-9 bytes
 	// and thus the length of the column family ID data will fit in a single

pkg/keys/keys_test.go

@@ -622,6 +622,7 @@ func TestEnsureSafeSplitKey(t *testing.T) {
 		expected roachpb.Key
 	}{
 		{es(), es()},                     // Not a table key
+		{es(1, 2), es(1, 2)},             // /Table/1/2 -> /Table/1/2
 		{es(1, 2, 0), es(1, 2)},          // /Table/1/2/0 -> /Table/1/2
 		{es(1, 2, 1), es(1)},             // /Table/1/2/1 -> /Table/1
 		{es(1, 2, 2), es()},              // /Table/1/2/2 -> /Table
@@ -631,6 +632,7 @@ func TestEnsureSafeSplitKey(t *testing.T) {
 		{es(1, 2, 3, 4, 1), es(1, 2, 3)}, // /Table/1/2/3/4/1 -> /Table/1/2/3
 		// Same test cases, but for tenant 5.
 		{e5(), e5()},                     // Not a table key
+		{e5(1, 2), e5(1, 2)},             // /Tenant5/Table/1/2 -> /Tenant5/Table/1/2
 		{e5(1, 2, 0), e5(1, 2)},          // /Tenant/5/Table/1/2/0 -> /Tenant/5/Table/1/2
 		{e5(1, 2, 1), e5(1)},             // /Tenant/5/Table/1/2/1 -> /Tenant/5/Table/1
 		{e5(1, 2, 2), e5()},              // /Tenant/5/Table/1/2/2 -> /Tenant/5/Table
@@ -673,24 +675,23 @@ func TestEnsureSafeSplitKey(t *testing.T) {
 		err string
 	}{
 		// Column ID suffix size is too large.
-		{es(1), "malformed table key"},
-		{es(1, 2), "malformed table key"},
+		{es(1), "not a valid table key"},
+		{es(1, 2, 5), "malformed table key"},
 		// The table ID is invalid.
-		{es(200)[:1], "insufficient bytes to decode uvarint value"},
-		// The index ID is invalid.
-		{es(1, 200)[:2], "insufficient bytes to decode uvarint value"},
+		{es(200)[:1], "not a valid table key"},
 		// The column ID suffix is invalid.
+		{es(1, 200)[:2], "insufficient bytes to decode uvarint value"},
 		{es(1, 2, 200)[:3], "insufficient bytes to decode uvarint value"},
 		// Exercises a former overflow bug. We decode a uint(18446744073709551610) which, if cast
 		// to int carelessly, results in -6.
-		{encoding.EncodeVarintAscending(tenSysCodec.TablePrefix(999), 322434), "malformed table key"},
+		{encoding.EncodeVarintAscending(encoding.EncodeVarintAscending(tenSysCodec.TablePrefix(999), 2), 322434), "malformed table key"},
 		// Same test cases, but for tenant 5.
-		{e5(1), "malformed table key"},
-		{e5(1, 2), "malformed table key"},
-		{e5(200)[:3], "insufficient bytes to decode uvarint value"},
+		{e5(1), "not a valid table key"},
+		{e5(1, 2, 5), "malformed table key"},
+		{e5(200)[:3], "not a valid table key"},
 		{e5(1, 200)[:4], "insufficient bytes to decode uvarint value"},
 		{e5(1, 2, 200)[:5], "insufficient bytes to decode uvarint value"},
-		{encoding.EncodeVarintAscending(ten5Codec.TablePrefix(999), 322434), "malformed table key"},
+		{encoding.EncodeVarintAscending(encoding.EncodeVarintAscending(ten5Codec.TablePrefix(999), 2), 322434), "malformed table key"},
 	}
 	for i, d := range errorData {
 		_, err := EnsureSafeSplitKey(d.in)

pkg/sql/apply_join.go

@@ -308,7 +308,7 @@ func runPlanInsidePlan(params runParams, plan *planComponents, resultWriter rowR
 	)
 	distributeType := DistributionType(DistributionTypeNone)
 	if distributePlan.WillDistribute() {
-		distributeType = DistributionTypeSystemTenantOnly
+		distributeType = DistributionTypeAlways
 	}
 	planCtx := params.p.extendedEvalCtx.ExecCfg.DistSQLPlanner.NewPlanningCtx(
 		params.ctx, evalCtx, &plannerCopy, params.p.txn, distributeType)

pkg/sql/conn_executor_exec.go

@@ -1139,7 +1139,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	}
 	distribute := DistributionType(DistributionTypeNone)
 	if distributePlan.WillDistribute() {
-		distribute = DistributionTypeSystemTenantOnly
+		distribute = DistributionTypeAlways
 	}
 	ex.sessionTracing.TraceExecStart(ctx, "distributed")
 	stats, err := ex.execWithDistSQLEngine(

pkg/sql/distsql_physical_planner.go

@@ -437,6 +437,12 @@ func mustWrapValuesNode(planCtx *PlanningCtx, specifiedInQuery bool) bool {
 func checkSupportForPlanNode(node planNode) (distRecommendation, error) {
 	switch n := node.(type) {
 	// Keep these cases alphabetized, please!
+	case *createStatsNode:
+		if n.runAsJob {
+			return cannotDistribute, planNodeNotSupportedErr
+		}
+		return shouldDistribute, nil
+
 	case *distinctNode:
 		return checkSupportForPlanNode(n.plan)
 
@@ -652,12 +658,6 @@ func checkSupportForPlanNode(node planNode) (distRecommendation, error) {
 		}
 		return shouldDistribute, nil
 
-	case *createStatsNode:
-		if n.runAsJob {
-			return cannotDistribute, planNodeNotSupportedErr
-		}
-		return shouldDistribute, nil
-
 	default:
 		return cannotDistribute, planNodeNotSupportedErr
 	}
@@ -1165,7 +1165,7 @@ func (dsp *DistSQLPlanner) partitionSpansTenant(
 		return partitions, nil
 	}
 	if dsp.sqlInstanceProvider == nil {
-		return nil, errors.New("sql instance provider not available in multi-tenant environment")
+		return nil, errors.AssertionFailedf("sql instance provider not available in multi-tenant environment")
 	}
 	// GetAllInstances only returns healthy instances.
 	instances, err := dsp.sqlInstanceProvider.GetAllInstances(ctx)
@@ -1183,11 +1183,24 @@ func (dsp *DistSQLPlanner) partitionSpansTenant(
 
 	// nodeMap maps a SQLInstanceID to an index inside the partitions array.
 	nodeMap := make(map[base.SQLInstanceID]int)
+	var lastKey roachpb.Key
+	var lastIdx int
 	for i := range spans {
 		span := spans[i]
 		if log.V(1) {
 			log.Infof(ctx, "partitioning span %s", span)
 		}
+		// Rows with column families may have been split into different spans. These
+		// spans should be assigned the same pod so that the pod can stitch together
+		// the rows correctly. Split rows are in adjacent spans.
+		if safeKey, err := keys.EnsureSafeSplitKey(span.Key); err == nil {
+			if safeKey.Equal(lastKey) {
+				partition := &partitions[lastIdx]
+				partition.Spans = append(partition.Spans, span)
+				continue
+			}
+			lastKey = safeKey
+		}
 		sqlInstanceID := instances[i%len(instances)].InstanceID
 		partitionIdx, inNodeMap := nodeMap[sqlInstanceID]
 		if !inNodeMap {
@@ -1197,6 +1210,13 @@ func (dsp *DistSQLPlanner) partitionSpansTenant(
 		}
 		partition := &partitions[partitionIdx]
 		partition.Spans = append(partition.Spans, span)
+		lastIdx = partitionIdx
+	}
+	// If spans were only assigned to one SQL instance, then assign them all to
+	// the gateway instance. The primary reason is to avoid an extra hop.
+	// TODO(harding): Don't do this if using an instance in another locality.
+	if len(partitions) == 1 && partitions[0].SQLInstanceID != dsp.gatewaySQLInstanceID {
+		partitions[0].SQLInstanceID = dsp.gatewaySQLInstanceID
 	}
 	return partitions, nil
 }
@@ -2902,6 +2922,20 @@ func (dsp *DistSQLPlanner) createPhysPlanForPlanNode(
 
 	switch n := node.(type) {
 	// Keep these cases alphabetized, please!
+	case *createStatsNode:
+		if n.runAsJob {
+			plan, err = dsp.wrapPlan(ctx, planCtx, n, false /* allowPartialDistribution */)
+		} else {
+			// Create a job record but don't actually start the job.
+			var record *jobs.Record
+			record, err = n.makeJobRecord(ctx)
+			if err != nil {
+				return nil, err
+			}
+			plan, err = dsp.createPlanForCreateStats(ctx, planCtx, 0, /* jobID */
+				record.Details.(jobspb.CreateStatsDetails))
+		}
+
 	case *distinctNode:
 		plan, err = dsp.createPlanForDistinct(ctx, planCtx, n)
 
@@ -3022,20 +3056,6 @@ func (dsp *DistSQLPlanner) createPhysPlanForPlanNode(
 	case *zigzagJoinNode:
 		plan, err = dsp.createPlanForZigzagJoin(planCtx, n)
 
-	case *createStatsNode:
-		if n.runAsJob {
-			plan, err = dsp.wrapPlan(ctx, planCtx, n, false /* allowPartialDistribution */)
-		} else {
-			// Create a job record but don't actually start the job.
-			var record *jobs.Record
-			record, err = n.makeJobRecord(ctx)
-			if err != nil {
-				return nil, err
-			}
-			plan, err = dsp.createPlanForCreateStats(ctx, planCtx, 0, /* jobID */
-				record.Details.(jobspb.CreateStatsDetails))
-		}
-
 	default:
 		// Can't handle a node? We wrap it and continue on our way.
 		plan, err = dsp.wrapPlan(ctx, planCtx, n, false /* allowPartialDistribution */)

pkg/sql/distsql_running.go

@@ -1255,7 +1255,7 @@ func (dsp *DistSQLPlanner) planAndRunSubquery(
 	).WillDistribute()
 	distribute := DistributionType(DistributionTypeNone)
 	if distributeSubquery {
-		distribute = DistributionTypeSystemTenantOnly
+		distribute = DistributionTypeAlways
 	}
 	subqueryPlanCtx := dsp.NewPlanningCtx(ctx, evalCtx, planner, planner.txn,
 		distribute)
@@ -1602,7 +1602,7 @@ func (dsp *DistSQLPlanner) planAndRunPostquery(
 	).WillDistribute()
 	distribute := DistributionType(DistributionTypeNone)
 	if distributePostquery {
-		distribute = DistributionTypeSystemTenantOnly
+		distribute = DistributionTypeAlways
 	}
 	postqueryPlanCtx := dsp.NewPlanningCtx(ctx, evalCtx, planner, planner.txn,
 		distribute)

pkg/sql/exec_util.go

@@ -1628,9 +1628,6 @@ func getPlanDistribution(
 		return physicalplan.LocalPlan
 	}
 
-	if _, singleTenant := nodeID.OptionalNodeID(); !singleTenant {
-		return physicalplan.LocalPlan
-	}
 	if distSQLMode == sessiondatapb.DistSQLOff {
 		return physicalplan.LocalPlan
 	}

pkg/sql/explain_vec.go

@@ -109,7 +109,7 @@ func newPlanningCtxForExplainPurposes(
 ) *PlanningCtx {
 	distribute := DistributionType(DistributionTypeNone)
 	if distribution.WillDistribute() {
-		distribute = DistributionTypeSystemTenantOnly
+		distribute = DistributionTypeAlways
 	}
 	planCtx := distSQLPlanner.NewPlanningCtx(params.ctx, params.extendedEvalCtx,
 		params.p, params.p.txn, distribute)

pkg/sql/logictest/logic.go

@@ -820,10 +820,11 @@ var logicTestConfigs = []testClusterConfig{
 		// logictest command.
 		// To run a logic test with this config as a directive, run:
 		// make test PKG=./pkg/ccl/logictestccl TESTS=TestTenantLogic//<test_name>
-		name:        threeNodeTenantConfigName,
-		numNodes:    3,
-		useTenant:   true,
-		isCCLConfig: true,
+		name:                threeNodeTenantConfigName,
+		numNodes:            3,
+		useTenant:           true,
+		isCCLConfig:         true,
+		overrideDistSQLMode: "on",
 	},
 	// Regions and zones below are named deliberately, and contain "-"'s to be reflective
 	// of the naming convention in public clouds.  "-"'s are handled differently in SQL
@@ -3403,7 +3404,7 @@ func (t *logicTest) verifyError(
 		} else {
 			newErr := errors.Errorf("%s: %s\nexpected error code %q, but found success",
 				pos, sql, expectErrCode)
-			return (err != nil), newErr
+			return err != nil, newErr
 		}
 	}
 	return true, nil

pkg/sql/logictest/testdata/logic_test/distsql_automatic_stats

@@ -1,8 +1,4 @@
-# LogicTest: !metamorphic !3node-tenant
-
-# Note: this test is disabled on 3node-tenant because it sometimes causes one of
-# the UPDATE statements below (where we update more than 20% of the table) to be
-# flaky. See comments there for details.
+# LogicTest: !metamorphic
 
 # Disable automatic stats
 statement ok
@@ -79,6 +75,7 @@ UPDATE data SET d = 12 WHERE d = 10
 # For some reason, 3node-tenant occasionally splits the UPDATE into 4 pieces,
 # with each one affecting at most 88 rows. Since 88 < 205, the refresh is not
 # guaranteed, making this test flaky.
+skipif config 3node-tenant
 query TTIII colnames,rowsort,retry
 SELECT DISTINCT ON (column_names) statistics_name, column_names, row_count, distinct_count, null_count
 FROM [SHOW STATISTICS FOR TABLE data] ORDER BY column_names ASC, created DESC

pkg/sql/logictest/testdata/logic_test/distsql_tenant

@@ -0,0 +1,22 @@
+# LogicTest: 3node-tenant
+
+statement ok
+CREATE TABLE t (k INT PRIMARY KEY, v INT, w INT, x INT,
+   FAMILY fam_0 (k),
+   FAMILY fam_1 (x),
+   FAMILY fam_2 (v, w)
+)
+
+statement ok
+INSERT INTO t VALUES (23, 1, 2, 3), (34, 3, 4, 8);
+
+query IIII
+SELECT * FROM t WHERE k < 10 OR (k > 20 AND k < 29) OR k > 40
+----
+23 1 2 3
+
+query II
+SELECT v, w FROM t WHERE k = 23
+----
+1  2
+

pkg/sql/logictest/testdata/logic_test/role

@@ -292,15 +292,15 @@ role_name  member     is_admin
 testrole   testuser   true
 testrole   testuser2  true
 
-query TTB colnames
+query TTB colnames,rowsort
 SHOW GRANTS ON ROLE FOR testuser, testuser2
 ----
 role_name  member     is_admin
 admin      testuser   false
 testrole   testuser   true
 testrole   testuser2  true
 
-query TTB colnames
+query TTB colnames,rowsort
 SHOW GRANTS ON ROLE admin, testrole FOR root, testuser2
 ----
 role_name  member     is_admin

pkg/sql/opt/exec/execbuilder/BUILD.bazel

@@ -80,3 +80,9 @@ go_test(
         "//pkg/util/randutil",
     ],
 )
+
+filegroup(
+    name = "testdata",
+    srcs = glob(["testdata/**"]),
+    visibility = ["//pkg/ccl/logictestccl:__subpackages__"],
+)