Skip to content

Commit cef3bec

Browse files
DrewKimballDrewKimball
committed
sql: track SET columns in routine dependencies
Previously, we did not track dependencies on target columns in the SET clause of an UPDATE statement while creating a routine. This meant that it was possible to break a routine by dropping such a column if it wasn't referenced elsewhere in the routine. This commit fixes the bug by adding target columns that are explicitly referenced by an UPDATE statement to the routine dependencies. The new behavior is gated behind the session variable `prevent_update_set_column_drop`. Fixes #158898 Release note (bug fix): Fixed a bug that allowed a column to be dropped from its table despite being referenced by a routine. The bug could happen when the column was only referenced as a target column in the SET clause of an UPDATE statement within the routine. This fix only applies to newly-created routines. In versions prior to v26.1, the fix must be enabled by setting the session variable `prevent_update_set_column_drop`.
1 parent 205e210 commit cef3bec

File tree

11 files changed

+85
-1
lines changed

11 files changed

+85
-1
lines changed

‎pkg/sql/logictest/testdata/logic_test/information_schema‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4234,6 +4234,7 @@ plan_cache_mode auto
42344234
plpgsql_use_strict_into off
42354235
prefer_lookup_joins_for_fks off
42364236
prepared_statements_cache_size 0 B
4237+
prevent_update_set_column_drop off
42374238
propagate_admission_header_to_leaf_transactions on
42384239
propagate_input_ordering off
42394240
recursion_depth_limit 1000

‎pkg/sql/logictest/testdata/logic_test/pg_catalog‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3188,6 +3188,7 @@ plan_cache_mode auto
31883188
plpgsql_use_strict_into off NULL NULL NULL string
31893189
prefer_lookup_joins_for_fks off NULL NULL NULL string
31903190
prepared_statements_cache_size 0 B NULL NULL NULL string
3191+
prevent_update_set_column_drop off NULL NULL NULL string
31913192
propagate_admission_header_to_leaf_transactions on NULL NULL NULL string
31923193
propagate_input_ordering off NULL NULL NULL string
31933194
recursion_depth_limit 1000 NULL NULL NULL string
@@ -3438,6 +3439,7 @@ plan_cache_mode auto
34383439
plpgsql_use_strict_into off NULL user NULL off off
34393440
prefer_lookup_joins_for_fks off NULL user NULL off off
34403441
prepared_statements_cache_size 0 B B user NULL 0 B 0 B
3442+
prevent_update_set_column_drop off NULL user NULL off off
34413443
propagate_admission_header_to_leaf_transactions on NULL user NULL on on
34423444
propagate_input_ordering off NULL user NULL off off
34433445
recursion_depth_limit 1000 NULL user NULL 1000 1000
@@ -3679,6 +3681,7 @@ plan_cache_mode NULL NULL
36793681
plpgsql_use_strict_into NULL NULL NULL NULL NULL
36803682
prefer_lookup_joins_for_fks NULL NULL NULL NULL NULL
36813683
prepared_statements_cache_size NULL NULL NULL NULL NULL
3684+
prevent_update_set_column_drop NULL NULL NULL NULL NULL
36823685
propagate_admission_header_to_leaf_transactions NULL NULL NULL NULL NULL
36833686
propagate_input_ordering NULL NULL NULL NULL NULL
36843687
recursion_depth_limit NULL NULL NULL NULL NULL

‎pkg/sql/logictest/testdata/logic_test/show_source‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ plan_cache_mode auto
205205
plpgsql_use_strict_into off
206206
prefer_lookup_joins_for_fks off
207207
prepared_statements_cache_size 0 B
208+
prevent_update_set_column_drop off
208209
propagate_admission_header_to_leaf_transactions on
209210
propagate_input_ordering off
210211
recursion_depth_limit 1000

‎pkg/sql/logictest/testdata/logic_test/udf_update‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -386,3 +386,29 @@ statement error pgcode 2BP01 pq: cannot drop column "a" because function "f_upda
386386
ALTER TABLE table_drop DROP COLUMN a;
387387

388388
subtest end
389+
390+
# Prevent drop of column used in SET clause within a routine.
391+
subtest regression_158898
392+
393+
statement ok
394+
SET prevent_update_set_column_drop = true;
395+
396+
statement ok
397+
CREATE TABLE t158898 (col1 INT8 NOT NULL);
398+
399+
statement ok
400+
CREATE OR REPLACE PROCEDURE p158898() LANGUAGE PLpgSQL AS $proc$
401+
BEGIN UPDATE t158898 SET col1 = 1; END;
402+
$proc$;
403+
404+
statement error pgcode 2BP01 pq: cannot drop column "col1" because function "p158898" depends on it
405+
ALTER TABLE t158898 DROP COLUMN col1;
406+
407+
statement ok
408+
DROP PROCEDURE p158898;
409+
DROP TABLE t158898;
410+
411+
statement ok
412+
RESET prevent_update_set_column_drop;
413+
414+
subtest end

