sql/schemachanger: fix lost dependencies in ALTER POLICY expressions

When altering only the USING expression of a policy that also had a
WITH CHECK expression, dependencies from the WITH CHECK expression
were incorrectly lost. This occurred because the upsertPolicyExpressions
function overwrote dependency sets instead of unioning them.

Fixes #153191

Epic: none

Release note (bug fix): Fixed ALTER POLICY incorrectly dropping dependency
tracking for functions, sequences, or types in policy expressions.

Author: spilchen

pkg/sql/logictest/testdata/logic_test/row_level_security

@@ -5533,6 +5533,64 @@ REVOKE SYSTEM CREATEDB FROM can_createdb_global;
 statement ok
 DROP ROLE can_createdb_global;
 
+# Test that dependencies are preserved when altering only the USING expression
+# of a policy that has both USING and WITH CHECK expressions. This is a repro
+# for GH issue 153191.
+subtest bug_fix_lost_dependencies_alter_policy_#153191
+
+statement ok
+CREATE TABLE dep_test (id INT PRIMARY KEY, value INT);
+
+statement ok
+CREATE SEQUENCE seq_with_check;
+
+statement ok
+CREATE FUNCTION func_using_old(n INT) RETURNS BOOL AS $$ SELECT n > 0; $$ LANGUAGE SQL;
+
+statement ok
+CREATE FUNCTION func_using_new(n INT) RETURNS BOOL AS $$ SELECT n >= 0; $$ LANGUAGE SQL;
+
+# Create a olicy that have different dependencies for USING and WITH CHECK
+# expressions.
+statement ok
+CREATE POLICY test_policy ON dep_test
+  FOR ALL
+  USING (func_using_old(value))
+  WITH CHECK (nextval('seq_with_check') < 1000);
+
+# Verify initial dependencies are tracked.
+statement error pq: cannot drop function "func_using_old" because other objects
+DROP FUNCTION func_using_old;
+
+statement error pq: cannot drop sequence seq_with_check because other objects depend on it
+DROP SEQUENCE seq_with_check;
+
+# Alter only the USING expression, which should preserve the WITH CHECK
+# dependency.
+statement ok
+ALTER POLICY test_policy ON dep_test USING (func_using_new(value));
+
+# Ensure dependency on WITH CHECK still exists.
+statement error pq: cannot drop sequence seq_with_check because other objects depend on it
+DROP SEQUENCE seq_with_check;
+
+# Ensure old dependency on USING doesn't exist anymore.
+statement ok
+DROP FUNCTION func_using_old;
+
+# And the new USING function should be protected.
+statement error pq: cannot drop function "func_using_new" because other objects
+DROP FUNCTION func_using_new;
+
+statement ok
+DROP TABLE dep_test CASCADE;
+
+statement ok
+DROP SEQUENCE seq_with_check;
+
+statement ok
+DROP FUNCTION func_using_new;
+
 subtest end
 
 subtest filter_pushdown_leaks

pkg/sql/schemachanger/scbuild/internal/scbuildstmt/create_policy.go

@@ -204,9 +204,9 @@ func upsertPolicyExpressions(
 			} else {
 				// If we aren't dropping the old expression, we need to keep track of
 				// the dependencies in case we replace PolicyDeps.
-				usesTypeIDs = catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...)
-				usesRelationIDs = catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...)
-				usesFunctionIDs = catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...)
+				usesTypeIDs = usesTypeIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...))
+				usesRelationIDs = usesRelationIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...))
+				usesFunctionIDs = usesFunctionIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...))
 			}
 		}
 	})
@@ -217,9 +217,9 @@ func upsertPolicyExpressions(
 			PolicyID:   policyID,
 			Expression: *expr,
 		})
-		usesTypeIDs = catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...)
-		usesRelationIDs = catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...)
-		usesFunctionIDs = catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...)
+		usesTypeIDs = usesTypeIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...))
+		usesRelationIDs = usesRelationIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...))
+		usesFunctionIDs = usesFunctionIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...))
 	}
 
 	policyElems.ForEach(func(current scpb.Status, target scpb.TargetStatus, e scpb.Element) {
@@ -230,9 +230,9 @@ func upsertPolicyExpressions(
 			} else {
 				// If we aren't dropping the old expression, we need to keep track of
 				// the dependencies in case we replace PolicyDeps.
-				usesTypeIDs = catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...)
-				usesRelationIDs = catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...)
-				usesFunctionIDs = catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...)
+				usesTypeIDs = usesTypeIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesTypeIDs...))
+				usesRelationIDs = usesRelationIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesSequenceIDs...))
+				usesFunctionIDs = usesFunctionIDs.Union(catalog.MakeDescriptorIDSet(expr.UsesFunctionIDs...))
 			}
 		}
 	})