sqlliveness: embed current region into session ID

Follow up on cockroachdb#90408 and cockroachdb#91019.

Previously, the sqlliveness subsystem was using enum.One when constructing the
session ID since it doesn't have the regional data yet. This commit implements
that missing part of it by ensuring that the regional representation is
plumbed all the way to the sqlliveness subsystem whenever the SQL pod is
started, and use it to construct the session ID. This will enable our REGIONAL
BY ROW work to put the data in the right region.

The multi-region enum for the system database will be read during startup for
that to work. Since doing this already requires loading the system DB's
descriptor (which indirectly tests whether the system DB has been bootstrapped)
for the tenant, we can remove the call to checkTenantExists.

Epic: None

Release note: None

pkg/ccl/multiregionccl/BUILD.bazel

@@ -69,6 +69,7 @@ go_test(
         "//pkg/sql/parser",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
+        "//pkg/sql/sqlliveness",
         "//pkg/sql/sqlliveness/slstorage",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/tests",

pkg/ccl/multiregionccl/multiregion_system_table_test.go

@@ -18,16 +18,22 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/ccl/multiregionccl/multiregionccltestutils"
 	"github.com/cockroachdb/cockroach/pkg/keys"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/enum"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slstorage"
+	"github.com/cockroachdb/cockroach/pkg/testutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/cockroach/pkg/util/uuid"
+	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 )
 
@@ -138,3 +144,91 @@ func TestRbrSqllivenessTable(t *testing.T) {
 		require.Equal(t, writeExpiration, readExpiration)
 	})
 }
+
+func TestTenantStartupWithMultiRegionEnum(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	defer envutil.TestSetEnv(t, "COCKROACH_MR_SYSTEM_DATABASE", "1")()
+
+	tc, db, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
+		t, 3 /*numServers*/, base.TestingKnobs{},
+	)
+	defer cleanup()
+	sqlDB := sqlutils.MakeSQLRunner(db)
+
+	tenID := roachpb.MakeTenantID(10)
+	ten, tSQL := serverutils.StartTenant(t, tc.Server(0), base.TestTenantArgs{
+		TenantID: tenID,
+		Locality: roachpb.Locality{
+			Tiers: []roachpb.Tier{
+				{Key: "region", Value: "us-east1"},
+			},
+		},
+	})
+	defer tSQL.Close()
+	tenSQLDB := sqlutils.MakeSQLRunner(tSQL)
+
+	// Set cluster setting override on the host, and wait for it to propagate.
+	sqlDB.Exec(t, fmt.Sprintf("ALTER TENANT $1 SET CLUSTER SETTING %s = 'true'",
+		sql.SecondaryTenantsMultiRegionAbstractionsEnabledSettingName), tenID.ToUint64())
+	testutils.SucceedsSoon(t, func() error {
+		var currentVal string
+		tenSQLDB.QueryRow(t, fmt.Sprintf("SHOW CLUSTER SETTING %s",
+			sql.SecondaryTenantsMultiRegionAbstractionsEnabledSettingName)).Scan(&currentVal)
+		if currentVal != "true" {
+			return errors.New("waiting for cluster setting to be set to true")
+		}
+		return nil
+	})
+
+	// Update system database with regions.
+	tenSQLDB.Exec(t, `SET descriptor_validation = read_only`)
+	tenSQLDB.Exec(t, `ALTER DATABASE system SET PRIMARY REGION "us-east1"`)
+	tenSQLDB.Exec(t, `SET descriptor_validation = on`)
+	tenSQLDB.Exec(t, `ALTER DATABASE system ADD REGION "us-east2"`)
+	tenSQLDB.Exec(t, `ALTER DATABASE system ADD REGION "us-east3"`)
+
+	ten2, tSQL2 := serverutils.StartTenant(t, tc.Server(2), base.TestTenantArgs{
+		TenantID: tenID,
+		Existing: true,
+		Locality: roachpb.Locality{
+			Tiers: []roachpb.Tier{
+				{Key: "region", Value: "us-east3"},
+			},
+		},
+	})
+	defer tSQL2.Close()
+	tenSQLDB2 := sqlutils.MakeSQLRunner(tSQL2)
+
+	// The sqlliveness entry created by the first SQL server has enum.One as the
+	// region as the system database hasn't been updated when it first started.
+	var sessionID string
+	tenSQLDB2.QueryRow(t, `SELECT session_id FROM system.sql_instances WHERE id = $1`,
+		ten.SQLInstanceID()).Scan(&sessionID)
+	region, id, err := slstorage.UnsafeDecodeSessionID(sqlliveness.SessionID(sessionID))
+	require.NoError(t, err)
+	require.NotNil(t, id)
+	require.Equal(t, enum.One, region)
+
+	// Ensure that the sqlliveness entry created by the second SQL server has
+	// the right region and session UUID.
+	tenSQLDB2.QueryRow(t, `SELECT session_id FROM system.sql_instances WHERE id = $1`,
+		ten2.SQLInstanceID()).Scan(&sessionID)
+	region, id, err = slstorage.UnsafeDecodeSessionID(sqlliveness.SessionID(sessionID))
+	require.NoError(t, err)
+	require.NotNil(t, id)
+	require.NotEqual(t, enum.One, region)
+
+	rows := tenSQLDB2.Query(t, `SELECT crdb_region, session_uuid FROM system.sqlliveness`)
+	defer rows.Close()
+	livenessMap := map[string][]byte{}
+	for rows.Next() {
+		var region, sessionUUID string
+		require.NoError(t, rows.Scan(&region, &sessionUUID))
+		livenessMap[sessionUUID] = []byte(region)
+	}
+	require.NoError(t, rows.Err())
+	r, ok := livenessMap[string(id)]
+	require.True(t, ok)
+	require.Equal(t, r, region)
+}

