Fix upgrade and downgrade paths for master/citus_update_table_statistics (#4805)

(cherry picked from commit 71a9f45513)

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';

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[]

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';

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';

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';

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$$;