sql: still more routine dependency improvements

pkg/sql/logictest/testdata/logic_test/check_constraints

@@ -494,3 +494,33 @@ ALTER TABLE t_91697 ADD CHECK (a < 123);
 
 statement error pgcode 23514 pq: failed to satisfy CHECK constraint \(a < 'public.t67100a'::REGCLASS\)
 INSERT INTO t_91697 VALUES (321);
+
+# Regression test for #158154. Don't add transitive column dependencies to
+# routines that build check constraints.
+subtest regression_158154
+
+statement ok
+CREATE TABLE t158154 (a INT, b INT, CONSTRAINT foo CHECK (b > 0));
+
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  INSERT INTO t158154 (a) VALUES (1);
+  UPDATE t158154 SET a = a + 1;
+$$;
+
+statement ok
+CALL p158154();
+
+# The procedure does not directly depend on column b, so dropping column b
+# should be allowed after dropping the check constraint.
+statement ok
+ALTER TABLE t158154 DROP CONSTRAINT foo, DROP COLUMN b;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+subtest end

pkg/sql/logictest/testdata/logic_test/enums

@@ -2116,3 +2116,35 @@ WHERE stat->>'columns' = '["a"]'
     "null_count": 0,
     "row_count": 5
 }
+
+# Regression test for #158154. Don't add transitive column dependencies to
+# routines that build check constraints for enum columns.
+subtest regression_158154
+
+statement ok
+CREATE TYPE typ158154 AS ENUM ('foo', 'bar');
+
+statement ok
+CREATE TABLE t158154 (a INT, b typ158154);
+
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  INSERT INTO t158154 (a) VALUES (1);
+  UPDATE t158154 SET a = a + 1;
+$$;
+
+statement ok
+CALL p158154();
+
+statement ok
+ALTER TABLE t158154 DROP COLUMN b;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+DROP TYPE typ158154;
+
+subtest end

pkg/sql/logictest/testdata/logic_test/hash_sharded_index

@@ -1435,3 +1435,37 @@ SELECT crdb_internal_d_shard_16 AS shard, count(*) < 20 FROM decimals GROUP BY 1
 15  false
 
 subtest end
+
+# Regression test for #158154. Don't add transitive column dependencies to
+# routines that build check constraints.
+subtest regression_158154
+
+statement ok
+CREATE TABLE t158154 (col1 INT8 NOT NULL, col2 TIMESTAMP NOT NULL);
+
+statement ok
+CREATE INDEX idx1 ON t158154 (col2 ASC) USING HASH WITH (bucket_count = 6);
+
+statement ok
+CREATE OR REPLACE PROCEDURE p158154 () LANGUAGE PLpgSQL AS $proc$
+  BEGIN
+    UPDATE t158154 SET col1 = col1 + 1;
+  END;
+$proc$;
+
+statement ok
+CALL p158154();
+
+# The procedure does not directly depend on the hash-sharded index computed
+# column, so dropping the index should succeed.
+statement ok
+DROP INDEX t158154@idx1;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+subtest end

pkg/sql/logictest/testdata/logic_test/partial_index

@@ -2497,3 +2497,97 @@ DROP INDEX idx97551;
 statement ok
 DROP TABLE t97551;
 DROP TYPE enum97551;
+
+subtest end
+
+# Regression test for #158154. Don't add transitive column dependencies to
+# routines that build partial index predicate expressions.
+subtest regression_158154
+
+# Simple case with INSERT and UPDATE into a table with a partial index.
+statement ok
+CREATE TABLE t158154 (a INT, b INT, INDEX b_idx (b) WHERE (b > 0));
+
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  INSERT INTO t158154 (a) VALUES (1);
+  UPDATE t158154 SET a = a + 1;
+$$;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP INDEX t158154@b_idx;
+
+statement ok
+ALTER TABLE t158154 DROP COLUMN b;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+# Case with a partial index used as an arbiter.
+statement ok
+CREATE TABLE t158154 (a INT, b INT);
+
+statement ok
+CREATE UNIQUE INDEX a_idx ON t158154 (a) WHERE (b > 0);
+
+# WHERE false implies any partial index predicate.
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  INSERT INTO t158154 (a) VALUES (1) ON CONFLICT (a) WHERE (false) DO NOTHING;
+  INSERT INTO t158154 (a) VALUES (2) ON CONFLICT (a) WHERE (false) DO UPDATE SET a = EXCLUDED.a + 1;
+$$;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP INDEX t158154@a_idx;
+
+statement ok
+ALTER TABLE t158154 DROP COLUMN b;
+
+statement ok
+CREATE UNIQUE INDEX a_idx ON t158154 (a);
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+# Verify that a SELECT statement inside a procedure does not add transitive
+# dependencies from the partial index predicate.
+statement ok
+CREATE TABLE t158154 (a INT, b INT, INDEX b_idx (b) WHERE (b > 0));
+
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  SELECT a FROM t158154;
+  SELECT count(*) FROM t158154 WHERE a > 5;
+$$;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP INDEX t158154@b_idx;
+
+statement ok
+ALTER TABLE t158154 DROP COLUMN b;
+
+statement ok
+CALL p158154();
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+subtest end

