From e583cb96e09b7a3122d8365362662b9249200332 Mon Sep 17 00:00:00 2001
From: Chris Seto <chriskseto@gmail.com>
Date: Wed, 6 Sep 2023 17:48:19 +0000
Subject: [PATCH] sql: update UDF references in `ALTER TABLE`

Previously, `ALTER TABLE ... ADD CONSTRAINT` statements  that referenced
a UDF did not appropriately update UDFs' `DependedOnBy` field to
reference the table/constraint that utilized it. This would surface as
the validation error reported in #109414, `relation X: depends-on
function X has no corresponding depended-on-by back reference`. This
commit ensures that the back references are updated by copying a hook
from `CREATE TABLE`.

Epic: None
Fixes: #109414
Release note (bug fix): `ALTER TABLE ... ADD CONSTRAINT CHECK ...`
statements that utilize a UDF in the CHECK no longer cause a validation
error.
---
 pkg/sql/alter_table.go                        |  8 ++++++
 .../testdata/logic_test/udf_in_constraints    | 28 +++++++++++++++++++
 2 files changed, 36 insertions(+)

diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go
index 1d677f9dbe62..e075afa4757e 100644
--- a/pkg/sql/alter_table.go
+++ b/pkg/sql/alter_table.go
@@ -865,6 +865,14 @@ func (n *alterTableNode) startExec(params runParams) error {
 		return err
 	}
 
+	// Replace all UDF names with OIDs in check constraints and update back
+	// references in functions used.
+	for _, ck := range n.tableDesc.CheckConstraints() {
+		if err := params.p.updateFunctionReferencesForCheck(params.ctx, n.tableDesc, ck.CheckDesc()); err != nil {
+			return err
+		}
+	}
+
 	// Record this table alteration in the event log. This is an auditable log
 	// event and is recorded in the same transaction as the table descriptor
 	// update.
diff --git a/pkg/sql/logictest/testdata/logic_test/udf_in_constraints b/pkg/sql/logictest/testdata/logic_test/udf_in_constraints
index 5810ce0f65f4..bf28116165c5 100644
--- a/pkg/sql/logictest/testdata/logic_test/udf_in_constraints
+++ b/pkg/sql/logictest/testdata/logic_test/udf_in_constraints
@@ -436,3 +436,31 @@ CREATE FUNCTION f_circle() RETURNS INT LANGUAGE SQL AS $$ SELECT a FROM t_circle
 # TODO(107369): This does not appear to error in postgres.
 statement error .*cannot add dependency from descriptor \d+ to function f_circle \(\d+\) because there will be a dependency cycle
 ALTER TABLE t_circle ADD CONSTRAINT ckb CHECK (b + f_circle() > 1);
+
+# Reproduction/regression test for https://github.com/cockroachdb/cockroach/issues/109414
+# Adding a check constraint with alter table doesn't appropriately update back references.
+subtest issue-109414-minimal
+
+statement ok
+CREATE TABLE accounts (id UUID NOT NULL);
+CREATE FUNCTION has_account(account_id UUID) RETURNS BOOL LANGUAGE SQL AS $$ SELECT EXISTS(SELECT * FROM accounts WHERE id = account_id) $$;
+CREATE TABLE account_ref ( account_id UUID NOT NULL );
+ALTER TABLE account_ref ADD CONSTRAINT diy_fk CHECK (has_account(account_id));
+
+# This is the original Reproduction of #109414. It's quite interesting in and
+# of itself, so it's been included as it may catch other regressions that the
+# minimal case wouldn't.
+subtest issue-109414-full
+statement ok
+CREATE TABLE accounts_a (id UUID NOT NULL);
+CREATE TABLE accounts_b (id UUID NOT NULL);
+CREATE FUNCTION is_a_or_b(account_id UUID, account_type TEXT) RETURNS BOOL LANGUAGE SQL AS $$ SELECT (CASE
+	WHEN account_type = 'type_a' THEN (SELECT EXISTS(SELECT * FROM accounts_a WHERE id = account_id))
+	WHEN account_type = 'type_b' THEN (SELECT EXISTS(SELECT * FROM accounts_b WHERE id = account_id))
+	ELSE false
+END) $$;
+CREATE TABLE a (
+  account_id UUID NOT NULL,
+	account_type TEXT NOT NULL
+);
+ALTER TABLE a ADD CONSTRAINT is_a_or_b CHECK (is_a_or_b(account_id, account_type));