sql: track routine dependencies on columns in RETURNING clause

Previously, while building a routine, we did not explicitly add dependencies
to columns referenced in the RETURNING clause of mutations within the body
statements. This was not a problem for INSERT or UPSERT due to #145098, but
could cause a routine with UPDATE or DELETE statements to be broken after
a referenced column was dropped. This commit fixes the bug by tracking
column references in the RETURNING clause.

Fixes #146414

Release note (bug fix): Fixed a bug that allowed a column to be dropped
from a table even if it was referenced in the RETURNING clause of an UPDATE
or DELETE statement in a routine. In releases prior to v25.3, the fix will
be off by default, and can be enabled by setting the session variable
`use_improved_routine_dependency_tracking`.

Author: DrewKimball

pkg/sql/logictest/testdata/logic_test/udf_delete

@@ -302,3 +302,25 @@ SELECT * FROM u_a;
 2  b  20
 
 subtest end
+
+subtest regression_146414
+
+statement ok
+CREATE TABLE t146414 (
+  a INT NOT NULL,
+  b INT AS (a + 1) VIRTUAL
+)
+
+statement ok
+CREATE FUNCTION f146414() RETURNS INT LANGUAGE SQL AS $$
+  DELETE FROM t146414 WHERE a = 1 RETURNING b;
+  SELECT 1;
+$$;
+
+statement error pgcode 2BP01 pq: cannot drop column "b" because function "f146414" depends on it
+ALTER TABLE t146414 DROP COLUMN b;
+
+statement ok
+SELECT f146414()
+
+subtest end

pkg/sql/logictest/testdata/logic_test/udf_insert

