From a63f83b619401ff16fce2f2d1fa71c3f0ed411b8 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 24 Aug 2023 16:41:55 -0700
Subject: [PATCH] server: use params for temp storage and settings for default
 test tenant

This commit makes it so that we use the settings and temp storage from
the server params (if available) when starting up the default test
tenant. The temp storage fix allows us to adjust one test to work with
the test tenant.

Release note: None
---
 pkg/server/testserver.go         | 11 +++++++++--
 pkg/sql/colflow/BUILD.bazel      |  2 ++
 pkg/sql/colflow/draining_test.go | 13 +++++++++++++
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/pkg/server/testserver.go b/pkg/server/testserver.go
index fb120df63367..dcb21a01bcd8 100644
--- a/pkg/server/testserver.go
+++ b/pkg/server/testserver.go
@@ -599,7 +599,14 @@ func (ts *testServer) maybeStartDefaultTestTenant(ctx context.Context) error {
 		return nil // nolint:returnerrcheck
 	}
 
-	tempStorageConfig := base.DefaultTestTempStorageConfig(cluster.MakeTestingClusterSettings())
+	tenantSettings := ts.params.Settings
+	if tenantSettings == nil {
+		tenantSettings = cluster.MakeTestingClusterSettings()
+	}
+	tempStorageConfig := ts.params.TempStorageConfig
+	if tempStorageConfig.Settings == nil {
+		tempStorageConfig = base.DefaultTestTempStorageConfig(tenantSettings)
+	}
 	params := base.TestTenantArgs{
 		// Currently, all the servers leverage the same tenant ID. We may
 		// want to change this down the road, for more elaborate testing.
@@ -614,7 +621,7 @@ func (ts *testServer) maybeStartDefaultTestTenant(ctx context.Context) error {
 		SSLCertsDir:               ts.params.SSLCertsDir,
 		TestingKnobs:              ts.params.Knobs,
 		StartDiagnosticsReporting: ts.params.StartDiagnosticsReporting,
-		Settings:                  ts.params.Settings,
+		Settings:                  tenantSettings,
 	}
 
 	// Since we're creating a tenant, it doesn't make sense to pass through the
diff --git a/pkg/sql/colflow/BUILD.bazel b/pkg/sql/colflow/BUILD.bazel
index 6603f085212c..a3e6461475d0 100644
--- a/pkg/sql/colflow/BUILD.bazel
+++ b/pkg/sql/colflow/BUILD.bazel
@@ -85,11 +85,13 @@ go_test(
     embed = [":colflow"],
     deps = [
         "//pkg/base",
+        "//pkg/ccl",
         "//pkg/col/coldata",
         "//pkg/col/coldataext",
         "//pkg/col/coldatatestutils",
         "//pkg/keys",
         "//pkg/kv",
+        "//pkg/multitenant/tenantcapabilities",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",
diff --git a/pkg/sql/colflow/draining_test.go b/pkg/sql/colflow/draining_test.go
index 987a1426c282..f92054ab21e5 100644
--- a/pkg/sql/colflow/draining_test.go
+++ b/pkg/sql/colflow/draining_test.go
@@ -17,7 +17,10 @@ import (
 	"testing"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/ccl"
+	"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
@@ -32,6 +35,7 @@ import (
 func TestDrainingAfterRemoteError(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
+	defer ccl.TestingEnableEnterprise()()
 
 	ctx := context.Background()
 	st := cluster.MakeTestingClusterSettings()
@@ -59,6 +63,15 @@ func TestDrainingAfterRemoteError(t *testing.T) {
 	tc := testcluster.StartTestCluster(t, 2 /* nodes */, args)
 	defer tc.Stopper().Stop(ctx)
 
+	if tc.StartedDefaultTestTenant() {
+		systemSqlDB := tc.Server(0).SystemLayer().SQLConn(t, "system")
+		_, err := systemSqlDB.Exec(`ALTER TENANT [$1] GRANT CAPABILITY can_admin_relocate_range=true`, serverutils.TestTenantID().ToUint64())
+		require.NoError(t, err)
+		serverutils.WaitForTenantCapabilities(t, tc.Server(0), serverutils.TestTenantID(), map[tenantcapabilities.ID]string{
+			tenantcapabilities.CanAdminRelocateRange: "true",
+		}, "")
+	}
+
 	// Create two tables, one with small values, and another with large rows.
 	// Relocate the range for the small table to node 2.
 	conn := tc.ServerConn(0)