pkg/sql/logictest/testdata/logic_test/row_level_security

@@ -5930,4 +5930,61 @@ DROP TABLE accounts, non_rls;
 statement ok
 DROP USER alice;
 
+statement ok
+RESET SESSION row_security;
+
+subtest end
+
+# Regression test for #158154. Don't add transitive column dependencies to
+# routines that build RLS constraints.
+subtest regression_158154
+
+statement ok
+CREATE TABLE t158154 (c0 INT, c1 TEXT DEFAULT 'foobarbaz', FAMILY (c0, c1));
+
+statement ok
+CREATE ROLE user_158154
+
+statement ok
+GRANT ALL ON t158154 TO user_158154;
+
+statement ok
+ALTER TABLE t158154 ENABLE ROW LEVEL SECURITY, FORCE ROW LEVEL SECURITY;
+
+statement ok
+ALTER TABLE t158154 OWNER TO user_158154;
+
+statement ok
+SET ROLE user_158154;
+
+statement ok
+CREATE POLICY IF NOT EXISTS p1 on t158154 WITH CHECK (c0 > 0 AND c1 = 'foobarbaz');
+
+statement ok
+CREATE PROCEDURE p158154() LANGUAGE SQL AS $$
+  UPDATE t158154 SET c0 = 1 WHERE true;
+$$;
+
+statement ok
+CALL p158154();
+
+# The procedure does not directly depend on column b, so dropping column b
+# should be allowed after dropping the RLS constraint.
+statement ok
+DROP POLICY p1 on t158154;
+ALTER TABLE t158154 DROP COLUMN c1;
+
+statement ok
+CALL p158154();
+
+statement ok
+SET ROLE root;
+
+statement ok
+DROP PROCEDURE p158154;
+DROP TABLE t158154;
+
+statement ok
+DROP ROLE user_158154;
+
 subtest end

pkg/sql/opt/optbuilder/insert.go

