From 6d8c5931d6b0a8a4b95d42070d7f03645da14451 Mon Sep 17 00:00:00 2001
From: Jelte Fennema <github-tech@jeltef.nl>
Date: Fri, 8 Apr 2022 15:09:33 +0200
Subject: [PATCH] Work around flaky test related to search_path (#5894)

For some reason search_path is not always set correctly on the worker
when calling a distributed function, this shows up when calling
`insert_document` in our distributed_triggers test. The underlying
reason is currently unknown and warrants deeper investigation.

Currently this test is one of the main causes for random CI failures. So
this change sets the search_path of each function explicitly, to reduce
these failures. So other devs can be more efficient, while I continue
investigating the root cause of this issue.

Also changes explicit `SET citus.enable_unsafe_triggers = false` to
`RESET citus.enable_unsafe_triggers` in passing.
---
 src/test/regress/expected/distributed_triggers.out | 13 +++++++++----
 src/test/regress/sql/distributed_triggers.sql      | 13 +++++++++----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/test/regress/expected/distributed_triggers.out b/src/test/regress/expected/distributed_triggers.out
index b4ad001f5..4cf0cfee8 100644
--- a/src/test/regress/expected/distributed_triggers.out
+++ b/src/test/regress/expected/distributed_triggers.out
@@ -314,6 +314,8 @@ ORDER BY shard_key_value, object_id, change_id;
 DROP TRIGGER bad_shardkey_record_change_trigger ON data;
 CREATE OR REPLACE FUNCTION remote_shardkey_record_change()
 RETURNS trigger
+SET search_path = 'distributed_triggers'
+LANGUAGE plpgsql
 AS $$
 DECLARE
     last_change_id bigint;
@@ -321,12 +323,15 @@ BEGIN
     UPDATE distributed_triggers.data_changes SET operation_type = TG_OP;
     RETURN NULL;
 END;
-$$ LANGUAGE plpgsql;
+$$;
 CREATE TRIGGER remote_shardkey_record_change_trigger
 AFTER INSERT OR UPDATE OR DELETE ON data
 FOR EACH ROW EXECUTE FUNCTION distributed_triggers.remote_shardkey_record_change();
 CREATE FUNCTION insert_document(key text, id text)
-RETURNS void LANGUAGE plpgsql AS $fn$
+RETURNS void
+LANGUAGE plpgsql
+SET search_path = 'distributed_triggers'
+AS $fn$
 BEGIN
 	INSERT INTO distributed_triggers.data VALUES (key, id, '{"id1":"id2"}');
 	DELETE FROM distributed_triggers.data where shard_key_value = key;
@@ -882,8 +887,8 @@ SELECT run_command_on_workers($$SELECT count(*) FROM pg_trigger WHERE tgname lik
 (2 rows)
 
 RESET client_min_messages;
-SET citus.enable_unsafe_triggers TO false;
-SELECT run_command_on_workers('ALTER SYSTEM SET citus.enable_unsafe_triggers TO false;');
+RESET citus.enable_unsafe_triggers;
+SELECT run_command_on_workers('ALTER SYSTEM RESET citus.enable_unsafe_triggers;');
        run_command_on_workers
 ---------------------------------------------------------------------
  (localhost,57637,t,"ALTER SYSTEM")
diff --git a/src/test/regress/sql/distributed_triggers.sql b/src/test/regress/sql/distributed_triggers.sql
index 33ecdbb8c..9acb1434c 100644
--- a/src/test/regress/sql/distributed_triggers.sql
+++ b/src/test/regress/sql/distributed_triggers.sql
@@ -175,6 +175,8 @@ DROP TRIGGER bad_shardkey_record_change_trigger ON data;
 
 CREATE OR REPLACE FUNCTION remote_shardkey_record_change()
 RETURNS trigger
+SET search_path = 'distributed_triggers'
+LANGUAGE plpgsql
 AS $$
 DECLARE
     last_change_id bigint;
@@ -182,14 +184,17 @@ BEGIN
     UPDATE distributed_triggers.data_changes SET operation_type = TG_OP;
     RETURN NULL;
 END;
-$$ LANGUAGE plpgsql;
+$$;
 
 CREATE TRIGGER remote_shardkey_record_change_trigger
 AFTER INSERT OR UPDATE OR DELETE ON data
 FOR EACH ROW EXECUTE FUNCTION distributed_triggers.remote_shardkey_record_change();
 
 CREATE FUNCTION insert_document(key text, id text)
-RETURNS void LANGUAGE plpgsql AS $fn$
+RETURNS void
+LANGUAGE plpgsql
+SET search_path = 'distributed_triggers'
+AS $fn$
 BEGIN
 	INSERT INTO distributed_triggers.data VALUES (key, id, '{"id1":"id2"}');
 	DELETE FROM distributed_triggers.data where shard_key_value = key;
@@ -465,8 +470,8 @@ SELECT tgrelid::regclass::text, tgname FROM pg_trigger WHERE tgname like 'insert
 SELECT run_command_on_workers($$SELECT count(*) FROM pg_trigger WHERE tgname like 'insert_99_trigger%'$$);
 
 RESET client_min_messages;
-SET citus.enable_unsafe_triggers TO false;
-SELECT run_command_on_workers('ALTER SYSTEM SET citus.enable_unsafe_triggers TO false;');
+RESET citus.enable_unsafe_triggers;
+SELECT run_command_on_workers('ALTER SYSTEM RESET citus.enable_unsafe_triggers;');
 SELECT run_command_on_workers('SELECT pg_reload_conf();');
 SET citus.log_remote_commands TO off;