Merge pull request cockroachdb#101504 from msirek/backport23.1.0-100977

release-23.1.0: execbuilder: fix internal error in handleRemoteLookupJoinError

pkg/ccl/logictestccl/testdata/logic_test/multi_region_remote_access_error

@@ -76,10 +76,6 @@ CREATE TABLE messages_rbt (
     INDEX msg_idx(message)
 ) LOCALITY REGIONAL BY TABLE
 
-# Sleep so that the follow_read_timestamp() used by phase 3 dynamic checking
-# of a query's home region isn't executing before the tables existed.
-sleep 5s
-
 statement ok
 CREATE TABLE messages_rbr (
     account_id INT NOT NULL,
@@ -197,12 +193,93 @@ CREATE TABLE json_arr2_rbt (
   a STRING[]
 ) LOCALITY REGIONAL BY TABLE
 
+# Sleep so that the follower_read_timestamp() used by phase 3 dynamic checking
+# of a query's home region isn't executing before the tables existed.
+sleep 5s
+
 statement ok
 SET enforce_home_region = true
 
 statement ok
 SET enforce_home_region_follower_reads_enabled = true
 
+# A lookup join with a CTE as input should succeed.
+statement ok
+WITH vtab AS (INSERT INTO parent VALUES (6) RETURNING p_id)
+SELECT * FROM vtab
+  INNER LOOKUP JOIN messages_rbt rbt on vtab.p_id = rbt.account_id
+
+# A lookup join with a scalar list of constants as input should succeed.
+retry
+statement ok
+SELECT * FROM (SELECT 1, 'Hello, Dude!') vtab(account_id, message)
+  INNER LOOKUP JOIN messages_rbt rbt on vtab.account_id = rbt.account_id
+
+# A scalar subquery with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT * FROM (SELECT (SELECT max(account_id) FROM messages_rbr), 'Hello, Dude!') vtab(account_id, message)
+  INNER LOOKUP JOIN messages_rbt rbt on vtab.account_id = rbt.account_id
+
+# A scalar subquery with a home region other than the gateway should fail.
+retry
+statement error pq: Query is not running in its home region\. Try running the query from region 'ca-central-1'\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT * FROM (SELECT (SELECT max(account_id) FROM messages_rbr WHERE crdb_region = 'ca-central-1'), 'Hello, Dude!') vtab(account_id, message)
+  INNER LOOKUP JOIN messages_rbt rbt on vtab.account_id = rbt.account_id
+
+# A scalar subquery with a home region should succeed.
+retry
+statement ok
+SELECT * FROM (SELECT (SELECT max(account_id) FROM messages_rbt), 'Hello, Dude!') vtab(account_id, message)
+  INNER LOOKUP JOIN messages_rbt rbt on vtab.account_id = rbt.account_id
+
+statement ok
+CREATE OR REPLACE FUNCTION rbr() RETURNS INT AS 'SELECT max(account_id) FROM messages_rbr' LANGUAGE SQL;
+
+# A UDF with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a filter on messages_rbr\.crdb_region and/or on key column \(multi_region_test_db\.public\.messages_rbr\.account_id\)\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT rbr()
+
+# An EXISTS subquery with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT * FROM (SELECT 1 WHERE EXISTS (SELECT max(account_id) FROM messages_rbr))
+
+# An apply join with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT * FROM messages_rbt rbt, LATERAL
+  (SELECT count(DISTINCT rbr.account_id + rbt.account_id + 1) AS g
+    FROM messages_rbr rbr WHERE rbr.account_id * rbt.account_id < 10)
+
+# An apply join with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a filter on rbr\.crdb_region and/or on key column \(rbr\.account_id\)\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT * FROM messages_rbr rbr, LATERAL
+  (SELECT count(DISTINCT rbr.account_id + rbt.account_id + 1) AS g
+    FROM messages_rbt rbt WHERE rbr.account_id * rbt.account_id < 10)
+
+# An array scalar expression with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a filter on messages_rbr\.crdb_region and/or on key column \(multi_region_test_db\.public\.messages_rbr\.account_id\)\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT ARRAY[rbr(),rbr()]
+
+# An ALL subquery with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT 1 WHERE 1 > ALL (SELECT max(account_id) FROM messages_rbr)
+
+# An ANY subquery with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT 1 WHERE 1 > ANY (SELECT max(account_id) FROM messages_rbr)
+
+# A SOME subquery with no home region should fail.
+retry
+statement error pq: Query has no home region\. Try adding a LIMIT clause\. For more information, see https://www.cockroachlabs.com/docs/stable/cost-based-optimizer.html#control-whether-queries-are-limited-to-a-single-region
+SELECT 1 WHERE 1 > ANY (SELECT max(account_id) FROM messages_rbr)
+
 ### Regression tests for issue #89875
 
 # Non-DML statements should not error out due to enforce_home_region.

pkg/sql/opt/exec/execbuilder/relational.go

@@ -2155,7 +2155,20 @@ func (b *Builder) handleRemoteLookupJoinError(join *memo.LookupJoinExpr) (err er
 	inputTableName := ""
 	// ScanExprs from global tables will have filled in a provided distribution
 	// of the gateway region by now.
-	queryHomeRegion, queryHasHomeRegion := input.(memo.RelExpr).ProvidedPhysical().Distribution.GetSingleRegion()
+	var queryHomeRegion string
+	var queryHasHomeRegion bool
+	if rel, ok := input.(memo.RelExpr); ok {
+		queryHomeRegion, queryHasHomeRegion = rel.ProvidedPhysical().Distribution.GetSingleRegion()
+	} else if _, ok = input.(*memo.ScalarListExpr); ok {
+		// A list of scalar constants doesn't access remote regions.
+		// If these aren't constants, such as scalar subqueries, checks for a home
+		// region are done elsewhere.
+		queryHasHomeRegion = true
+		queryHomeRegion = gatewayRegion
+	} else {
+		return errors.AssertionFailedf("unexpected expression kind while checking home region of input to lookup join: %v", input)
+	}
+
 	var inputTableMeta *opt.TableMeta
 	var inputTable cat.Table
 	var inputIndexOrdinal cat.IndexOrdinal

pkg/sql/opt/exec/execbuilder/scalar.go

@@ -13,19 +13,23 @@ package execbuilder
 import (
 	"context"
 
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/xform"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
+	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/builtins/builtinsregistry"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treebin"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree/treecmp"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/volatility"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -710,6 +714,26 @@ func (b *Builder) buildSubquery(
 	subquery := scalar.(*memo.SubqueryExpr)
 	input := subquery.Input
 
+	if b.evalCtx.SessionData().EnforceHomeRegion && b.IsANSIDML {
+		inputDistributionProvidedPhysical := input.ProvidedPhysical()
+		if homeRegion, ok := inputDistributionProvidedPhysical.Distribution.GetSingleRegion(); ok {
+			if gatewayRegion, ok := b.evalCtx.GetLocalRegion(); ok {
+				if homeRegion != gatewayRegion {
+					return nil, pgerror.Newf(pgcode.QueryNotRunningInHomeRegion,
+						`%s. Try running the query from region '%s'. %s`,
+						execinfra.QueryNotRunningInHomeRegionMessagePrefix,
+						homeRegion,
+						sqlerrors.EnforceHomeRegionFurtherInfo,
+					)
+				}
+			}
+		} else {
+			return nil, pgerror.Newf(pgcode.QueryHasNoHomeRegion,
+				"Query has no home region. Try adding a LIMIT clause. %s",
+				sqlerrors.EnforceHomeRegionFurtherInfo)
+		}
+	}
+
 	// TODO(radu): for now we only support the trivial projection.
 	cols := input.Relational().OutputCols
 	if cols.Len() != 1 {