From 2b4be535de51749878d046d8a1db9659865a0dfa Mon Sep 17 00:00:00 2001
From: Onur Tirtir <onurcantirtir@gmail.com>
Date: Mon, 16 Jan 2023 18:11:00 +0300
Subject: [PATCH] Do clean-up before upgrade_columnar_before to make it
 runnable multiple times

So that flaky test detector can run upgrade_columnar_before.sql multiple
times.
---
 .../expected/upgrade_columnar_before.out      | 22 +++++++++++++++++
 .../regress/sql/upgrade_columnar_before.sql   | 24 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/src/test/regress/expected/upgrade_columnar_before.out b/src/test/regress/expected/upgrade_columnar_before.out
index 28c252e30..a4895c770 100644
--- a/src/test/regress/expected/upgrade_columnar_before.out
+++ b/src/test/regress/expected/upgrade_columnar_before.out
@@ -1,5 +1,27 @@
 -- Test if relying on topological sort of the objects, not their names, works
 -- fine when re-creating objects during pg_upgrade.
+DO
+$$
+BEGIN
+IF EXISTS (SELECT * FROM pg_namespace WHERE nspname = 'upgrade_columnar')
+THEN
+    -- Drop the the table leftover from the earlier run of
+    -- upgrade_columnar_before.sql. Similarly, drop the fake public schema
+    -- created before and rename the original one (renamed to citus_schema)
+    -- back to public.
+    --
+    -- This can only happen if upgrade_columnar_before.sql is run multiple
+    -- times for flaky test detection.
+    DROP TABLE citus_schema.new_columnar_table;
+    DROP SCHEMA public CASCADE;
+    ALTER SCHEMA citus_schema RENAME TO public;
+
+    SET LOCAL client_min_messages TO WARNING;
+    DROP SCHEMA upgrade_columnar CASCADE;
+END IF;
+END
+$$
+LANGUAGE plpgsql;
 ALTER SCHEMA public RENAME TO citus_schema;
 SET search_path TO citus_schema;
 -- As mentioned in https://github.com/citusdata/citus/issues/5447, it
diff --git a/src/test/regress/sql/upgrade_columnar_before.sql b/src/test/regress/sql/upgrade_columnar_before.sql
index ea71dba02..6f39f4234 100644
--- a/src/test/regress/sql/upgrade_columnar_before.sql
+++ b/src/test/regress/sql/upgrade_columnar_before.sql
@@ -1,5 +1,29 @@
 -- Test if relying on topological sort of the objects, not their names, works
 -- fine when re-creating objects during pg_upgrade.
+
+DO
+$$
+BEGIN
+IF EXISTS (SELECT * FROM pg_namespace WHERE nspname = 'upgrade_columnar')
+THEN
+    -- Drop the the table leftover from the earlier run of
+    -- upgrade_columnar_before.sql. Similarly, drop the fake public schema
+    -- created before and rename the original one (renamed to citus_schema)
+    -- back to public.
+    --
+    -- This can only happen if upgrade_columnar_before.sql is run multiple
+    -- times for flaky test detection.
+    DROP TABLE citus_schema.new_columnar_table;
+    DROP SCHEMA public CASCADE;
+    ALTER SCHEMA citus_schema RENAME TO public;
+
+    SET LOCAL client_min_messages TO WARNING;
+    DROP SCHEMA upgrade_columnar CASCADE;
+END IF;
+END
+$$
+LANGUAGE plpgsql;
+
 ALTER SCHEMA public RENAME TO citus_schema;
 SET search_path TO citus_schema;