opt: session setting to enforce queries access only home region rows

Fixes #86228

This commit adds a new session setting, `enforce_home_region`, which
causes queries to error out if they need to talk to regions other than
the gateway region to answer the query.

A home region specifies the region(s) from which consistent reads from a
set of rows in a table can be served locally. The home region for a set
of rows in a multiregion table is determined differently depending on
the type of multiregion table involved:
| Locality | Home Region |
| -------- | ----------- |
| REGIONAL BY ROW | Home region determined by crdb_region column value |
| REGIONAL BY TABLE | All rows share a single home region |
| GLOBAL | Any region can act as the home region |

When `enforce_home_region` is true, and a query has no home region
(for example, reading from different home regions in a REGIONAL BY ROW
table), error code 42899 (`QueryHasNoHomeRegion`) is returned.
When `enforce_home_region` is true, and a query's home region differs
from the gateway region, error code 42898
(`QueryNotRunningInHomeRegion`) is returned.
The error message, in some instances, provides a useful tip on possible
steps to take to allow the query to run entirely in the gateway region,
e.g.,
```
Query is not running in its home region. Try running the query from region 'ap-southeast-2'.
Query has no home region. Try removing the join expression.
Query has no home region. Try accessing only tables in multi-region databases with ZONE survivability.
```

Support for this new session mode is being added in 3 phases.
This commit consists of phase 1, which include only simple static checks
during query compilation for the following allowed cases:
- A scan of a table with `LOCALITY REGIONAL BY TABLE` with primary
region matching the gateway region
- A scan of a table with `LOCALITY GLOBAL`
- A scan of a table with `LOCALITY REGIONAL BY ROW` using only local
constraints (e.g. crdb_region = 'ca-central-1')
- A scan of a table with `LOCALITY REGIONAL BY ROW` using
locality-optimized search.
- A lookup join into a table with `LOCALITY REGIONAL BY ROW` using
locality-optimized lookup.

Only tables in multiregion databases with ZONE survivability may be
scanned without error because with REGION survivability, ranges in a
down region may be served non-local to the gateway region, so are not
guaranteed to have low latency.

Note that locality-optimized search and lookup join are not guaranteed
to scan no remote rows, but are still allowed.

Release note (sql change): A new session setting, enforce_home_region,
is added, which when true causes queries which may scan rows via a
database connection outside of the query's home region to error out.
Also, only tables in multiregion databases with ZONE survivability may
be scanned without error when this setting is true because with REGION
survivability, ranges in a down region may be served non-local to the
gateway region, so are not guaranteed to have low latency.

pkg/sql/exec_util.go

@@ -3281,6 +3281,10 @@ func (m *sessionDataMutator) SetTroubleshootingModeEnabled(val bool) {
 	m.data.TroubleshootingMode = val
 }
 
+func (m *sessionDataMutator) SetEnforceHomeRegion(val bool) {
+	m.data.EnforceHomeRegion = val
+}
+
 // Utility functions related to scrubbing sensitive information on SQL Stats.
 
 // quantizeCounts ensures that the Count field in the

pkg/sql/faketreeeval/BUILD.bazel

