sqlliveness: embed current region into session ID

pkg/ccl/multiregionccl/BUILD.bazel

@@ -31,6 +31,7 @@ go_test(
         "multiregion_system_table_test.go",
         "multiregion_test.go",
         "region_test.go",
+        "region_util_test.go",
         "regional_by_row_system_database_test.go",
         "regional_by_row_test.go",
         "roundtrips_test.go",
@@ -67,11 +68,13 @@ go_test(
         "//pkg/sql/catalog/desctestutils",
         "//pkg/sql/catalog/tabledesc",
         "//pkg/sql/catalog/typedesc",
+        "//pkg/sql/enum",
         "//pkg/sql/execinfra",
         "//pkg/sql/parser",
         "//pkg/sql/rowenc",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sqlliveness",
+        "//pkg/sql/sqlliveness/slstorage",
         "//pkg/sql/sqltestutils",
         "//pkg/sql/tests",
         "//pkg/sql/types",

pkg/ccl/multiregionccl/multiregion_system_table_test.go

@@ -26,8 +26,10 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
+	"github.com/cockroachdb/cockroach/pkg/sql/enum"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlliveness/slstorage"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -176,10 +178,10 @@ func TestMrSystemDatabase(t *testing.T) {
 
 		t.Run("InUse", func(t *testing.T) {
 			query := `
-				SELECT id, addr, session_id, locality, crdb_region
-				FROM system.sql_instances
-				WHERE session_id IS NOT NULL
-			`
+                SELECT id, addr, session_id, locality, crdb_region
+                FROM system.sql_instances
+                WHERE session_id IS NOT NULL
+            `
 			rows := tDB.Query(t, query)
 			require.True(t, rows.Next())
 			for {
@@ -207,10 +209,10 @@ func TestMrSystemDatabase(t *testing.T) {
 
 		t.Run("Preallocated", func(t *testing.T) {
 			query := `
-				SELECT id, addr, session_id, locality, crdb_region
-				FROM system.sql_instances
-				WHERE session_id IS NULL
-			`
+                SELECT id, addr, session_id, locality, crdb_region
+                FROM system.sql_instances
+                WHERE session_id IS NULL
+            `
 			rows := tDB.Query(t, query)
 			require.True(t, rows.Next())
 			for {
@@ -234,3 +236,82 @@ func TestMrSystemDatabase(t *testing.T) {
 		})
 	})
 }
+
+func TestTenantStartupWithMultiRegionEnum(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	defer envutil.TestSetEnv(t, "COCKROACH_MR_SYSTEM_DATABASE", "1")()
+
+	// Enable settings required for configuring a tenant's system database as multi-region.
+	cs := cluster.MakeTestingClusterSettings()
+	sql.SecondaryTenantsMultiRegionAbstractionsEnabled.Override(context.Background(), &cs.SV, true)
+	sql.SecondaryTenantZoneConfigsEnabled.Override(context.Background(), &cs.SV, true)
+
+	tc, _, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
+		t, 3 /*numServers*/, base.TestingKnobs{}, multiregionccltestutils.WithSettings(cs),
+	)
+	defer cleanup()
+
+	tenID := roachpb.MustMakeTenantID(10)
+	ten, tSQL := serverutils.StartTenant(t, tc.Server(0), base.TestTenantArgs{
+		Settings: cs,
+		TenantID: tenID,
+		Locality: roachpb.Locality{
+			Tiers: []roachpb.Tier{
+				{Key: "region", Value: "us-east1"},
+			},
+		},
+	})
+	defer tSQL.Close()
+	tenSQLDB := sqlutils.MakeSQLRunner(tSQL)
+
+	// Update system database with regions.
+	tenSQLDB.Exec(t, `ALTER DATABASE system SET PRIMARY REGION "us-east1"`)
+	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{
+		Settings: cs,
+		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/multiregionccl/region_util_test.go

@@ -0,0 +1,108 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Licensed as a CockroachDB Enterprise file under the Cockroach Community
+// License (the "License"); you may not use this file except in compliance with
+// the License. You may obtain a copy of the License at
+//
+//     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
+
+package multiregionccl_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/ccl/multiregionccl/multiregionccltestutils"
+	"github.com/cockroachdb/cockroach/pkg/kv"
+	"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/catalog/descs"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
+)
+
+// TestGetLocalityRegionEnumPhysicalRepresentation is in the ccl package since
+// it utilizes adding regions to a database.
+func TestGetLocalityRegionEnumPhysicalRepresentation(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	tc, sqlDB, cleanup := multiregionccltestutils.TestingCreateMultiRegionCluster(
+		t, 3 /* numServers */, base.TestingKnobs{})
+	defer cleanup()
+
+	tDB := sqlutils.MakeSQLRunner(sqlDB)
+	tDB.Exec(t, `CREATE DATABASE foo PRIMARY REGION "us-east1" REGIONS "us-east1", "us-east2", "us-east3"`)
+
+	s0 := tc.ServerTyped(0)
+	ief := s0.InternalExecutorFactory().(descs.TxnManager)
+	dbID := sqlutils.QueryDatabaseID(t, sqlDB, "foo")
+
+	t.Run("with locality that exists", func(t *testing.T) {
+		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
+			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{
+				Tiers: []roachpb.Tier{{Key: "region", Value: "us-east2"}},
+			},
+		)
+		require.NoError(t, err)
+
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		require.NotEmpty(t, enumMembers)
+		require.Equal(t, enumMembers["us-east2"], regionEnum)
+	})
+
+	t.Run("with non-existent locality", func(t *testing.T) {
+		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
+			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{
+				Tiers: []roachpb.Tier{{Key: "region", Value: "europe-west1"}},
+			},
+		)
+		require.NoError(t, err)
+
+		// Fallback to primary region if the locality is provided, but non-existent.
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		require.NotEmpty(t, enumMembers)
+		require.Equal(t, enumMembers["us-east1"], regionEnum)
+	})
+
+	t.Run("without locality", func(t *testing.T) {
+		regionEnum, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
+			ctx, ief, s0.DB(), descpb.ID(dbID), roachpb.Locality{})
+		require.NoError(t, err)
+
+		// Fallback to primary region is locality information is missing.
+		enumMembers := getEnumMembers(t, ctx, tc.Server(0), descpb.ID(dbID))
+		require.NotEmpty(t, enumMembers)
+		require.Equal(t, enumMembers["us-east1"], regionEnum)
+	})
+}
+
+func getEnumMembers(
+	t *testing.T, ctx context.Context, ts serverutils.TestServerInterface, dbID descpb.ID,
+) map[string][]byte {
+	t.Helper()
+	enumMembers := make(map[string][]byte)
+	err := sql.TestingDescsTxn(ctx, ts, func(ctx context.Context, txn *kv.Txn, descsCol *descs.Collection) error {
+		_, dbDesc, err := descsCol.GetImmutableDatabaseByID(ctx, txn, dbID,
+			tree.DatabaseLookupFlags{Required: true})
+		require.NoError(t, err)
+		regionEnumID, err := dbDesc.MultiRegionEnumID()
+		require.NoError(t, err)
+		regionEnumDesc, err := descsCol.GetImmutableTypeByID(ctx, txn, regionEnumID,
+			tree.ObjectLookupFlags{CommonLookupFlags: tree.CommonLookupFlags{Required: true}})
+		require.NoError(t, err)
+		for ord := 0; ord < regionEnumDesc.NumEnumMembers(); ord++ {
+			enumMembers[regionEnumDesc.GetMemberLogicalRepresentation(ord)] = regionEnumDesc.GetMemberPhysicalRepresentation(ord)
+		}
+		return nil
+	})
+	require.NoError(t, err)
+	return enumMembers
+}

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/server_sql.go

@@ -71,7 +71,6 @@ import (
 	"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 +1255,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 +1274,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 +1282,19 @@ 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).
+	regionPhysicalRep, err := sql.GetLocalityRegionEnumPhysicalRepresentation(
+		ctx, s.internalExecutorFactory, s.execCfg.DB, keys.SystemDatabaseID, s.distSQLServer.Locality)
+	if err != nil && !errors.Is(err, sql.ErrNotMultiRegionDatabase) {
+		return err
 	}
 
 	// 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