From 84bb802685c33d4f7cdf55f49fd7754278eaa4cc 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       | 24 ++++++++++++++++++
 3 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/add_column.go b/pkg/sql/add_column.go
index cee5cbcda560..ba3e9aaccffa 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 1d4248c3b4e8..9681f7b6cc5f 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 80533aeab439..10a3cfb53718 100644
--- a/pkg/sql/logictest/testdata/logic_test/virtual_columns
+++ b/pkg/sql/logictest/testdata/logic_test/virtual_columns
@@ -1312,3 +1312,27 @@ CREATE TABLE public.t_added (
 
 statement ok
 DROP TABLE t_added
+
+# 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;