From 2f201a98bb7be0f6ffa78e8588395ac58fd28f06 Mon Sep 17 00:00:00 2001
From: Marco Slot <marco.slot@gmail.com>
Date: Tue, 22 Oct 2019 19:42:30 +0200
Subject: [PATCH] Revoke usage from the citus schema

(cherry picked from commit 04040e0a37d0e967304b00b55e03b90d6f037ee3)
---
 .../distributed/sql/citus--9.0-1--9.0-2.sql   | 10 +++++++++
 .../distributed_tables_colocated/9.0-2.sql    | 22 +++++++++++++++++++
 .../distributed_tables_colocated/latest.sql   | 22 +++++++++++++++++++
 3 files changed, 54 insertions(+)
 create mode 100644 src/backend/distributed/sql/citus--9.0-1--9.0-2.sql
 create mode 100644 src/backend/distributed/sql/udfs/distributed_tables_colocated/9.0-2.sql
 create mode 100644 src/backend/distributed/sql/udfs/distributed_tables_colocated/latest.sql

diff --git a/src/backend/distributed/sql/citus--9.0-1--9.0-2.sql b/src/backend/distributed/sql/citus--9.0-1--9.0-2.sql
new file mode 100644
index 000000000..15f4a3eaf
--- /dev/null
+++ b/src/backend/distributed/sql/citus--9.0-1--9.0-2.sql
@@ -0,0 +1,10 @@
+-- Using the citus schema is a bad idea since many environments use "citus"
+-- as the main user and the "citus" schema then sits in front of the
+-- search_path.
+REVOKE USAGE ON SCHEMA citus FROM public;
+
+-- redefine distributed_tables_colocated to avoid using citus schema
+#include "udfs/distributed_tables_colocated/9.0-2.sql"
+
+-- type was used in old version of distributed_tables_colocated
+DROP TYPE citus.colocation_placement_type;
diff --git a/src/backend/distributed/sql/udfs/distributed_tables_colocated/9.0-2.sql b/src/backend/distributed/sql/udfs/distributed_tables_colocated/9.0-2.sql
new file mode 100644
index 000000000..ed6bb7444
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/distributed_tables_colocated/9.0-2.sql
@@ -0,0 +1,22 @@
+--
+-- distributed_tables_colocated returns true if given tables are co-located, false otherwise.
+-- The function checks shard definitions, matches shard placements for given tables.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.distributed_tables_colocated(table1 regclass,
+																   table2 regclass)
+RETURNS bool
+LANGUAGE plpgsql
+AS $function$
+DECLARE
+	table1_colocationid int;
+	table2_colocationid int;
+BEGIN
+	SELECT colocationid INTO table1_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table1;
+
+	SELECT colocationid INTO table2_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table2;
+
+	RETURN table1_colocationid = table2_colocationid;
+END;
+$function$;
diff --git a/src/backend/distributed/sql/udfs/distributed_tables_colocated/latest.sql b/src/backend/distributed/sql/udfs/distributed_tables_colocated/latest.sql
new file mode 100644
index 000000000..ed6bb7444
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/distributed_tables_colocated/latest.sql
@@ -0,0 +1,22 @@
+--
+-- distributed_tables_colocated returns true if given tables are co-located, false otherwise.
+-- The function checks shard definitions, matches shard placements for given tables.
+--
+CREATE OR REPLACE FUNCTION pg_catalog.distributed_tables_colocated(table1 regclass,
+																   table2 regclass)
+RETURNS bool
+LANGUAGE plpgsql
+AS $function$
+DECLARE
+	table1_colocationid int;
+	table2_colocationid int;
+BEGIN
+	SELECT colocationid INTO table1_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table1;
+
+	SELECT colocationid INTO table2_colocationid
+	FROM pg_catalog.pg_dist_partition WHERE logicalrelid = table2;
+
+	RETURN table1_colocationid = table2_colocationid;
+END;
+$function$;