sql: audit DistSQL supportability for streamer

In c17591e we made so that if we have
a plan that cannot be distributed for any reason, we would disable the
streamer out of caution. This was known to be a super-set of problematic
scenarios, and this commit addresses that oversight.

In particular, it revamps `checkSupportForPlanNode` function so that we
accumulate all possible causes for DistSQL being prohibited, and then we
examine each one to decide whether it's safe to use the streamer given
the causes found. The causes are accumulated into a bit-map, and the
methods have been adjusted to do a full walk over the planNode tree (as
well as over all expressions) to capture all causes - this should have
negligible overhead.

An additional improvement is that we now replace errors with different
bits and we can log all causes for DistSQL not being supported.

What prompted this work was a customer query where we didn't use the
streamer because of "unsupported planNode" distsql error. This was
attempted to be fixed in f6038a1, but
the fix was incomplete, and now it is.

Release note: None

Author: yuzefovich

pkg/gen/stringer.bzl

@@ -46,6 +46,7 @@ STRINGER_SRCS = [
     "//pkg/sql/sem/tree:statementreturntype_string.go",
     "//pkg/sql/sem/tree:statementtype_string.go",
     "//pkg/sql:advancecode_string.go",
+    "//pkg/sql:distsqlprohibitedcause_string.go",
     "//pkg/sql:nodestatus_string.go",
     "//pkg/sql:txneventtype_string.go",
     "//pkg/sql:txntype_string.go",

pkg/sql/BUILD.bazel

@@ -298,6 +298,7 @@ go_library(
         "zone_config.go",
         "zone_config_helper.go",
         ":gen-advancecode-stringer",  # keep
+        ":gen-distsqlprohibitedcause-stringer",  # keep
         ":gen-nodestatus-stringer",  # keep
         ":gen-txneventtype-stringer",  # keep
         ":gen-txntype-stringer",  # keep
@@ -680,6 +681,7 @@ go_test(
         "delete_preserving_index_test.go",
         "descriptor_mutation_test.go",
         "descriptor_test.go",
+        "distsql_check_test.go",
         "distsql_leaf_txn_test.go",
         "distsql_physical_planner_test.go",
         "distsql_plan_backfill_test.go",
@@ -871,6 +873,7 @@ go_test(
         "//pkg/sql/execstats",
         "//pkg/sql/flowinfra",
         "//pkg/sql/gcjob",
+        "//pkg/sql/inverted",
         "//pkg/sql/isql",
         "//pkg/sql/lexbase",
         "//pkg/sql/mutations",
@@ -1025,6 +1028,12 @@ stringer(
     typ = "advanceCode",
 )
 
+stringer(
+    name = "gen-distsqlprohibitedcause-stringer",
+    src = "distsql_check.go",
+    typ = "distSQLProhibitedCause",
+)
+
 stringer(
     name = "gen-nodestatus-stringer",
     src = "distsql_physical_planner.go",

pkg/sql/apply_join.go

@@ -333,14 +333,14 @@ func runPlanInsidePlan(
 		}
 	}
 
-	distributePlan, distSQLProhibitedErr := plannerCopy.getPlanDistribution(ctx, plan.main, notPostquery)
+	distributePlan, noDistSQLCauses := plannerCopy.getPlanDistribution(ctx, plan.main, notPostquery)
 	distributeType := DistributionType(LocalDistribution)
 	if distributePlan.WillDistribute() {
 		distributeType = FullDistribution
 	}
 	evalCtx := evalCtxFactory()
 	planCtx := execCfg.DistSQLPlanner.NewPlanningCtx(ctx, evalCtx, &plannerCopy, plannerCopy.txn, distributeType)
-	planCtx.distSQLProhibitedErr = distSQLProhibitedErr
+	planCtx.distSQLProhibitedCauses = noDistSQLCauses
 	planCtx.stmtType = recv.stmtType
 	if sqlStatsBuilder != nil && plannerCopy.instrumentation.ShouldSaveFlows() {
 		planCtx.collectExecStats = true

pkg/sql/conn_executor.go

@@ -3339,7 +3339,7 @@ func (ex *connExecutor) execCopyIn(
 				defer p.curPlan.close(ctx)
 				_, err := ex.execWithDistSQLEngine(
 					ctx, p, tree.RowsAffected, res, LocalDistribution,
-					nil /* progressAtomic */, nil, /* distSQLProhibitedErr */
+					nil /* progressAtomic */, 0, /* distSQLProhibitedCauses */
 				)
 				return err
 			},

pkg/sql/conn_executor_exec.go

@@ -2879,7 +2879,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 			ex.sessionData().DistSQLMode = origDistSQLMode
 		}
 	}
-	distributePlan, distSQLProhibitedErr := planner.getPlanDistribution(ctx, planner.curPlan.main, notPostquery)
+	distributePlan, noDistSQLCauses := planner.getPlanDistribution(ctx, planner.curPlan.main, notPostquery)
 	if afterGetPlanDistribution != nil {
 		afterGetPlanDistribution()
 	}
@@ -2927,7 +2927,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 	}
 	ex.sessionTracing.TraceExecStart(ctx, "distributed")
 	stats, err := ex.execWithDistSQLEngine(
-		ctx, planner, stmt.AST.StatementReturnType(), res, distribute, progAtomic, distSQLProhibitedErr,
+		ctx, planner, stmt.AST.StatementReturnType(), res, distribute, progAtomic, noDistSQLCauses,
 	)
 	if ppInfo := getPausablePortalInfo(planner); ppInfo != nil {
 		// For pausable portals, we log the stats when closing the portal, so we need
@@ -3379,7 +3379,7 @@ func (ex *connExecutor) execWithDistSQLEngine(
 	res RestrictedCommandResult,
 	distribute DistributionType,
 	progressAtomic *uint64,
-	distSQLProhibitedErr error,
+	distSQLProhibitedCauses distSQLProhibitedCauses,
 ) (topLevelQueryStats, error) {
 	defer planner.curPlan.savePlanInfo()
 	recv := MakeDistSQLReceiver(
@@ -3404,7 +3404,7 @@ func (ex *connExecutor) execWithDistSQLEngine(
 		evalCtx := planner.ExtendedEvalContext()
 		planCtx := ex.server.cfg.DistSQLPlanner.NewPlanningCtx(ctx, evalCtx, planner, planner.txn, distribute)
 		planCtx.setUpForMainQuery(ctx, planner, recv)
-		planCtx.distSQLProhibitedErr = distSQLProhibitedErr
+		planCtx.distSQLProhibitedCauses = distSQLProhibitedCauses
 
 		var evalCtxFactory func(usedConcurrently bool) *extendedEvalContext
 		if len(planner.curPlan.subqueryPlans) != 0 ||

pkg/sql/create_stats.go

@@ -661,7 +661,7 @@ func createStatsDefaultColumns(
 		}
 		var distSQLVisitor distSQLExprCheckVisitor
 		for i, col := range desc.PublicColumns() {
-			cannotDistribute[i] = col.IsVirtual() && checkExprForDistSQL(exprs[i], &distSQLVisitor) != nil
+			cannotDistribute[i] = col.IsVirtual() && checkExprForDistSQL(exprs[i], &distSQLVisitor) != 0
 		}
 	}