pkg/ccl/serverccl/server_sql_test.go

@@ -198,8 +198,7 @@ func TestNonExistentTenant(t *testing.T) {
 			DisableCreateTenant: true,
 			SkipTenantCheck:     true,
 		})
-	require.Error(t, err)
-	require.Equal(t, "system DB uninitialized, check if tenant is non existent", err.Error())
+	require.EqualError(t, err, `database "[1]" does not exist`)
 }
 
 // TestTenantRowIDs confirms `unique_rowid()` works as expected in a

pkg/server/BUILD.bazel

@@ -165,6 +165,7 @@ go_library(
         "//pkg/sql/cacheutil",
         "//pkg/sql/catalog/bootstrap",
         "//pkg/sql/catalog/catalogkeys",
+        "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/catsessiondata",
         "//pkg/sql/catalog/colinfo",
         "//pkg/sql/catalog/descidgen",

pkg/server/server_sql.go

@@ -66,12 +66,12 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/sql/cacheutil"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkeys"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catsessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descidgen"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/hydrateddesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
-	"github.com/cockroachdb/cockroach/pkg/sql/catalog/systemschema"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexec"
 	"github.com/cockroachdb/cockroach/pkg/sql/consistencychecker"
 	"github.com/cockroachdb/cockroach/pkg/sql/contention"
@@ -1256,28 +1256,6 @@ func newSQLServer(ctx context.Context, cfg sqlServerArgs) (*SQLServer, error) {
 	}, nil
 }
 
-// Checks if tenant exists. This function does a very superficial check to see if the system db
-// has been bootstrapped for the tenant. This is not a complete check and is only sufficient
-// to be used in the dev environment.
-func checkTenantExists(ctx context.Context, codec keys.SQLCodec, db *kv.DB) error {
-	if codec.ForSystemTenant() {
-		return errors.AssertionFailedf("asked to check for tenant but system codec specified")
-	}
-
-	key := catalogkeys.MakeDatabaseNameKey(codec, systemschema.SystemDatabaseName)
-	result, err := db.Get(ctx, key)
-	if err != nil {
-		return err
-	}
-	if result.Value == nil || result.ValueInt() != keys.SystemDatabaseID {
-		return errors.New("system DB uninitialized, check if tenant is non existent")
-	}
-	// Tenant has been confirmed to be bootstrapped successfully
-	// as the system database, which is a part of the bootstrap data for
-	// a tenant keyspace, exists in the namespace table.
-	return nil
-}
-
 func (s *SQLServer) setInstanceID(
 	ctx context.Context, instanceID base.SQLInstanceID, sessionID sqlliveness.SessionID,
 ) error {
@@ -1297,7 +1275,6 @@ func (s *SQLServer) preStart(
 	pgL net.Listener,
 	orphanedLeasesTimeThresholdNanos int64,
 ) error {
-
 	// If necessary, start the tenant proxy first, to ensure all other
 	// components can properly route to KV nodes. The Start method will block
 	// until a connection is established to the cluster and its ID has been
@@ -1306,17 +1283,46 @@ func (s *SQLServer) preStart(
 		if err := s.tenantConnect.Start(ctx); err != nil {
 			return err
 		}
-		// Confirm tenant exists prior to initialization. This is a sanity
-		// check for the dev environment to ensure that a tenant has been
-		// successfully created before attempting to initialize a SQL
-		// server for it.
-		if err := checkTenantExists(ctx, s.execCfg.Codec, s.execCfg.DB); err != nil {
-			return err
+	}
+
+	// Load the multi-region enum by reading the system database's descriptor.
+	// This also serves as a simple check to see if a tenant exist (i.e. by
+	// checking whether the system db has been bootstrapped).
+	var enumReps map[catpb.RegionName][]byte
+	var primaryRegion catpb.RegionName
+	if err := s.internalExecutorFactory.DescsTxn(ctx, s.execCfg.DB, func(
+		ctx context.Context, txn *kv.Txn, descsCol *descs.Collection,
+	) error {
+		enumReps, primaryRegion = nil, ""
+		var err error
+		enumReps, primaryRegion, err = sql.GetRegionEnumRepresentations(
+			ctx, txn, keys.SystemDatabaseID, descsCol)
+		if errors.Is(err, sql.ErrNotMultiRegionDatabase) {
+			return nil
+		}
+		return err
+	}); err != nil {
+		return err
+	}
+
+	var regionPhysicalRep []byte
+	if enumReps != nil {
+		// The primary region will be used if no region was provided through
+		// the locality flag.
+		currentRegion, _ := s.distSQLServer.Locality.Find("region")
+		if currentRegion == "" {
+			currentRegion = primaryRegion.String()
+		}
+		for regionName, phyRep := range enumReps {
+			if regionName == catpb.RegionName(currentRegion) {
+				regionPhysicalRep = phyRep
+				break
+			}
 		}
 	}
 
 	// Start the sql liveness subsystem. We'll need it to get a session.
-	s.sqlLivenessProvider.Start(ctx)
+	s.sqlLivenessProvider.Start(ctx, regionPhysicalRep)
 
 	_, isMixedSQLAndKVNode := s.sqlIDContainer.OptionalNodeID()
 	isTenant := !isMixedSQLAndKVNode