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 blockers for DistSQL, and then we examine each
one to decide whether it's safe to use the streamer given the blockers
found. The blockers 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 blockers - this should have negligible
overhead.

An additional improvement is that we now replace errors with different
bits and we can log all blockers 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:distsqlblocker_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-distsqlblocker-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-distsqlblocker-stringer",
+    src = "distsql_check.go",
+    typ = "distSQLBlocker",
+)
+
 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, blockers := 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.distSQLBlockers = blockers
 	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, /* distSQLBlockers */
 				)
 				return err
 			},

pkg/sql/conn_executor_exec.go

@@ -2893,7 +2893,7 @@ func (ex *connExecutor) dispatchToExecutionEngine(
 			ex.sessionData().DistSQLMode = origDistSQLMode
 		}
 	}
-	distributePlan, distSQLProhibitedErr := planner.getPlanDistribution(ctx, planner.curPlan.main, notPostquery)
+	distributePlan, blockers := planner.getPlanDistribution(ctx, planner.curPlan.main, notPostquery)
 	if afterGetPlanDistribution != nil {
 		afterGetPlanDistribution()
 	}
@@ -2941,7 +2941,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, blockers,
 	)
 	if ppInfo := getPausablePortalInfo(planner); ppInfo != nil {
 		// For pausable portals, we log the stats when closing the portal, so we need
@@ -3393,7 +3393,7 @@ func (ex *connExecutor) execWithDistSQLEngine(
 	res RestrictedCommandResult,
 	distribute DistributionType,
 	progressAtomic *uint64,
-	distSQLProhibitedErr error,
+	distSQLBlockers distSQLBlockers,
 ) (topLevelQueryStats, error) {
 	defer planner.curPlan.savePlanInfo()
 	recv := MakeDistSQLReceiver(
@@ -3418,7 +3418,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.distSQLBlockers = distSQLBlockers
 
 		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
 		}
 	}