From c2460a1c31bade57a08bf1bb29eb1158d5474665 Mon Sep 17 00:00:00 2001
From: Onder Kalaci <onder@citusdata.com>
Date: Wed, 23 Oct 2019 11:59:28 +0200
Subject: [PATCH] Add upgrade test for distributed functions

Simply make sure that Citus can pushdown functions after pg upgrade.
---
 src/test/regress/after_pg_upgrade_schedule    |  2 +-
 src/test/regress/before_pg_upgrade_schedule   |  2 +-
 .../upgrade_distributed_function_after.out    | 25 +++++++
 .../upgrade_distributed_function_before.out   | 69 +++++++++++++++++++
 .../upgrade_distributed_function_after.sql    | 11 +++
 .../upgrade_distributed_function_before.sql   | 39 +++++++++++
 6 files changed, 146 insertions(+), 2 deletions(-)
 create mode 100644 src/test/regress/expected/upgrade_distributed_function_after.out
 create mode 100644 src/test/regress/expected/upgrade_distributed_function_before.out
 create mode 100644 src/test/regress/sql/upgrade_distributed_function_after.sql
 create mode 100644 src/test/regress/sql/upgrade_distributed_function_before.sql

diff --git a/src/test/regress/after_pg_upgrade_schedule b/src/test/regress/after_pg_upgrade_schedule
index 5861b07aa..e06ae46a4 100644
--- a/src/test/regress/after_pg_upgrade_schedule
+++ b/src/test/regress/after_pg_upgrade_schedule
@@ -1 +1 @@
-test: upgrade_basic_after upgrade_type_after upgrade_ref2ref_after
+test: upgrade_basic_after upgrade_type_after upgrade_ref2ref_after upgrade_distributed_function_after
diff --git a/src/test/regress/before_pg_upgrade_schedule b/src/test/regress/before_pg_upgrade_schedule
index 48ee4424e..b527ff257 100644
--- a/src/test/regress/before_pg_upgrade_schedule
+++ b/src/test/regress/before_pg_upgrade_schedule
@@ -1,3 +1,3 @@
 # The basic tests runs analyze which depends on shard numbers
 test: upgrade_basic_before