@@ -264,3 +264,25 @@ statement error pgcode 23514 pq: failed to satisfy CHECK constraint \(b > 1:::IN
 SELECT f_checkb();
 
 subtest end
+
+subtest regression_146414
+
+statement ok
+CREATE TABLE t146414 (
+  a INT NOT NULL,
+  b INT AS (a + 1) VIRTUAL
+)
+
+statement ok
+CREATE FUNCTION f146414() RETURNS INT LANGUAGE SQL AS $$
+  INSERT INTO t146414 (a) VALUES (100) RETURNING b;
+  SELECT 1;
+$$;
+
+statement error pgcode 2BP01 pq: cannot drop column "b" because function "f146414" depends on it
+ALTER TABLE t146414 DROP COLUMN b;
+
+statement ok
+SELECT f146414()
+
+subtest end

pkg/sql/logictest/testdata/logic_test/udf_update

@@ -326,3 +326,25 @@ SELECT * FROM kv;
 5 5
 
 subtest end
+
+subtest regression_146414
+
+statement ok
+CREATE TABLE t146414 (
+  a INT NOT NULL,
+  b INT AS (a + 1) VIRTUAL
+)
+
+statement ok
+CREATE FUNCTION f146414() RETURNS INT LANGUAGE SQL AS $$
+  UPDATE t146414 SET a = a + 1 WHERE a = 1 RETURNING b;
+  SELECT 1;
+$$;
+
+statement error pgcode 2BP01 pq: cannot drop column "b" because function "f146414" depends on it
+ALTER TABLE t146414 DROP COLUMN b;
+
+statement ok
+SELECT f146414()
+
+subtest end

pkg/sql/logictest/testdata/logic_test/udf_upsert

@@ -238,3 +238,25 @@ statement error pgcode 22001 value too long for type CHAR\(3\)
 SELECT f_check_colerr_vals(NULL, 'abcd')
 
 subtest end
+
+subtest regression_146414
+
+statement ok
+CREATE TABLE t146414 (
+  a INT NOT NULL,
+  b INT AS (a + 1) VIRTUAL
+)
+
+statement ok
+CREATE FUNCTION f146414() RETURNS INT LANGUAGE SQL AS $$
+  UPSERT INTO t146414 (a) VALUES (100) RETURNING b;
+  SELECT 1;
+$$;
+
+statement error pgcode 2BP01 pq: cannot drop column "b" because function "f146414" depends on it
+ALTER TABLE t146414 DROP COLUMN b;
+
+statement ok
+SELECT f146414()
+
+subtest end

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -1702,7 +1702,7 @@ func (mb *mutationBuilder) buildReturningScopes(
 	//   4. Project columns in same order as defined in table schema.
 	//
 	inScope = mb.outScope.replace()
-	inScope.appendOrdinaryColumnsFromTable(mb.md.TableMeta(mb.tabID), &mb.alias)
+	mb.b.appendOrdinaryColumnsFromTable(inScope, mb.md.TableMeta(mb.tabID), &mb.alias)
 
 	// extraAccessibleCols contains all the columns that the RETURNING
 	// clause can refer to in addition to the table columns. This is useful for

pkg/sql/opt/optbuilder/partial_index.go

@@ -57,7 +57,7 @@ func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, scan
 	// Construct a scan as the tableScope expr so that logical properties of the
 	// scan can be used to fully normalize the index predicate.
 	tableScope := b.allocScope()
-	tableScope.appendOrdinaryColumnsFromTable(tabMeta, &tabMeta.Alias)
+	b.appendOrdinaryColumnsFromTable(tableScope, tabMeta, &tabMeta.Alias)
 
 	// If the optional scan argument was provided and it outputs all of the
 	// ordinary table columns, we use it as tableScope.expr. Otherwise, we must

pkg/sql/opt/optbuilder/scope.go

@@ -14,7 +14,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/build"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
-	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
@@ -230,28 +229,6 @@ func (s *scope) appendColumnsFromScope(src *scope) {
 	}
 }
 
-// appendOrdinaryColumnsFromTable adds all non-mutation and non-system columns from the
-// given table metadata to this scope.
-func (s *scope) appendOrdinaryColumnsFromTable(tabMeta *opt.TableMeta, alias *tree.TableName) {
-	tab := tabMeta.Table
-	if s.cols == nil {
-		s.cols = make([]scopeColumn, 0, tab.ColumnCount())
-	}
-	for i, n := 0, tab.ColumnCount(); i < n; i++ {
-		tabCol := tab.Column(i)
-		if tabCol.Kind() != cat.Ordinary {
-			continue
-		}
-		s.cols = append(s.cols, scopeColumn{
-			name:       scopeColName(tabCol.ColName()),
-			table:      *alias,
-			typ:        tabCol.DatumType(),
-			id:         tabMeta.MetaID.ColumnID(i),
-			visibility: columnVisibility(tabCol.Visibility()),
-		})
-	}
-}
-
 // appendColumns adds newly bound variables to this scope.
 // The expressions in the new columns are reset to nil.
 func (s *scope) appendColumns(cols []scopeColumn) {

pkg/sql/opt/optbuilder/select.go

@@ -828,7 +828,7 @@ func (b *Builder) addCheckConstraintsForTable(tabMeta *opt.TableMeta) {
 
 	// Create a scope that can be used for building the scalar expressions.
 	tableScope := b.allocScope()
-	tableScope.appendOrdinaryColumnsFromTable(tabMeta, &tabMeta.Alias)
+	b.appendOrdinaryColumnsFromTable(tableScope, tabMeta, &tabMeta.Alias)
 	// Synthesized CHECK expressions, e.g., for columns of ENUM types, may
 	// reference inaccessible columns. This can happen when the type of an
 	// indexed expression is an ENUM. We make these columns visible so that they
@@ -925,7 +925,7 @@ func (b *Builder) addComputedColsForTable(
 
 		if tableScope == nil {
 			tableScope = b.allocScope()
-			tableScope.appendOrdinaryColumnsFromTable(tabMeta, &tabMeta.Alias)
+			b.appendOrdinaryColumnsFromTable(tableScope, tabMeta, &tabMeta.Alias)
 		}
 
 		colType := tabCol.DatumType()

pkg/sql/opt/optbuilder/util.go

@@ -851,3 +851,41 @@ func (b *Builder) makePLpgSQLRaiseFn(args memo.ScalarListExpr) opt.ScalarExpr {
 		},
 	)
 }
+
+// appendOrdinaryColumnsFromTable adds all non-mutation and non-system columns
+// from the given table metadata to the given scope. References to these columns
+// will be tracked in the schema dependencies, if trackSchemaDeps is set.
+func (b *Builder) appendOrdinaryColumnsFromTable(
+	s *scope, tabMeta *opt.TableMeta, alias *tree.TableName,
+) {
+	tab := tabMeta.Table
+	if s.cols == nil {
+		s.cols = make([]scopeColumn, 0, tab.ColumnCount())
+	}
+	for i, n := 0, tab.ColumnCount(); i < n; i++ {
+		tabCol := tab.Column(i)
+		if tabCol.Kind() != cat.Ordinary {
+			continue
+		}
+		s.cols = append(s.cols, scopeColumn{
+			name:       scopeColName(tabCol.ColName()),
+			table:      *alias,
+			typ:        tabCol.DatumType(),
+			id:         tabMeta.MetaID.ColumnID(i),
+			visibility: columnVisibility(tabCol.Visibility()),
+		})
+	}
+	if b.trackSchemaDeps {
+		dep := opt.SchemaDep{DataSource: tab}
+		for i, n := 0, tab.ColumnCount(); i < n; i++ {
+			if tab.Column(i).Kind() != cat.Ordinary {
+				continue
+			}
+			if dep.ColumnIDToOrd == nil {
+				dep.ColumnIDToOrd = make(map[opt.ColumnID]int)
+			}
+			dep.ColumnIDToOrd[tabMeta.MetaID.ColumnID(i)] = i
+		}
+		b.schemaDeps = append(b.schemaDeps, dep)
+	}
+}