sql: audit usages of SystemSQLCodec in tests

One of the common modifications we need to make for tests to work with
secondary tenants is using the right `keys.SQLCodec` (as opposed to the
one belonging to the system tenant). This commit audits all tests in
`pkg/sql` directory to ensure that the right codec is passed. Note that
tests still might need further modifications. Also note that tests that
don't start a test server / a test cluster weren't touched (since they
don't benefit from the tenant randomization behavior).

Release note: None

Author: yuzefovich

pkg/server/server_systemlog_gc.go

@@ -8,7 +8,7 @@ package server
 import (
 	"context"
 	"fmt"
-	math_rand "math/rand"
+	"math/rand"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
@@ -288,5 +288,5 @@ func startSystemLogsGC(ctx context.Context, sqlServer *SQLServer) error {
 // jitteredInterval returns a randomly jittered (+/-25%) duration
 // from the interval.
 func jitteredInterval(interval time.Duration) time.Duration {
-	return time.Duration(float64(interval) * (0.75 + 0.5*math_rand.Float64()))
+	return time.Duration(float64(interval) * (0.75 + 0.5*rand.Float64()))
 }

pkg/sql/catalog/bootstrap/splits_test.go

@@ -51,6 +51,9 @@ func TestInitialSplitPoints(t *testing.T) {
 
 		ctx := context.Background()
 		tc := testcluster.StartTestCluster(t, 1, base.TestClusterArgs{
+			ServerArgs: base.TestServerArgs{
+				DefaultTestTenant: base.TestControlsTenantsExplicitly,
+			},
 			ReplicationMode: base.ReplicationManual,
 		})
 		defer tc.Stopper().Stop(ctx)

pkg/sql/check_test.go

@@ -15,7 +15,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobstest"
-	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
@@ -42,12 +41,12 @@ func TestValidateTTLScheduledJobs(t *testing.T) {
 
 	testCases := []struct {
 		desc          string
-		setup         func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.TestServerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID)
+		setup         func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.ApplicationLayerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID)
 		expectedErrRe func(tableID descpb.ID, scheduleID jobspb.ScheduleID) string
 	}{
 		{
 			desc: "not pointing at a valid scheduled job",
-			setup: func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.TestServerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID) {
+			setup: func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.ApplicationLayerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID) {
 				require.NoError(t, sqltestutils.TestingDescsTxn(ctx, s, func(ctx context.Context, txn isql.Txn, col *descs.Collection) (err error) {
 					// We need the collection to read the descriptor from storage for
 					// the subsequent write to succeed.
@@ -66,7 +65,7 @@ func TestValidateTTLScheduledJobs(t *testing.T) {
 		},
 		{
 			desc: "scheduled job points at an different table",
-			setup: func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.TestServerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID) {
+			setup: func(t *testing.T, sqlDB *gosql.DB, kvDB *kv.DB, s serverutils.ApplicationLayerInterface, tableDesc *tabledesc.Mutable, scheduleID jobspb.ScheduleID) {
 				db := s.InternalDB().(isql.DB)
 				require.NoError(t, db.Txn(ctx, func(ctx context.Context, txn isql.Txn) error {
 					schedules := jobs.ScheduledJobTxn(txn)
@@ -103,13 +102,14 @@ func TestValidateTTLScheduledJobs(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
-			s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
-			defer s.Stopper().Stop(ctx)
+			srv, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
+			defer srv.Stopper().Stop(ctx)
+			s := srv.ApplicationLayer()
 
 			_, err := sqlDB.Exec(`CREATE TABLE t () WITH (ttl_expire_after = '10 mins')`)
 			require.NoError(t, err)
 
-			tableDesc := desctestutils.TestingGetMutableExistingTableDescriptor(kvDB, keys.SystemSQLCodec, "defaultdb", "t")
+			tableDesc := desctestutils.TestingGetMutableExistingTableDescriptor(kvDB, s.Codec(), "defaultdb", "t")
 			require.NotNil(t, tableDesc.GetRowLevelTTL())
 			scheduleID := tableDesc.GetRowLevelTTL().ScheduleID
 

pkg/sql/conn_executor_internal_test.go

@@ -274,6 +274,7 @@ func startConnExecutor(
 ) {
 	// A lot of boilerplate for creating a connExecutor.
 	stopper := stop.NewStopper()
+	codec := keys.SystemSQLCodec
 	clock := hlc.NewClockForTesting(nil)
 	factory := kv.MakeMockTxnSenderFactory(
 		func(context.Context, *roachpb.Transaction, *kvpb.BatchRequest,
@@ -297,7 +298,7 @@ func startConnExecutor(
 	})
 	// This pool should never be Stop()ed because, if the test is failing, memory
 	// is not properly released.
-	collectionFactory := descs.NewBareBonesCollectionFactory(st, keys.SystemSQLCodec)
+	collectionFactory := descs.NewBareBonesCollectionFactory(st, codec)
 	registry := stmtdiagnostics.NewRegistry(nil, st)
 	cfg := &ExecutorConfig{
 		AmbientCtx: ambientCtx,
@@ -313,7 +314,7 @@ func startConnExecutor(
 			NodeID:           nodeID,
 			LogicalClusterID: func() uuid.UUID { return uuid.UUID{} },
 		},
-		Codec: keys.SystemSQLCodec,
+		Codec: codec,
 		DistSQLPlanner: NewDistSQLPlanner(
 			ctx, st, 1, /* sqlInstanceID */
 			nil, /* rpcCtx */
@@ -339,7 +340,7 @@ func startConnExecutor(
 			nil, /* connHealthCheckerSystem */
 			nil, /* instanceConnHealthChecker */
 			nil, /* sqlInstanceDialer */
-			keys.SystemSQLCodec,
+			codec,
 			nil, /* sqlAddressResolver */
 			clock,
 		),

pkg/sql/crdb_internal_test.go

@@ -1595,8 +1595,8 @@ func TestVirtualPTSTableDeprecated(t *testing.T) {
 			Mode:      ptpb.PROTECT_AFTER,
 			DeprecatedSpans: []roachpb.Span{
 				{
-					Key:    keys.SystemSQLCodec.TablePrefix(42),
-					EndKey: keys.SystemSQLCodec.TablePrefix(42).PrefixEnd(),
+					Key:    s.Codec().TablePrefix(42),
+					EndKey: s.Codec().TablePrefix(42).PrefixEnd(),
 				},
 			},
 			MetaType: "foo",

pkg/sql/database_test.go

@@ -42,8 +42,9 @@ func TestDatabaseHasChildSchemas(t *testing.T) {
 	defer log.Scope(t).Close(t)
 
 	ctx := context.Background()
-	s, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
-	defer s.Stopper().Stop(ctx)
+	srv, sqlDB, kvDB := serverutils.StartServer(t, base.TestServerArgs{})
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
 
 	// Create a database and schema.
 	if _, err := sqlDB.Exec(`
@@ -55,7 +56,7 @@ CREATE SCHEMA sc;
 	}
 
 	// Now get the database descriptor from disk.
-	db := desctestutils.TestingGetDatabaseDescriptor(kvDB, keys.SystemSQLCodec, "d")
+	db := desctestutils.TestingGetDatabaseDescriptor(kvDB, s.Codec(), "d")
 	if db.GetSchemaID("sc") == descpb.InvalidID {
 		t.Fatal("expected to find child schema sc in db")
 	}
@@ -65,7 +66,7 @@ CREATE SCHEMA sc;
 		t.Fatal(err)
 	}
 
-	db = desctestutils.TestingGetDatabaseDescriptor(kvDB, keys.SystemSQLCodec, "d")
+	db = desctestutils.TestingGetDatabaseDescriptor(kvDB, s.Codec(), "d")
 	if db.GetSchemaID("sc2") == descpb.InvalidID {
 		t.Fatal("expected to find child schema sc2 in db")
 	}

pkg/sql/distsql_physical_planner_test.go

@@ -1372,15 +1372,16 @@ func TestPartitionSpans(t *testing.T) {
 	// We need a mock Gossip to contain addresses for the nodes. Otherwise the
 	// DistSQLPlanner will not plan flows on them.
 	ctx := context.Background()
-	s := serverutils.StartServerOnly(t, base.TestServerArgs{
+	srv := serverutils.StartServerOnly(t, base.TestServerArgs{
 		Knobs: base.TestingKnobs{
 			DistSQL: &execinfra.TestingKnobs{
 				MinimumNumberOfGatewayPartitions: 1,
 			},
 		},
 	})
-	defer s.Stopper().Stop(ctx)
-	mockGossip := gossip.NewTest(roachpb.NodeID(1), s.Stopper(), metric.NewRegistry())
+	defer srv.Stopper().Stop(ctx)
+	s := srv.ApplicationLayer()
+	mockGossip := gossip.NewTest(roachpb.NodeID(1), srv.Stopper(), metric.NewRegistry())
 	var nodeDescs []*roachpb.NodeDescriptor
 	mockInstances := make(mockAddressResolver)
 	for i := 1; i <= 10; i++ {
@@ -1459,7 +1460,7 @@ func TestPartitionSpans(t *testing.T) {
 						TestingKnobs: execinfra.TestingKnobs{MinimumNumberOfGatewayPartitions: 1},
 					},
 				},
-				codec:     keys.SystemSQLCodec,
+				codec:     s.Codec(),
 				nodeDescs: mockGossip,
 			}
 
@@ -1468,7 +1469,7 @@ func TestPartitionSpans(t *testing.T) {
 				require.NoError(t, locFilter.Set(tc.locFilter))
 			}
 			evalCtx := &eval.Context{
-				Codec: keys.SystemSQLCodec,
+				Codec: s.Codec(),
 				SessionDataStack: sessiondata.NewStack(&sessiondata.SessionData{
 					SessionData: sessiondatapb.SessionData{
 						DistsqlPlanGatewayBias: 2,
@@ -1761,6 +1762,7 @@ func TestPartitionSpansSkipsNodesNotInGossip(t *testing.T) {
 
 	stopper := stop.NewStopper()
 	defer stopper.Stop(context.Background())
+	codec := keys.SystemSQLCodec
 
 	mockGossip := gossip.NewTest(roachpb.NodeID(1), stopper, metric.NewRegistry())
 	var nodeDescs []*roachpb.NodeDescriptor
@@ -1815,14 +1817,14 @@ func TestPartitionSpansSkipsNodesNotInGossip(t *testing.T) {
 				return true
 			},
 		},
-		codec: keys.SystemSQLCodec,
+		codec: codec,
 	}
 
 	ctx := context.Background()
 	// This test is specific to gossip-based planning.
 	useGossipPlanning.Override(ctx, &st.SV, true)
 	planCtx := dsp.NewPlanningCtx(
-		ctx, &extendedEvalContext{Context: eval.Context{Codec: keys.SystemSQLCodec, Settings: st}},
+		ctx, &extendedEvalContext{Context: eval.Context{Codec: codec, Settings: st}},
 		nil /* planner */, nil /* txn */, FullDistribution,
 	)
 	partitions, err := dsp.PartitionSpans(ctx, planCtx, roachpb.Spans{span}, PartitionSpansBoundDefault)

pkg/sql/distsql_plan_backfill_test.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
-	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/desctestutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
@@ -55,35 +54,35 @@ func TestDistBackfill(t *testing.T) {
 			},
 		})
 	defer tc.Stopper().Stop(context.Background())
-	cdb := tc.Server(0).DB()
+	s := tc.ApplicationLayer(0)
+	conn, cdb := s.SQLConn(t, serverutils.DBName("test")), s.DB()
 
 	sqlutils.CreateTable(
-		t, tc.ServerConn(0), "numtosquare", "x INT PRIMARY KEY, xsquared INT",
+		t, conn, "numtosquare", "x INT PRIMARY KEY, xsquared INT",
 		n,
 		sqlutils.ToRowFn(sqlutils.RowIdxFn, func(row int) tree.Datum {
 			return tree.NewDInt(tree.DInt(row * row))
 		}),
 	)
 
 	sqlutils.CreateTable(
-		t, tc.ServerConn(0), "numtostr", "y INT PRIMARY KEY, str STRING",
+		t, conn, "numtostr", "y INT PRIMARY KEY, str STRING",
 		n*n,
 		sqlutils.ToRowFn(sqlutils.RowIdxFn, sqlutils.RowEnglishFn),
 	)
 	// Split the table into multiple ranges.
-	descNumToStr := desctestutils.TestingGetPublicTableDescriptor(cdb, keys.SystemSQLCodec, "test", "numtostr")
+	descNumToStr := desctestutils.TestingGetPublicTableDescriptor(cdb, s.Codec(), "test", "numtostr")
 	var sps []serverutils.SplitPoint
 	// for i := 1; i <= numNodes-1; i++ {
 	for i := numNodes - 1; i > 0; i-- {
 		sps = append(sps, serverutils.SplitPoint{TargetNodeIdx: i, Vals: []interface{}{n * n / numNodes * i}})
 	}
 	tc.SplitTable(t, descNumToStr, sps)
 
-	db := tc.ServerConn(0)
-	db.SetMaxOpenConns(1)
-	r := sqlutils.MakeSQLRunner(db)
+	conn.SetMaxOpenConns(1)
+	r := sqlutils.MakeSQLRunner(conn)
 	r.Exec(t, "SET DISTSQL = OFF")
-	if _, err := tc.ServerConn(0).Exec(`CREATE INDEX foo ON numtostr (str)`); err != nil {
+	if _, err := conn.Exec(`CREATE INDEX foo ON numtostr (str)`); err != nil {
 		t.Fatal(err)
 	}
 	r.Exec(t, "SET DISTSQL = ALWAYS")