Merge pull request #83353 from ajwerner/backport22.1-81693

release-22.1: sql: enforce NOT NULL when adding a virtual computed column
cockroachdb · Jun 27, 2022 · 8f7d8c3 · 8f7d8c3
2 parents da742f8 + 98020b6
commit 8f7d8c3
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 10 deletions.
diff --git 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
@@ -995,17 +995,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() {
@@ -1043,7 +1035,10 @@ func applyColumnMutation(
 
 		// Add a check constraint equivalent to the non-null constraint and drop
 		// it in the schema changer.
-		check := tabledesc.MakeNotNullCheckConstraint(col.GetName(), col.GetID(), tableDesc.GetNextConstraintID(), inuseNames, descpb.ConstraintValidity_Dropping)
+		check := tabledesc.MakeNotNullCheckConstraint(
+			col.GetName(), col.GetID(), tableDesc.GetNextConstraintID(),
+			inuseNames, descpb.ConstraintValidity_Dropping,
+		)
 		tableDesc.Checks = append(tableDesc.Checks, check)
 		tableDesc.NextConstraintID++
 		tableDesc.AddNotNullMutation(check, descpb.DescriptorMutation_DROP)
@@ -1062,6 +1057,24 @@ 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
@@ -1256,3 +1256,27 @@ 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;