sql: only fetch specific columns needed to validate check constraints…

… for UPDATES

This addresses a longstanding TODO to only request the columns used
in check constraints when performing an UPDATE instead of blindly
requesting all columns when a check constraint is present.

8d37935 did all of the heavy lifting for this. This change just plugs
it in.

The new logic tests fail if we don't correctly request the columns
we need for the check constraint. Interestingly, this isn't because
we don't fetch the other columns from KV. Instead, it's because we
don't decode them and pass them up the stack.

Release note: None

pkg/sql/logictest/testdata/logic_test/check_constraints

@@ -273,3 +273,77 @@ CREATE TABLE test2.t (
   CHECK (t.a > 0),
   CHECK (test2.t.a > 0)
 )
+
+# Use multiple column families.
+
+statement ok
+CREATE TABLE t9 (
+  a INT PRIMARY KEY,
+  b INT,
+  c INT,
+  d INT,
+  e INT,
+  FAMILY (a),
+  FAMILY (b),
+  FAMILY (c),
+  FAMILY (d, e),
+  CHECK (a > b),
+  CHECK (d IS NULL)
+)
+
+statement ok
+INSERT INTO t9 VALUES (5, 3)
+
+statement error pgcode 23514 failed to satisfy CHECK constraint \(a > b\)
+INSERT INTO t9 VALUES (6, 7)
+
+statement ok
+UPDATE t9 SET b = 4 WHERE a = 5
+
+statement error pgcode 23514 failed to satisfy CHECK constraint \(a > b\)
+UPDATE t9 SET b = 6 WHERE a = 5
+
+# Only column families that are needed to validate check constraints are fetched.
+query TTTTT
+EXPLAIN (VERBOSE) UPDATE t9 SET b = 2 WHERE a = 5
+----
+count                ·         ·                          ()                                 ·
+ └── update          ·         ·                          ()                                 ·
+      │              table     t9                         ·                                  ·
+      │              set       b                          ·                                  ·
+      │              check 0   t9.a > t9.b                ·                                  ·
+      │              check 1   t9.d IS NULL               ·                                  ·
+      └── render     ·         ·                          (a, b, d, "?column?")              a=CONST; "?column?"=CONST; key()
+           │         render 0  test.public.t9.a           ·                                  ·
+           │         render 1  test.public.t9.b           ·                                  ·
+           │         render 2  test.public.t9.d           ·                                  ·
+           │         render 3  2                          ·                                  ·
+           └── scan  ·         ·                          (a, b, c[omitted], d, e[omitted])  a=CONST; key()
+·                    table     t9@primary                 ·                                  ·
+·                    spans     /5/0-/5/1/2 /5/3/1-/5/3/2  ·                                  ·
+
+statement ok
+UPDATE t9 SET a = 7 WHERE a = 4
+
+statement error pgcode 23514 failed to satisfy CHECK constraint \(a > b\)
+UPDATE t9 SET a = 2 WHERE a = 5
+
+query TTTTT
+EXPLAIN (VERBOSE) UPDATE t9 SET a = 2 WHERE a = 5
+----
+count                ·         ·                 ()                           ·
+ └── update          ·         ·                 ()                           ·
+      │              table     t9                ·                            ·
+      │              set       a                 ·                            ·
+      │              check 0   t9.a > t9.b       ·                            ·
+      │              check 1   t9.d IS NULL      ·                            ·
+      └── render     ·         ·                 (a, b, c, d, e, "?column?")  a=CONST; "?column?"=CONST; key()
+           │         render 0  test.public.t9.a  ·                            ·
+           │         render 1  test.public.t9.b  ·                            ·
+           │         render 2  test.public.t9.c  ·                            ·
+           │         render 3  test.public.t9.d  ·                            ·
+           │         render 4  test.public.t9.e  ·                            ·
+           │         render 5  2                 ·                            ·
+           └── scan  ·         ·                 (a, b, c, d, e)              a=CONST; key()
+·                    table     t9@primary        ·                            ·
+·                    spans     /5-/5/#           ·                            ·

pkg/sql/update.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/types"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/pkg/errors"
 )
@@ -198,12 +199,42 @@ func (p *planner) Update(
 	rowsNeeded := resultsNeeded(n.Returning)
 
 	var requestedCols []sqlbase.ColumnDescriptor
-	if rowsNeeded || len(desc.Checks) > 0 {
+	if rowsNeeded {
 		// TODO(dan): This could be made tighter, just the rows needed for RETURNING
 		// exprs.
-		// TODO(nvanbenschoten): This could be made tighter, just the rows needed for
-		// the CHECK exprs.
 		requestedCols = desc.Columns
+	} else if len(desc.Checks) > 0 {
+		// Request any columns we'll need when validating check constraints. We
+		// could be smarter and only validate check constraints which depend on
+		// columns that are being modified in the UPDATE statement, in which
+		// case we'd only need to request the columns used by that subset of
+		// check constraints, but that doesn't seem worth the effort.
+		//
+		// TODO(nvanbenschoten): These conditions shouldn't be mutually
+		// exclusive, but since rowsNeeded implies that requestedCols =
+		// desc.Columns, there's no reason to enter this block if rowsNeeded is
+		// true. Remove this when the TODO above is addressed.
+		var requestedColSet util.FastIntSet
+		for _, col := range requestedCols {
+			requestedColSet.Add(int(col.ID))
+		}
+		for _, ck := range desc.Checks {
+			cols, err := ck.ColumnsUsed(desc)
+			if err != nil {
+				return nil, err
+			}
+			for _, colID := range cols {
+				if !requestedColSet.Contains(int(colID)) {
+					col, err := desc.FindColumnByID(colID)
+					if err != nil {
+						return nil, errors.Wrapf(err, "error finding column %d in table %s",
+							colID, desc.Name)
+					}
+					requestedCols = append(requestedCols, *col)
+					requestedColSet.Add(int(colID))
+				}
+			}
+		}
 	}
 
 	// Create the table updater, which does the bulk of the work.