-test: upgrade_type_before upgrade_ref2ref_before
+test: upgrade_type_before upgrade_ref2ref_before upgrade_distributed_function_before
diff --git a/src/test/regress/expected/upgrade_distributed_function_after.out b/src/test/regress/expected/upgrade_distributed_function_after.out
new file mode 100644
index 000000000..21d8456d3
--- /dev/null
+++ b/src/test/regress/expected/upgrade_distributed_function_after.out
@@ -0,0 +1,25 @@
+SET search_path TO upgrade_distributed_function_before, public;
+-- make sure that the metadata synced
+SELECT bool_and(metadatasynced) FROM pg_dist_node WHERE isactive AND noderole = 'primary';
+  bool_and
+ ----------
+  t
+ (1 row)
+ 
+SET client_min_messages TO DEBUG1;
+-- these are simple select functions, so doesn't have any
+-- side effects, safe to be called without BEGIN;..;ROLLBACK;
+SELECT count_values(11);
+DEBUG:  pushing down the function call
+  count_values
+ ---------------------------------------------------------------------
+             1
+ (1 row)
+ 
+SELECT count_values(12);
+DEBUG:  pushing down the function call
+  count_values
+ ---------------------------------------------------------------------
+             1
+ (1 row)
+
diff --git a/src/test/regress/expected/upgrade_distributed_function_before.out b/src/test/regress/expected/upgrade_distributed_function_before.out
new file mode 100644
index 000000000..7273196ed
--- /dev/null
+++ b/src/test/regress/expected/upgrade_distributed_function_before.out
@@ -0,0 +1,69 @@
+CREATE SCHEMA upgrade_distributed_function_before;
+SET search_path TO upgrade_distributed_function_before, public;
+SET citus.replication_model TO streaming;
+SET citus.shard_replication_factor TO 1;
+-- set sync intervals to less than 15s so wait_until_metadata_sync never times out
+ALTER SYSTEM SET citus.metadata_sync_interval TO 3000;
+ALTER SYSTEM SET citus.metadata_sync_retry_interval TO 500;
+SELECT pg_reload_conf();
+ pg_reload_conf 
+----------------
+ t
+(1 row)
+
+CREATE FUNCTION wait_until_metadata_sync(timeout INTEGER DEFAULT 15000)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
+CREATE TABLE t1 (a int PRIMARY KEY, b int);
+SELECT create_distributed_table('t1','a');
+ create_distributed_table 
+--------------------------
+ 
+(1 row)
+
+INSERT INTO t1 VALUES (11), (12);
+-- create a very simple distributed function colocated with the table
+CREATE FUNCTION count_values(input int) RETURNS int AS
+$$
+    DECLARE
+       cnt int := 0;
+    BEGIN
+        SELECT count(*) INTO cnt FROM upgrade_distributed_function_before.t1 WHERE a = $1;
+        RETURN cnt;
+    END;
+$$ LANGUAGE plpgsql;
+SELECT create_distributed_function('count_values(int)', '$1', colocate_with:='t1');
+ create_distributed_function 
+-----------------------------
+ 
+(1 row)
+
+-- make sure that the metadata synced before running the queries
+SELECT wait_until_metadata_sync();
+ wait_until_metadata_sync 
+--------------------------
+ 
+(1 row)
+
+SELECT bool_and(metadatasynced) FROM pg_dist_node WHERE isactive AND noderole = 'primary';
+ bool_and 
+----------
+ t
+(1 row)
+
+SET client_min_messages TO DEBUG1;
+SELECT count_values(11);
+DEBUG:  pushing down the function call
+ count_values 
+--------------
+            1
+(1 row)
+
+SELECT count_values(12);
+DEBUG:  pushing down the function call
+ count_values 
+--------------
+            1
+(1 row)
+
diff --git a/src/test/regress/sql/upgrade_distributed_function_after.sql b/src/test/regress/sql/upgrade_distributed_function_after.sql
new file mode 100644
index 000000000..cb11c58d0
--- /dev/null
+++ b/src/test/regress/sql/upgrade_distributed_function_after.sql
@@ -0,0 +1,11 @@
+SET search_path TO upgrade_distributed_function_before, public;
+
+-- make sure that the metadata synced
+SELECT bool_and(metadatasynced) FROM pg_dist_node WHERE isactive AND noderole = 'primary';
+SET client_min_messages TO DEBUG1;
+
+-- these are simple select functions, so doesn't have any
+-- side effects, safe to be called without BEGIN;..;ROLLBACK;
+SELECT count_values(11);
+SELECT count_values(12);
+
diff --git a/src/test/regress/sql/upgrade_distributed_function_before.sql b/src/test/regress/sql/upgrade_distributed_function_before.sql
new file mode 100644
index 000000000..d65b91d68
--- /dev/null
+++ b/src/test/regress/sql/upgrade_distributed_function_before.sql
@@ -0,0 +1,39 @@
+CREATE SCHEMA upgrade_distributed_function_before;
+SET search_path TO upgrade_distributed_function_before, public;
+SET citus.replication_model TO streaming;
+SET citus.shard_replication_factor TO 1;
+
+-- set sync intervals to less than 15s so wait_until_metadata_sync never times out
+ALTER SYSTEM SET citus.metadata_sync_interval TO 3000;
+ALTER SYSTEM SET citus.metadata_sync_retry_interval TO 500;
+SELECT pg_reload_conf();
+
+CREATE FUNCTION wait_until_metadata_sync(timeout INTEGER DEFAULT 15000)
+    RETURNS void
+    LANGUAGE C STRICT
+    AS 'citus';
+
+
+CREATE TABLE t1 (a int PRIMARY KEY, b int);
+SELECT create_distributed_table('t1','a');
+INSERT INTO t1 VALUES (11), (12);
+
+-- create a very simple distributed function colocated with the table
+CREATE FUNCTION count_values(input int) RETURNS int AS
+$$
+    DECLARE
+       cnt int := 0;
+    BEGIN
+        SELECT count(*) INTO cnt FROM upgrade_distributed_function_before.t1 WHERE a = $1;
+        RETURN cnt;
+    END;
+$$ LANGUAGE plpgsql;
+SELECT create_distributed_function('count_values(int)', '$1', colocate_with:='t1');
+
+-- make sure that the metadata synced before running the queries
+SELECT wait_until_metadata_sync();
+SELECT bool_and(metadatasynced) FROM pg_dist_node WHERE isactive AND noderole = 'primary';
+SET client_min_messages TO DEBUG1;
+
+SELECT count_values(11);
+SELECT count_values(12);