sql: unconditionally check DistSQL supportability

This commit hardens c17591e. In
particular, in that change we disabled the usage of the Streamer
whenever some part of the plan isn't distributable. However, population
of `distSQLProhibitedErr` was left on a best-effort basis. In
particular, if we have `distsql=off` (as well as a couple of other
conditions), we'd skip the supportability check. This meant that we
could still end up using the Streamer in some illegal cases - one
example that I came up with is if we have a routine (which currently
prohibits distsql), it might access the RootTxn concurrently with the
LeafTxn access by the Streamer, which is no bueno. This commit makes it
so that we do DistSQL supportability check unconditionally. Note that
this shouldn't really have a performance impact - after all, we do
expect this check to be done pretty much all the time (unless someone
runs with `distsql=off`).

Additionally, this change happens to fix a nil pointer error around the
recent fix to top-level query stats in presence of routines. Namely,
that patch only set the metadata forwarder for local plans if we end up
using the RootTxn, but in a query where we happened to use the streamer
with routines, (before this patch) we'd end up with LeafTxn and unset
metadata forwarder. Now we'll have streamer disabled, so we'll have the
RootTxn and the forwarder will be set.

Release note: None

Author: yuzefovich

pkg/sql/distsql_physical_planner.go

@@ -977,7 +977,8 @@ type PlanningCtx struct {
 	// isLocal is set to true if we're planning this query on a single node.
 	isLocal bool
 	// distSQLProhibitedErr, if set, indicates why the plan couldn't be
-	// distributed.
+	// distributed. If any part of the plan isn't distributable, then this is
+	// guaranteed to be non-nil.
 	distSQLProhibitedErr error
 	planner              *planner
 

pkg/sql/distsql_running.go

@@ -875,15 +875,10 @@ func (dsp *DistSQLPlanner) Run(
 			// that we might have a plan where some expression (e.g. a cast to
 			// an Oid type) uses the planner's txn (which is the RootTxn), so
 			// it'd be illegal to use LeafTxns for a part of such plan.
-			// TODO(yuzefovich): this check is both excessive and insufficient.
-			// For example:
-			// - it disables the usage of the Streamer when a subquery has an
-			// Oid type, but that would have no impact on usage of the Streamer
-			// in the main query;
-			// - it might allow the usage of the Streamer even when the internal
-			// executor is used by a part of the plan, and the IE would use the
-			// RootTxn. Arguably, this would be a bug in not prohibiting the
-			// DistSQL altogether.
+			// TODO(yuzefovich): this check could be excessive. For example, it
+			// disables the usage of the Streamer when a subquery has an Oid
+			// type (due to a serialization issue), but that would have no
+			// impact on usage of the Streamer in the main query.
 			if !containsLocking && !mustUseRootTxn && planCtx.distSQLProhibitedErr == nil {
 				if evalCtx.SessionData().StreamerEnabled {
 					for _, proc := range plan.Processors {

pkg/sql/distsql_running_test.go

@@ -1221,12 +1221,13 @@ func TestTopLevelQueryStats(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
 
+	ctx := context.Background()
 	// testQuery will be updated throughout the test to the current target.
 	var testQuery atomic.Value
 	// The callback will send number of rows read and rows written (for each
 	// ProducerMetadata.Metrics object) on these channels, respectively.
 	rowsReadCh, rowsWrittenCh := make(chan int64), make(chan int64)
-	s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
+	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
 		Knobs: base.TestingKnobs{
 			SQLExecutor: &ExecutorTestingKnobs{
 				DistSQLReceiverPushCallbackFactory: func(_ context.Context, query string) func(rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) (rowenc.EncDatumRow, coldata.Batch, *execinfrapb.ProducerMetadata) {
@@ -1244,11 +1245,13 @@ func TestTopLevelQueryStats(t *testing.T) {
 			},
 		},
 	})
-	defer s.Stopper().Stop(context.Background())
+	defer srv.Stopper().Stop(ctx)
+	conn, err := sqlDB.Conn(ctx)
+	require.NoError(t, err)
 
 	if _, err := sqlDB.Exec(`
-CREATE TABLE t (k INT PRIMARY KEY);
-INSERT INTO t SELECT generate_series(1, 10);
+CREATE TABLE t (k INT PRIMARY KEY, i INT, v INT, INDEX(i));
+INSERT INTO t SELECT i, 1, 1 FROM generate_series(1, 10) AS g(i);
 CREATE FUNCTION no_reads() RETURNS INT AS 'SELECT 1' LANGUAGE SQL;
 CREATE FUNCTION reads() RETURNS INT AS 'SELECT count(*) FROM t' LANGUAGE SQL;
 CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x' LANGUAGE SQL;
@@ -1259,6 +1262,7 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 	for _, tc := range []struct {
 		name           string
 		query          string
+		setup, cleanup string // optional
 		expRowsRead    int64
 		expRowsWritten int64
 	}{
@@ -1268,6 +1272,16 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 			expRowsRead:    10,
 			expRowsWritten: 0,
 		},
+		{
+			name:    "routine and index join (used to be powered by streamer)",
+			query:   "SELECT v FROM t@t_i_idx WHERE reads() > 0",
+			setup:   "SET distsql=off",
+			cleanup: "RESET distsql",
+			// 10 rows for secondary index, 10 for index join into primary, and
+			// then for each row do ten-row-scan in the routine.
+			expRowsRead:    120,
+			expRowsWritten: 0,
+		},
 		{
 			name:           "simple write",
 			query:          "INSERT INTO t SELECT generate_series(11, 42)",
@@ -1309,13 +1323,23 @@ CREATE FUNCTION write(x INT) RETURNS INT AS 'INSERT INTO t VALUES (x); SELECT x'
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
+			if tc.setup != "" {
+				_, err := conn.ExecContext(ctx, tc.setup)
+				require.NoError(t, err)
+			}
+			if tc.cleanup != "" {
+				defer func() {
+					_, err := conn.ExecContext(ctx, tc.cleanup)
+					require.NoError(t, err)
+				}()
+			}
 			testQuery.Store(tc.query)
 			errCh := make(chan error)
 			// Spin up the worker goroutine which will actually execute the
 			// query.
 			go func() {
 				defer close(errCh)
-				_, err := sqlDB.Exec(tc.query)
+				_, err := conn.ExecContext(ctx, tc.query)
 				errCh <- err
 			}()
 			// In the main goroutine, loop until the query is completed while

pkg/sql/exec_util.go

@@ -2225,8 +2225,10 @@ func shouldDistributeGivenRecAndMode(
 // the plan.
 //
 // The returned error, if any, indicates why we couldn't distribute the plan.
-// Note that it's possible that we choose to not distribute the plan while
-// nil error is returned.
+// Note that it's possible that we choose to not distribute the plan while nil
+// error is returned (but it's guaranteed that if some part of the plan isn't
+// distributable, then non-nil error is returned).
+//
 // WARNING: in some cases when this method returns
 // physicalplan.FullyDistributedPlan, the plan might actually run locally. This
 // is the case when
@@ -2245,23 +2247,9 @@ func (p *planner) getPlanDistribution(
 		return plan.physPlan.Distribution, nil
 	}
 
-	// If this transaction has modified or created any types, it is not safe to
-	// distribute due to limitations around leasing descriptors modified in the
-	// current transaction.
-	if p.Descriptors().HasUncommittedDescriptors() {
-		return physicalplan.LocalPlan, nil
-	}
-
+	// Check DistSQL-supportability as the first order of business in order to
+	// find whether usage of DistSQL is prohibited by features of the plan.
 	sd := p.SessionData()
-	if sd.DistSQLMode == sessiondatapb.DistSQLOff {
-		return physicalplan.LocalPlan, nil
-	}
-
-	// Don't try to run empty nodes (e.g. SET commands) with distSQL.
-	if _, ok := plan.planNode.(*zeroNode); ok {
-		return physicalplan.LocalPlan, nil
-	}
-
 	// Determine whether the txn has buffered some writes.
 	txnHasBufferedWrites := p.txn.HasBufferedWrites()
 	if sd.BufferedWritesEnabled && p.curPlan.main == plan {
@@ -2277,11 +2265,26 @@ func (p *planner) getPlanDistribution(
 	}
 	rec, err := checkSupportForPlanNode(ctx, plan.planNode, &p.distSQLVisitor, sd, txnHasBufferedWrites)
 	if err != nil {
-		// Don't use distSQL for this request.
 		log.VEventf(ctx, 1, "query not supported for distSQL: %s", err)
 		return physicalplan.LocalPlan, err
 	}
 
+	// If this transaction has modified or created any types, it is not safe to
+	// distribute due to limitations around leasing descriptors modified in the
+	// current transaction.
+	if p.Descriptors().HasUncommittedDescriptors() {
+		return physicalplan.LocalPlan, nil
+	}
+
+	if sd.DistSQLMode == sessiondatapb.DistSQLOff {
+		return physicalplan.LocalPlan, nil
+	}
+
+	// Don't try to run empty nodes (e.g. SET commands) with distSQL.
+	if _, ok := plan.planNode.(*zeroNode); ok {
+		return physicalplan.LocalPlan, nil
+	}
+
 	if shouldDistributeGivenRecAndMode(rec, sd.DistSQLMode) {
 		return physicalplan.FullyDistributedPlan, nil
 	}

pkg/sql/opt/exec/execbuilder/testdata/inverted_join_geospatial

@@ -336,7 +336,7 @@ quality of service: regular
 │       │ estimated max sql temp disk usage: 0 B
 │       │ group by: lk
 │       │
-│       └── • lookup join (left outer) (streamer)
+│       └── • lookup join (left outer)
 │           │ sql nodes: <hidden>
 │           │ regions: <hidden>
 │           │ actual row count: 0

pkg/sql/opt/exec/execbuilder/testdata/jsonb_path_query

@@ -153,19 +153,19 @@ query T kvtrace
 SELECT jsonb_path_query(b, '$.a.b') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a;
 ----
 Scan /Table/106/2/{???-"a"/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/}
-Scan /Table/106/1/7/0, /Table/106/1/11/0, /Table/106/1/8/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0, /Table/106/1/12/0, /Table/106/1/3/0, /Table/106/1/9/0, /Table/106/1/4/0, /Table/106/1/6/0
+Scan /Table/106/1/3/0, /Table/106/1/4/0, /Table/106/1/6/0, /Table/106/1/7/0, /Table/106/1/8/0, /Table/106/1/9/0, /Table/106/1/11/0, /Table/106/1/12/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0
 
 query T kvtrace
 SELECT jsonb_path_query(b, '$.a[*].b') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a[*].b') ORDER BY a;
 ----
 Scan /Table/106/2/{???-"a"/Arr/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/}
-Scan /Table/106/1/7/0, /Table/106/1/11/0, /Table/106/1/8/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0, /Table/106/1/12/0, /Table/106/1/3/0, /Table/106/1/9/0, /Table/106/1/4/0, /Table/106/1/6/0
+Scan /Table/106/1/3/0, /Table/106/1/4/0, /Table/106/1/6/0, /Table/106/1/7/0, /Table/106/1/8/0, /Table/106/1/9/0, /Table/106/1/11/0, /Table/106/1/12/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0
 
 query T kvtrace
 SELECT jsonb_path_query(b, '$.a[*][*].b') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a[*][*].b') ORDER BY a;
 ----
 Scan /Table/106/2/{???-"a"/Arr/Arr/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/}
-Scan /Table/106/1/7/0, /Table/106/1/11/0, /Table/106/1/8/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0, /Table/106/1/12/0, /Table/106/1/3/0, /Table/106/1/9/0, /Table/106/1/4/0, /Table/106/1/6/0
+Scan /Table/106/1/3/0, /Table/106/1/4/0, /Table/106/1/6/0, /Table/106/1/7/0, /Table/106/1/8/0, /Table/106/1/9/0, /Table/106/1/11/0, /Table/106/1/12/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0
 
 query T kvtrace
 SELECT jsonb_path_query(b, '$.a.b[*].c') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b[*].c') ORDER BY a;
@@ -234,7 +234,7 @@ query T kvtrace
 SELECT a, jsonb_path_query(b, '$.a.b') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b') ORDER BY a;
 ----
 Scan /Table/106/2/{???-"a"/Arr/Arr/}, /Table/106/2/{???-"a"/Arr/}
-Scan /Table/106/1/7/0, /Table/106/1/11/0, /Table/106/1/8/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0, /Table/106/1/12/0, /Table/106/1/3/0, /Table/106/1/9/0, /Table/106/1/4/0, /Table/106/1/6/0
+Scan /Table/106/1/3/0, /Table/106/1/4/0, /Table/106/1/6/0, /Table/106/1/7/0, /Table/106/1/8/0, /Table/106/1/9/0, /Table/106/1/11/0, /Table/106/1/12/0, /Table/106/1/13/0, /Table/106/1/14/0, /Table/106/1/15/0, /Table/106/1/16/0
 
 query T kvtrace
 SELECT a, jsonb_path_query(b, '$.a.b'), jsonb_path_query(b, '$.a.d') FROM json_tab@foo_inv WHERE jsonb_path_exists(b, '$.a.b') AND jsonb_path_exists(b, '$.a.d') ORDER BY a;