From 9dac4e49fa20c1ffa78c6ea8cfa662055d89cae5 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    | 57 +++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go
index cd0e33534fb9..8b22bccb08e5 100644
--- a/pkg/sql/alter_table.go
+++ b/pkg/sql/alter_table.go
@@ -858,6 +858,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 278a0a2047fd..b9a3fca2a205 100644
--- a/pkg/sql/logictest/testdata/logic_test/udf_in_constraints
+++ b/pkg/sql/logictest/testdata/logic_test/udf_in_constraints
@@ -425,3 +425,60 @@ CREATE FUNCTION f_circle() RETURNS INT LANGUAGE SQL AS $$ SELECT a FROM t_circle
 
 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 FUNCTION true_is_true() RETURNS BOOL LANGUAGE SQL AS $$ SELECT true = true $$;
+
+statement ok
+BEGIN;
+CREATE TABLE alter_add_check_constraint();
+ALTER TABLE alter_add_check_constraint ADD CONSTRAINT noop CHECK (true_is_true());
+COMMIT;
+
+query T
+SELECT create_statement FROM [SHOW CREATE TABLE alter_add_check_constraint];
+----
+CREATE TABLE public.alter_add_check_constraint (
+  rowid INT8 NOT VISIBLE NOT NULL DEFAULT unique_rowid(),
+  CONSTRAINT alter_add_check_constraint_pkey PRIMARY KEY (rowid ASC),
+  CONSTRAINT noop CHECK (public.true_is_true())
+)
+
+# 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, FAMILY "primary" (id, rowid));
+CREATE TABLE accounts_b (id UUID NOT NULL, FAMILY "primary" (id, rowid));
+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) $$;
+
+statement ok
+BEGIN;
+CREATE TABLE a (
+  account_id UUID NOT NULL,
+  account_type TEXT NOT NULL,
+  FAMILY "primary" (account_id, account_type, rowid)
+);
+ALTER TABLE a ADD CONSTRAINT is_a_or_b CHECK (is_a_or_b(account_id, account_type));
+COMMIT;
+
+query T
+SELECT create_statement FROM [SHOW CREATE TABLE a];
+----
+CREATE TABLE public.a (
+	account_id UUID NOT NULL,
+	account_type STRING NOT NULL,
+	rowid INT8 NOT VISIBLE NOT NULL DEFAULT unique_rowid(),
+	CONSTRAINT a_pkey PRIMARY KEY (rowid ASC),
+	CONSTRAINT is_a_or_b CHECK (public.is_a_or_b(account_id, account_type))
+)