From ae220c126ac313041df12125109b2cad983574a8 Mon Sep 17 00:00:00 2001
From: Andrew Werner <awerner32@gmail.com>
Date: Mon, 23 May 2022 16:59:35 -0400
Subject: [PATCH] sql: enforce NOT NULL when adding a virtual computed column

Before this change, we did not ever validate that newly added
virtual computed columns which were marked NOT NULL actually were
not NULL. This is because there historically, in the legacy schema changer,
we'd validate the `NULL` nature of the column when performing a column
backfill. In the declarative schema changer, we don't do a column backfill,
so we have bugs aplenty related to the general failure to check the `NULL`
disposition of the data (see #81679).

This change only affects the legacy schema changer so that it can be
backported. A separate issue (#81690) has been filed for parity in the
declarative schema changer.

Release note (bug fix): Before this change, not validation was performed
when adding a virtual computed column which was marked `NOT NULL`. This meant
that it was possible to have a virtual computed column with an active `NOT
NULL`
constraint despite having rows in the table for which the column was `NULL`.
This has now been fixed.
---
 pkg/sql/add_column.go                         | 11 ++++++++
 pkg/sql/alter_table.go                        | 25 ++++++++++++-------
 .../testdata/logic_test/virtual_columns       | 25 +++++++++++++++++++
 3 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/add_column.go b/pkg/sql/add_column.go
index b15e9a4d6c75..3311a1258d84 100644
--- a/pkg/sql/add_column.go
+++ b/pkg/sql/add_column.go
@@ -206,6 +206,17 @@ func (p *planner) addColumnImpl(
 		}
 	}
 
+	if col.Virtual && !col.Nullable {
+		colName := tree.Name(col.Name)
+		newCol, err := n.tableDesc.FindColumnWithName(colName)
+		if err != nil {
+			return errors.NewAssertionErrorWithWrappedErrf(err, "failed to find newly added column %v", colName)
+		}
+		if err := addNotNullConstraintMutationForCol(n.tableDesc, newCol); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go
index 881bc11d28ce..2f01f0e610ae 100644
--- a/pkg/sql/alter_table.go
+++ b/pkg/sql/alter_table.go
@@ -998,17 +998,9 @@ func applyColumnMutation(
 			}
 		}
 
-		info, err := tableDesc.GetConstraintInfo()
-		if err != nil {
+		if err := addNotNullConstraintMutationForCol(tableDesc, col); err != nil {
 			return err
 		}
-		inuseNames := make(map[string]struct{}, len(info))
-		for k := range info {
-			inuseNames[k] = struct{}{}
-		}
-		check := tabledesc.MakeNotNullCheckConstraint(col.GetName(), col.GetID(), tableDesc.GetNextConstraintID(), inuseNames, descpb.ConstraintValidity_Validating)
-		tableDesc.AddNotNullMutation(check, descpb.DescriptorMutation_ADD)
-		tableDesc.NextConstraintID++
 
 	case *tree.AlterTableDropNotNull:
 		if col.IsNullable() {
@@ -1065,6 +1057,21 @@ func applyColumnMutation(
 	return nil
 }
 
+func addNotNullConstraintMutationForCol(tableDesc *tabledesc.Mutable, col catalog.Column) error {
+	info, err := tableDesc.GetConstraintInfo()
+	if err != nil {
+		return err
+	}
+	inuseNames := make(map[string]struct{}, len(info))
+	for k := range info {
+		inuseNames[k] = struct{}{}
+	}
+	check := tabledesc.MakeNotNullCheckConstraint(col.GetName(), col.GetID(), tableDesc.GetNextConstraintID(), inuseNames, descpb.ConstraintValidity_Validating)
+	tableDesc.AddNotNullMutation(check, descpb.DescriptorMutation_ADD)
+	tableDesc.NextConstraintID++
+	return nil
+}
+
 func labeledRowValues(cols []catalog.Column, values tree.Datums) string {
 	var s bytes.Buffer
 	for i := range cols {
diff --git a/pkg/sql/logictest/testdata/logic_test/virtual_columns b/pkg/sql/logictest/testdata/logic_test/virtual_columns
index f00fcde5cecb..48338da8dfcd 100644
--- a/pkg/sql/logictest/testdata/logic_test/virtual_columns
+++ b/pkg/sql/logictest/testdata/logic_test/virtual_columns
@@ -1256,3 +1256,28 @@ JOIN t75147 AS t2 ON
   AND t1.v1 = t2.v1
   AND t1.b = t2.b
 JOIN t75147 AS t3 ON t1.a = t3.a;
+
+
+# Regression test for #81675. The schema change logic must validate that
+# NOT NULL virtual columns indeed validate to non-NULL values for the existing
+# data in the table.
+subtest adding_not_null_virtual_column_validates_81675
+
+statement ok
+CREATE TABLE t81675 (i INT);
+INSERT INTO t81675 VALUES (1), (2), (NULL)
+
+statement ok
+ALTER TABLE t81675 ADD COLUMN j INT GENERATED ALWAYS AS (i+1) VIRTUAL;
+
+statement ok
+ALTER TABLE t81675 DROP COLUMN j;
+
+# Note that the pgcode here ought to be 23502 (not null violation) but is
+# instead 23514 (check constraint violation) because of implementation details.
+onlyif config local-legacy-schema-changer
+statement error pgcode 23514 validation of NOT NULL constraint failed
+ALTER TABLE t81675 ADD COLUMN j INT GENERATED ALWAYS AS (i+1) VIRTUAL NOT NULL;
+
+statement ok
+DROP TABLE t81675;