From a6435b7f6b9ff92409d3688231d1a27bfd1bf69c Mon Sep 17 00:00:00 2001
From: Naisila Puka <37271756+naisila@users.noreply.github.com>
Date: Thu, 11 Mar 2021 14:52:40 +0300
Subject: [PATCH] Fix upgrade and downgrade paths for
 master/citus_update_table_statistics (#4805)

(cherry picked from commit 71a9f45513392c956e4e13b8d2085b8bf8d22069)
---
 .../distributed/sql/citus--10.0-2--10.0-3.sql | 10 ++++++++
 src/backend/distributed/sql/citus--8.0-1.sql  | 21 ++++++++++++----
 .../sql/downgrades/citus--10.0-3--10.0-2.sql  | 24 +++++++++++++++++++
 .../citus_update_table_statistics/10.0-1.sql  | 19 +++++++++++----
 .../citus_update_table_statistics/10.0-3.sql  |  6 +++++
 .../citus_update_table_statistics/latest.sql  |  2 +-
 6 files changed, 72 insertions(+), 10 deletions(-)
 create mode 100644 src/backend/distributed/sql/citus--10.0-2--10.0-3.sql
 create mode 100644 src/backend/distributed/sql/downgrades/citus--10.0-3--10.0-2.sql
 create mode 100644 src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-3.sql

diff --git a/src/backend/distributed/sql/citus--10.0-2--10.0-3.sql b/src/backend/distributed/sql/citus--10.0-2--10.0-3.sql
new file mode 100644
index 000000000..d4f1312fb
--- /dev/null
+++ b/src/backend/distributed/sql/citus--10.0-2--10.0-3.sql
@@ -0,0 +1,10 @@
+-- citus--10.0-2--10.0-3
+
+#include "udfs/citus_update_table_statistics/10.0-3.sql"
+
+CREATE OR REPLACE FUNCTION master_update_table_statistics(relation regclass)
+RETURNS VOID
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$citus_update_table_statistics$$;
+COMMENT ON FUNCTION pg_catalog.master_update_table_statistics(regclass)
+	IS 'updates shard statistics of the given table';
diff --git a/src/backend/distributed/sql/citus--8.0-1.sql b/src/backend/distributed/sql/citus--8.0-1.sql
index eced5cc60..5647e01fd 100644
--- a/src/backend/distributed/sql/citus--8.0-1.sql
+++ b/src/backend/distributed/sql/citus--8.0-1.sql
@@ -1394,11 +1394,22 @@ COMMENT ON FUNCTION master_update_node(node_id int, new_node_name text, new_node
 
 -- shard statistics
 CREATE OR REPLACE FUNCTION master_update_table_statistics(relation regclass)
-RETURNS VOID
-    LANGUAGE C STRICT
-    AS 'MODULE_PATHNAME', $$citus_update_table_statistics$$;
-COMMENT ON FUNCTION pg_catalog.master_update_table_statistics(regclass)
-	IS 'updates shard statistics of the given table';
+RETURNS VOID AS $$
+DECLARE
+	colocated_tables regclass[];
+BEGIN
+	SELECT get_colocated_table_array(relation) INTO colocated_tables;
+
+	PERFORM
+		master_update_shard_statistics(shardid)
+	FROM
+		pg_dist_shard
+	WHERE
+		logicalrelid = ANY (colocated_tables);
+END;
+$$ LANGUAGE 'plpgsql';
+COMMENT ON FUNCTION master_update_table_statistics(regclass)
+	IS 'updates shard statistics of the given table and its colocated tables';
 
 CREATE OR REPLACE FUNCTION get_colocated_shard_array(bigint)
 	RETURNS BIGINT[]
diff --git a/src/backend/distributed/sql/downgrades/citus--10.0-3--10.0-2.sql b/src/backend/distributed/sql/downgrades/citus--10.0-3--10.0-2.sql
new file mode 100644
index 000000000..f75a4642f
--- /dev/null
+++ b/src/backend/distributed/sql/downgrades/citus--10.0-3--10.0-2.sql
@@ -0,0 +1,24 @@
+-- citus--10.0-3--10.0-2
+-- this is a downgrade path that will revert the changes made in citus--10.0-2--10.0-3.sql
+
+DROP FUNCTION pg_catalog.citus_update_table_statistics(regclass);
+
+#include "../udfs/citus_update_table_statistics/10.0-1.sql"
+
+CREATE OR REPLACE FUNCTION master_update_table_statistics(relation regclass)
+RETURNS VOID AS $$
+DECLARE
+	colocated_tables regclass[];
+BEGIN
+	SELECT get_colocated_table_array(relation) INTO colocated_tables;
+
+	PERFORM
+		master_update_shard_statistics(shardid)
+	FROM
+		pg_dist_shard
+	WHERE
+		logicalrelid = ANY (colocated_tables);
+END;
+$$ LANGUAGE 'plpgsql';
+COMMENT ON FUNCTION master_update_table_statistics(regclass)
+	IS 'updates shard statistics of the given table and its colocated tables';
diff --git a/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-1.sql b/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-1.sql
index d4e07d24d..13f4a680e 100644
--- a/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-1.sql
+++ b/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-1.sql
@@ -1,6 +1,17 @@
 CREATE FUNCTION pg_catalog.citus_update_table_statistics(relation regclass)
-	RETURNS VOID
-    LANGUAGE C STRICT
-    AS 'MODULE_PATHNAME', $$citus_update_table_statistics$$;
+RETURNS VOID AS $$
+DECLARE
+	colocated_tables regclass[];
+BEGIN
+	SELECT get_colocated_table_array(relation) INTO colocated_tables;
+
+	PERFORM
+		master_update_shard_statistics(shardid)
+	FROM
+		pg_dist_shard
+	WHERE
+		logicalrelid = ANY (colocated_tables);
+END;
+$$ LANGUAGE 'plpgsql';
 COMMENT ON FUNCTION pg_catalog.citus_update_table_statistics(regclass)
-	IS 'updates shard statistics of the given table';
+	IS 'updates shard statistics of the given table and its colocated tables';
diff --git a/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-3.sql b/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-3.sql
new file mode 100644
index 000000000..753b1ce7f
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_update_table_statistics/10.0-3.sql
@@ -0,0 +1,6 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_update_table_statistics(relation regclass)
+	RETURNS VOID
+    LANGUAGE C STRICT
+    AS 'MODULE_PATHNAME', $$citus_update_table_statistics$$;
+COMMENT ON FUNCTION pg_catalog.citus_update_table_statistics(regclass)
+	IS 'updates shard statistics of the given table';
diff --git a/src/backend/distributed/sql/udfs/citus_update_table_statistics/latest.sql b/src/backend/distributed/sql/udfs/citus_update_table_statistics/latest.sql
index d4e07d24d..753b1ce7f 100644
--- a/src/backend/distributed/sql/udfs/citus_update_table_statistics/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_update_table_statistics/latest.sql
@@ -1,4 +1,4 @@
-CREATE FUNCTION pg_catalog.citus_update_table_statistics(relation regclass)
+CREATE OR REPLACE FUNCTION pg_catalog.citus_update_table_statistics(relation regclass)
 	RETURNS VOID
     LANGUAGE C STRICT
     AS 'MODULE_PATHNAME', $$citus_update_table_statistics$$;