‎pkg/sql/opt/memo/memo.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ type Memo struct {
217217
useMaxFrequencySelectivity bool
218218
usingHintInjection bool
219219
useSwapMutations bool
220+
preventUpdateSetColumnDrop bool
220221

221222
// txnIsoLevel is the isolation level under which the plan was created. This
222223
// affects the planning of some locking operations, so it must be included in
@@ -334,6 +335,7 @@ func (m *Memo) Init(ctx context.Context, evalCtx *eval.Context) {
334335
useMaxFrequencySelectivity: evalCtx.SessionData().OptimizerUseMaxFrequencySelectivity,
335336
usingHintInjection: evalCtx.Planner != nil && evalCtx.Planner.UsingHintInjection(),
336337
useSwapMutations: evalCtx.SessionData().UseSwapMutations,
338+
preventUpdateSetColumnDrop: evalCtx.SessionData().PreventUpdateSetColumnDrop,
337339
txnIsoLevel: evalCtx.TxnIsoLevel,
338340
}
339341
m.metadata.Init()
@@ -515,6 +517,7 @@ func (m *Memo) IsStale(
515517
m.useMaxFrequencySelectivity != evalCtx.SessionData().OptimizerUseMaxFrequencySelectivity ||
516518
m.usingHintInjection != (evalCtx.Planner != nil && evalCtx.Planner.UsingHintInjection()) ||
517519
m.useSwapMutations != evalCtx.SessionData().UseSwapMutations ||
520+
m.preventUpdateSetColumnDrop != evalCtx.SessionData().PreventUpdateSetColumnDrop ||
518521
m.txnIsoLevel != evalCtx.TxnIsoLevel {
519522
return true, nil
520523
}

‎pkg/sql/opt/memo/memo_test.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -611,6 +611,12 @@ func TestMemoIsStale(t *testing.T) {
611611
evalCtx.SessionData().OptimizerClampInequalitySelectivity = false
612612
notStale()
613613

614+
// Stale prevent_update_set_column_drop.
615+
evalCtx.SessionData().PreventUpdateSetColumnDrop = true
616+
stale()
617+
evalCtx.SessionData().PreventUpdateSetColumnDrop = false
618+
notStale()
619+
614620
// User no longer has access to view.
615621
catalog.View(tree.NewTableNameWithSchema("t", catconstants.PublicSchemaName, "abcview")).Revoked = true
616622
_, err = o.Memo().IsStale(ctx, &evalCtx, catalog)

‎pkg/sql/opt/optbuilder/mutation_builder.go‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,13 @@ type mutationBuilder struct {
7575
// targetColSet contains the same column IDs as targetColList, but as a set.
7676
targetColSet opt.ColSet
7777

78+
// explicitTargetColOrds contains the ordinals of the columns in targetColSet
79+
// that were explicitly specified by the user, e.g. in the target column list
80+
// of an INSERT or the SET clause of an UPDATE. This is used to determine
81+
// which columns to add to the list of dependencies if b.trackSchemaDeps is
82+
// true.
83+
explicitTargetColOrds intsets.Fast
84+
7885
// insertColIDs lists the input column IDs providing values to insert. Its
7986
// length is always equal to the number of columns in the target table,
8087
// including mutation columns. Table columns which will not have values
@@ -603,8 +610,18 @@ func (mb *mutationBuilder) addTargetCol(ord int) {
603610
"multiple assignments to the same column %q", tabCol.ColName()))
604611
}
605612
mb.targetColSet.Add(colID)
606-
607613
mb.targetColList = append(mb.targetColList, colID)
614+
mb.explicitTargetColOrds.Add(ord)
615+
}
616+
617+
// trackTargetColDeps adds column dependencies for the target columns that were
618+
// explicitly specified by the user.
619+
func (mb *mutationBuilder) trackTargetColDeps() {
620+
if mb.b.trackSchemaDeps && mb.b.evalCtx.SessionData().PreventUpdateSetColumnDrop {
621+
dep := opt.SchemaDep{DataSource: mb.tab}
622+
dep.ColumnOrdinals.CopyFrom(mb.explicitTargetColOrds)
623+
mb.b.schemaDeps = append(mb.b.schemaDeps, dep)
624+
}
608625
}
609626

610627
// extractValuesInput tests whether the given input is a VALUES clause with no

‎pkg/sql/opt/optbuilder/update.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,8 @@ func (b *Builder) buildUpdate(upd *tree.Update, inScope *scope) (outScope *scope
124124
}
125125
mb.buildUpdate(returningExpr, cat.PolicyScopeUpdate, &exprColRefs)
126126

127+
mb.trackTargetColDeps()
128+
127129
return mb.outScope
128130
}
129131

‎pkg/sql/sessiondatapb/local_only_session_data.proto‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -762,6 +762,10 @@ message LocalOnlySessionData {
762762
// UseSwapMutations, when true, enables use of the update swap and delete swap
763763
// operators.
764764
bool use_swap_mutations = 194;
765+
// PreventUpdateSetColumnDrop, when true, enables a fix to ensure that columns
766+
// referenced only as target columns in an UPDATE statement within a routine
767+
// are correctly tracked, preventing a drop that breaks the routine.
768+
bool prevent_update_set_column_drop = 195;
765769

766770
///////////////////////////////////////////////////////////////////////////
767771
// WARNING: consider whether a session parameter you're adding needs to //

‎pkg/sql/sessionmutator/mutator.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1114,3 +1114,7 @@ func (m *SessionDataMutator) SetCanaryStatsMode(val sessiondatapb.CanaryStatsMod
11141114
func (m *SessionDataMutator) SetUseSwapMutations(val bool) {
11151115
m.Data.UseSwapMutations = val
11161116
}
1117+
1118+
func (m *SessionDataMutator) SetPreventUpdateSetColumnDrop(val bool) {
1119+
m.Data.PreventUpdateSetColumnDrop = val
1120+
}

0 commit comments

Comments
 (0)