@@ -10,6 +10,7 @@ go_library(
         "//pkg/clusterversion",
         "//pkg/security/username",
         "//pkg/sql/catalog/catpb",
+        "//pkg/sql/catalog/descpb",
         "//pkg/sql/parser",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",

pkg/sql/faketreeeval/evalctx.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
@@ -124,29 +125,29 @@ var _ eval.RegionOperator = &DummyRegionOperator{}
 var errRegionOperator = unimplemented.NewWithIssue(42508,
 	"cannot evaluate scalar expressions containing region operations in this context")
 
-// CurrentDatabaseRegionConfig is part of the eval.DatabaseCatalog interface.
+// CurrentDatabaseRegionConfig is part of the eval.RegionOperator interface.
 func (so *DummyRegionOperator) CurrentDatabaseRegionConfig(
 	_ context.Context,
 ) (eval.DatabaseRegionConfig, error) {
 	return nil, errors.WithStack(errRegionOperator)
 }
 
-// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.DatabaseCatalog interface.
+// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.RegionOperator interface.
 func (so *DummyRegionOperator) ValidateAllMultiRegionZoneConfigsInCurrentDatabase(
 	_ context.Context,
 ) error {
 	return errors.WithStack(errRegionOperator)
 }
 
-// ResetMultiRegionZoneConfigsForTable is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForTable is part of the eval.RegionOperator
 // interface.
 func (so *DummyRegionOperator) ResetMultiRegionZoneConfigsForTable(
 	_ context.Context, id int64,
 ) error {
 	return errors.WithStack(errRegionOperator)
 }
 
-// ResetMultiRegionZoneConfigsForDatabase is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForDatabase is part of the eval.RegionOperator
 // interface.
 func (so *DummyRegionOperator) ResetMultiRegionZoneConfigsForDatabase(
 	_ context.Context, id int64,
@@ -310,28 +311,28 @@ var _ eval.Planner = &DummyEvalPlanner{}
 var errEvalPlanner = pgerror.New(pgcode.ScalarOperationCannotRunWithoutFullSessionContext,
 	"cannot evaluate scalar expressions using table lookups in this context")
 
-// CurrentDatabaseRegionConfig is part of the eval.DatabaseCatalog interface.
+// CurrentDatabaseRegionConfig is part of the eval.RegionOperator interface.
 func (ep *DummyEvalPlanner) CurrentDatabaseRegionConfig(
 	_ context.Context,
 ) (eval.DatabaseRegionConfig, error) {
 	return nil, errors.WithStack(errEvalPlanner)
 }
 
-// ResetMultiRegionZoneConfigsForTable is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForTable is part of the eval.RegionOperator
 // interface.
 func (ep *DummyEvalPlanner) ResetMultiRegionZoneConfigsForTable(_ context.Context, _ int64) error {
 	return errors.WithStack(errEvalPlanner)
 }
 
-// ResetMultiRegionZoneConfigsForDatabase is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForDatabase is part of the eval.RegionOperator
 // interface.
 func (ep *DummyEvalPlanner) ResetMultiRegionZoneConfigsForDatabase(
 	_ context.Context, _ int64,
 ) error {
 	return errors.WithStack(errEvalPlanner)
 }
 
-// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.DatabaseCatalog interface.
+// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.RegionOperator interface.
 func (ep *DummyEvalPlanner) ValidateAllMultiRegionZoneConfigsInCurrentDatabase(
 	_ context.Context,
 ) error {
@@ -428,6 +429,11 @@ func (ep *DummyEvalPlanner) IsActive(_ context.Context, _ clusterversion.Key) bo
 	return true
 }
 
+// GetMultiregionConfig is part of the eval.Planner interface.
+func (ep *DummyEvalPlanner) GetMultiregionConfig(databaseID descpb.ID) (interface{}, bool) {
+	return nil /* regionConfig */, false
+}
+
 // DummyPrivilegedAccessor implements the tree.PrivilegedAccessor interface by returning errors.
 type DummyPrivilegedAccessor struct{}
 

pkg/sql/importer/import_table_creation.go

@@ -256,29 +256,29 @@ func (i importDatabaseRegionConfig) PrimaryRegionString() string {
 
 var _ eval.DatabaseRegionConfig = &importDatabaseRegionConfig{}
 
-// CurrentDatabaseRegionConfig is part of the eval.DatabaseCatalog interface.
+// CurrentDatabaseRegionConfig is part of the eval.RegionOperator interface.
 func (so *importRegionOperator) CurrentDatabaseRegionConfig(
 	_ context.Context,
 ) (eval.DatabaseRegionConfig, error) {
 	return importDatabaseRegionConfig{primaryRegion: so.primaryRegion}, nil
 }
 
-// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.DatabaseCatalog interface.
+// ValidateAllMultiRegionZoneConfigsInCurrentDatabase is part of the eval.RegionOperator interface.
 func (so *importRegionOperator) ValidateAllMultiRegionZoneConfigsInCurrentDatabase(
 	_ context.Context,
 ) error {
 	return errors.WithStack(errRegionOperator)
 }
 
-// ResetMultiRegionZoneConfigsForTable is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForTable is part of the eval.RegionOperator
 // interface.
 func (so *importRegionOperator) ResetMultiRegionZoneConfigsForTable(
 	_ context.Context, _ int64,
 ) error {
 	return errors.WithStack(errRegionOperator)
 }
 
-// ResetMultiRegionZoneConfigsForDatabase is part of the eval.DatabaseCatalog
+// ResetMultiRegionZoneConfigsForDatabase is part of the eval.RegionOperator
 // interface.
 func (so *importRegionOperator) ResetMultiRegionZoneConfigsForDatabase(
 	_ context.Context, _ int64,

pkg/sql/internal.go

@@ -649,6 +649,9 @@ func (ie *InternalExecutor) execInternal(
 	if ie.sessionDataStack != nil {
 		// TODO(andrei): Properly clone (deep copy) ie.sessionData.
 		sd = ie.sessionDataStack.Top().Clone()
+		// Even if session queries are told to error on non-home region accesses,
+		// internal queries spawned from the same context should never do so.
+		sd.LocalOnlySessionData.EnforceHomeRegion = false
 	} else {
 		sd = ie.s.newSessionData(SessionArgs{})
 	}

pkg/sql/logictest/testdata/logic_test/information_schema

@@ -4690,6 +4690,7 @@ enable_multiregion_placement_policy                   off
 enable_seqscan                                        on
 enable_super_regions                                  off
 enable_zigzag_join                                    on
+enforce_home_region                                   off
 escape_string_warning                                 on
 expect_and_ignore_not_visible_columns_in_copy         off
 experimental_computed_column_rewrites                 ·

pkg/sql/logictest/testdata/logic_test/pg_catalog

@@ -4174,6 +4174,7 @@ enable_multiregion_placement_policy                   off                 NULL
 enable_seqscan                                        on                  NULL      NULL        NULL        string
 enable_super_regions                                  off                 NULL      NULL        NULL        string
 enable_zigzag_join                                    on                  NULL      NULL        NULL        string
+enforce_home_region                                   off                 NULL      NULL        NULL        string
 escape_string_warning                                 on                  NULL      NULL        NULL        string
 expect_and_ignore_not_visible_columns_in_copy         off                 NULL      NULL        NULL        string
 experimental_distsql_planning                         off                 NULL      NULL        NULL        string
@@ -4299,6 +4300,7 @@ enable_multiregion_placement_policy                   off                 NULL
 enable_seqscan                                        on                  NULL  user     NULL      on                  on
 enable_super_regions                                  off                 NULL  user     NULL      off                 off
 enable_zigzag_join                                    on                  NULL  user     NULL      on                  on
+enforce_home_region                                   off                 NULL  user     NULL      off                 off
 escape_string_warning                                 on                  NULL  user     NULL      on                  on
 expect_and_ignore_not_visible_columns_in_copy         off                 NULL  user     NULL      off                 off
 experimental_distsql_planning                         off                 NULL  user     NULL      off                 off
@@ -4420,6 +4422,7 @@ enable_multiregion_placement_policy                   NULL    NULL     NULL
 enable_seqscan                                        NULL    NULL     NULL     NULL        NULL
 enable_super_regions                                  NULL    NULL     NULL     NULL        NULL
 enable_zigzag_join                                    NULL    NULL     NULL     NULL        NULL
+enforce_home_region                                   NULL    NULL     NULL     NULL        NULL
 escape_string_warning                                 NULL    NULL     NULL     NULL        NULL
 expect_and_ignore_not_visible_columns_in_copy         NULL    NULL     NULL     NULL        NULL
 experimental_distsql_planning                         NULL    NULL     NULL     NULL        NULL

pkg/sql/logictest/testdata/logic_test/show_source

@@ -60,6 +60,7 @@ enable_multiregion_placement_policy                   off
 enable_seqscan                                        on
 enable_super_regions                                  off
 enable_zigzag_join                                    on
+enforce_home_region                                   off
 escape_string_warning                                 on
 expect_and_ignore_not_visible_columns_in_copy         off
 experimental_distsql_planning                         off

pkg/sql/opt/BUILD.bazel

@@ -24,7 +24,10 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/server/telemetry",
+        "//pkg/sql/catalog/catpb",
         "//pkg/sql/catalog/colinfo",
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/catalog/multiregion",
         "//pkg/sql/opt/cat",
         "//pkg/sql/opt/partition",
         "//pkg/sql/pgwire/pgcode",

pkg/sql/opt/cat/table.go

@@ -13,6 +13,7 @@ package cat
 import (
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
@@ -149,6 +150,21 @@ type Table interface {
 	// created with the NO DATA option and is yet to be refreshed. Accessing
 	// such a view prior to running refresh returns an error.
 	IsRefreshViewRequired() bool
+
+	// HomeRegion returns the home region of the table, if any, for example if
+	// a table is defined with LOCALITY REGIONAL BY TABLE.
+	HomeRegion() (region string, ok bool)
+
+	// IsGlobalTable returns true if the table is defined with LOCALITY GLOBAL.
+	IsGlobalTable() bool
+
+	// IsRegionalByRow returns true if the table is defined with LOCALITY REGIONAL
+	// BY ROW.
+	IsRegionalByRow() bool
+
+	// GetDatabaseID returns the owning database id of the table, or zero, if the
+	// owning database could not be determined.
+	GetDatabaseID() descpb.ID
 }
 
 // CheckConstraint contains the SQL text and the validity status for a check

pkg/sql/opt/constraint/constraint.go

@@ -830,3 +830,27 @@ func (c *Constraint) CollectFirstColumnValues(
 	}
 	return values, hasNullValue, true
 }
+
+// HasRemoteSpans returns true if any of the constraint spans in `c` belong to
+// partitions remote to the gateway region.
+func (c *Constraint) HasRemoteSpans(ps *partition.PrefixSorter) bool {
+	if c.IsUnconstrained() {
+		return true
+	}
+	if c.IsContradiction() {
+		return false
+	}
+	// Iterate through the spans and determine whether each one matches
+	// with a prefix from a remote partition.
+	for i, n := 0, c.Spans.Count(); i < n; i++ {
+		span := c.Spans.Get(i)
+		if match, ok := FindMatch(span, ps); ok {
+			if !match.IsLocal {
+				return true
+			}
+		} else {
+			return true
+		}
+	}
+	return false
+}

pkg/sql/opt/distribution/BUILD.bazel

@@ -7,9 +7,14 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/sql/opt/distribution",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/sql/catalog/descpb",
+        "//pkg/sql/opt",
         "//pkg/sql/opt/memo",
         "//pkg/sql/opt/props/physical",
+        "//pkg/sql/pgwire/pgcode",
+        "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/sem/eval",
+        "//pkg/sql/sem/tree",
         "//pkg/util/buildutil",
         "@com_github_cockroachdb_errors//:errors",
     ],