Merge cockroachdb#109482

109482: server: use params for temp storage and settings for default test tenant r=yuzefovich a=yuzefovich 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. Fixes: cockroachdb#109456. Release note: None Co-authored-by: Yahor Yuzefovich <[email protected]>
aadityasondhi · Aug 25, 2023 · 37f97b6 · 37f97b6
2 parents 7cef593 + a63f83b
commit 37f97b6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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)