@@ -846,6 +846,13 @@ func (mb *mutationBuilder) buildInputForDoNothing(
 	// TODO(andyk): do we need to do more here?
 	mb.outScope.ordering = nil
 
+	if !mb.arbiters.Empty() &&
+		mb.b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+		// Avoid resolving transitive dependencies on columns and UDTs referenced by
+		// any partial index predicates among the arbiters.
+		defer mb.b.DisableSchemaDepTracking()()
+	}
+
 	// Create an anti-join for each arbiter.
 	mb.arbiters.ForEach(func(
 		name string, conflictOrds intsets.Fast, pred tree.Expr, canaryOrd int, uniqueWithoutIndex bool, uniqueOrd int,
@@ -902,6 +909,13 @@ func (mb *mutationBuilder) buildInputForUpsert(
 	// Ignore any ordering requested by the input.
 	mb.outScope.ordering = nil
 
+	if !mb.arbiters.Empty() &&
+		mb.b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+		// Avoid resolving transitive dependencies on columns and UDTs referenced by
+		// any partial index predicates among the arbiters.
+		defer mb.b.DisableSchemaDepTracking()()
+	}
+
 	// Create an UpsertDistinctOn and a left-join for the single arbiter.
 	mb.arbiters.ForEach(func(
 		name string, conflictOrds intsets.Fast, pred tree.Expr, canaryOrd int, uniqueWithoutIndex bool, uniqueOrd int,

pkg/sql/opt/optbuilder/mutation_builder.go

@@ -1097,6 +1097,13 @@ func (mb *mutationBuilder) maybeAddRegionColLookup(op opt.Operator) {
 func (mb *mutationBuilder) addCheckConstraintCols(
 	isUpdate bool, policyCmdScope cat.PolicyCommandScope, includeSelectPolicies bool,
 ) {
+	if mb.b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+		// Avoid adding unnecessary dependencies on columns that are referenced by
+		// check expressions. We only need to track columns that were explicitly
+		// specified by the user, e.g. those in SET, WHERE or RETURNING clause, or
+		// the target columns of an INSERT.
+		defer mb.b.DisableSchemaDepTracking()()
+	}
 	if mb.tab.CheckCount() != 0 {
 		projectionsScope := mb.outScope.replace()
 		projectionsScope.appendColumnsFromScope(mb.outScope)
@@ -1508,6 +1515,11 @@ func (mb *mutationBuilder) projectPartialIndexPutAndDelCols() {
 // projectPartialIndexDelCols, or projectPartialIndexPutAndDelCols.
 func (mb *mutationBuilder) projectPartialIndexColsImpl(putScope, delScope *scope) {
 	if partialIndexCount(mb.tab) > 0 {
+		if mb.b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+			// Avoid unnecessary dependencies on columns and UDTs referenced by the
+			// partial index expression.
+			defer mb.b.DisableSchemaDepTracking()()
+		}
 		projectionScope := mb.outScope.replace()
 		projectionScope.appendColumnsFromScope(mb.outScope)
 

pkg/sql/opt/optbuilder/partial_index.go

@@ -27,11 +27,17 @@ import (
 // scan. A scan and its logical properties are required in order to fully
 // normalize the partial index predicates.
 func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, scan memo.RelExpr) {
-	// We do not want to track view/function deps here, otherwise a view/function
-	// depending on a table with a partial index predicate using an UDT will
-	// result in a type dependency being added between the view/function and the
-	// UDT.
-	defer b.DisableSchemaDepTracking()()
+	if !b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+		// We do not want to track view/function deps here, otherwise a
+		// view/function depending on a table with a partial index predicate using
+		// a UDT will result in a type dependency being added between the
+		// view/function and the UDT.
+		//
+		// This is the legacy path; with the session setting on, we will disable
+		// dependency tracking in buildPartialIndexPredicate below.
+		defer b.DisableSchemaDepTracking()()
+	}
+
 	tab := tabMeta.Table
 	numIndexes := tab.DeletableIndexCount()
 
@@ -115,6 +121,14 @@ func (b *Builder) addPartialIndexPredicatesForTable(tabMeta *opt.TableMeta, scan
 func (b *Builder) buildPartialIndexPredicate(
 	tabMeta *opt.TableMeta, tableScope *scope, expr tree.Expr, context string,
 ) (memo.FiltersExpr, error) {
+	if b.evalCtx.SessionData().UseImprovedRoutineDepsTriggersAndComputedCols {
+		// We do not want to track view/function deps here, otherwise a view/function
+		// depending on a table with a partial index predicate will add transitive
+		// dependencies on any UDTs or columns referenced by the predicate. The
+		// partial index will already prevent dropping such UDTs or columns.
+		defer b.DisableSchemaDepTracking()()
+	}
+
 	texpr := tableScope.resolveAndRequireType(expr, types.Bool)
 
 	var scalar opt.ScalarExpr

pkg/sql/opt/optbuilder/util.go

@@ -882,7 +882,8 @@ func (b *Builder) appendOrdinaryColumnsFromTable(
 			visibility: columnVisibility(tabCol.Visibility()),
 		})
 	}
-	if b.trackSchemaDeps && b.evalCtx.SessionData().UseImprovedRoutineDependencyTracking {
+	if !tab.IsVirtualTable() && b.trackSchemaDeps &&
+		b.evalCtx.SessionData().UseImprovedRoutineDependencyTracking {
 		dep := opt.SchemaDep{DataSource: tab}
 		for i, n := 0, tab.ColumnCount(); i < n; i++ {
 			if tab.Column(i).Kind() != cat.Ordinary {

pkg/sql/schemachanger/scbuild/testdata/drop_index

@@ -114,7 +114,7 @@ DROP INDEX idx3 CASCADE
 - [[UserPrivileges:{DescID: 105, Name: root}, ABSENT], PUBLIC]
   {descriptorId: 105, privileges: "2", userName: root, withGrantOption: "2"}
 - [[View:{DescID: 105}, ABSENT], PUBLIC]
-  {forwardReferences: [{columnIds: [2], indexId: 6, toId: 104}], usesRelationIds: [104], viewId: 105}
+  {forwardReferences: [{columnIds: [2], toId: 104}, {columnIds: [2], indexId: 6, toId: 104}], usesRelationIds: [104], viewId: 105}
 - [[SchemaChild:{DescID: 105, ReferencedDescID: 101}, ABSENT], PUBLIC]
   {childObjectId: 105, schemaId: 101}
 - [[Column:{DescID: 105, ColumnID: 1}, ABSENT], PUBLIC]