From a20cc3b36afc7f9b87e9827bda03e26c1f743f3b Mon Sep 17 00:00:00 2001
From: SaitTalhaNisanci <s.talhanisanci@gmail.com>
Date: Fri, 28 May 2021 11:33:48 +0300
Subject: [PATCH] Only consider shard state 1 in citus shards (#4970)

---
 .../distributed/sql/citus--10.0-3--10.1-1.sql |  1 +
 .../sql/downgrades/citus--10.1-1--10.0-3.sql  | 36 +++++++++++++++++++
 .../sql/udfs/citus_shards/10.1-1.sql          | 36 +++++++++++++++++++
 .../sql/udfs/citus_shards/latest.sql          |  5 +--
 4 files changed, 76 insertions(+), 2 deletions(-)
 create mode 100644 src/backend/distributed/sql/udfs/citus_shards/10.1-1.sql

diff --git a/src/backend/distributed/sql/citus--10.0-3--10.1-1.sql b/src/backend/distributed/sql/citus--10.0-3--10.1-1.sql
index d787a8fdb..0720750e2 100644
--- a/src/backend/distributed/sql/citus--10.0-3--10.1-1.sql
+++ b/src/backend/distributed/sql/citus--10.0-3--10.1-1.sql
@@ -43,3 +43,4 @@ SET repmodel = 's'
 WHERE repmodel = 'c'
     AND partmethod = 'h'
     AND logicalrelid NOT IN (SELECT * FROM replicated_relations);
+#include "udfs/citus_shards/10.1-1.sql"
diff --git a/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql b/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
index c9fd0c895..668c9eb7b 100644
--- a/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
+++ b/src/backend/distributed/sql/downgrades/citus--10.1-1--10.0-3.sql
@@ -41,3 +41,39 @@ ALTER TABLE pg_catalog.pg_dist_rebalance_strategy DROP COLUMN improvement_thresh
 -- so it doesn't have a DROP. This is why we DROP manually here.
 DROP FUNCTION pg_catalog.get_rebalance_progress;
 #include "../udfs/get_rebalance_progress/9.0-1.sql"
+
+CREATE OR REPLACE VIEW pg_catalog.citus_shards AS
+WITH shard_sizes AS (SELECT * FROM pg_catalog.citus_shard_sizes())
+SELECT
+     pg_dist_shard.logicalrelid AS table_name,
+     pg_dist_shard.shardid,
+     shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) as shard_name,
+     CASE WHEN partkey IS NOT NULL THEN 'distributed' WHEN repmodel = 't' THEN 'reference' ELSE 'local' END AS citus_table_type,
+     colocationid AS colocation_id,
+     pg_dist_node.nodename,
+     pg_dist_node.nodeport,
+     (SELECT size FROM shard_sizes WHERE
+       shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) = table_name
+       OR
+       'public.' || shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) = table_name
+      LIMIT 1) as shard_size
+FROM
+   pg_dist_shard
+JOIN
+   pg_dist_placement
+ON
+   pg_dist_shard.shardid = pg_dist_placement.shardid
+JOIN
+   pg_dist_node
+ON
+   pg_dist_placement.groupid = pg_dist_node.groupid
+JOIN
+   pg_dist_partition
+ON
+   pg_dist_partition.logicalrelid = pg_dist_shard.logicalrelid
+ORDER BY
+   pg_dist_shard.logicalrelid::text, shardid
+;
+
+#include "../udfs/citus_finish_pg_upgrade/10.0-1.sql"
+
diff --git a/src/backend/distributed/sql/udfs/citus_shards/10.1-1.sql b/src/backend/distributed/sql/udfs/citus_shards/10.1-1.sql
new file mode 100644
index 000000000..4ce7fc40e
--- /dev/null
+++ b/src/backend/distributed/sql/udfs/citus_shards/10.1-1.sql
@@ -0,0 +1,36 @@
+CREATE OR REPLACE VIEW pg_catalog.citus_shards AS
+WITH shard_sizes AS (SELECT * FROM pg_catalog.citus_shard_sizes())
+SELECT
+     pg_dist_shard.logicalrelid AS table_name,
+     pg_dist_shard.shardid,
+     shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) as shard_name,
+     CASE WHEN partkey IS NOT NULL THEN 'distributed' WHEN repmodel = 't' THEN 'reference' ELSE 'local' END AS citus_table_type,
+     colocationid AS colocation_id,
+     pg_dist_node.nodename,
+     pg_dist_node.nodeport,
+     (SELECT size FROM shard_sizes WHERE
+       shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) = table_name
+       OR
+       'public.' || shard_name(pg_dist_shard.logicalrelid, pg_dist_shard.shardid) = table_name
+      LIMIT 1) as shard_size
+FROM
+   pg_dist_shard
+JOIN
+   pg_dist_placement
+ON
+   pg_dist_shard.shardid = pg_dist_placement.shardid
+JOIN
+   pg_dist_node
+ON
+   pg_dist_placement.groupid = pg_dist_node.groupid
+JOIN
+   pg_dist_partition
+ON
+   pg_dist_partition.logicalrelid = pg_dist_shard.logicalrelid
+WHERE
+   pg_dist_placement.shardstate = 1
+ORDER BY
+   pg_dist_shard.logicalrelid::text, shardid
+;
+
+GRANT SELECT ON pg_catalog.citus_shards TO public;
diff --git a/src/backend/distributed/sql/udfs/citus_shards/latest.sql b/src/backend/distributed/sql/udfs/citus_shards/latest.sql
index 9f09199ba..4ce7fc40e 100644
--- a/src/backend/distributed/sql/udfs/citus_shards/latest.sql
+++ b/src/backend/distributed/sql/udfs/citus_shards/latest.sql
@@ -1,4 +1,4 @@
-CREATE OR REPLACE VIEW citus.citus_shards AS
+CREATE OR REPLACE VIEW pg_catalog.citus_shards AS
 WITH shard_sizes AS (SELECT * FROM pg_catalog.citus_shard_sizes())
 SELECT
      pg_dist_shard.logicalrelid AS table_name,
@@ -27,9 +27,10 @@ JOIN
    pg_dist_partition
 ON
    pg_dist_partition.logicalrelid = pg_dist_shard.logicalrelid
+WHERE
+   pg_dist_placement.shardstate = 1
 ORDER BY
    pg_dist_shard.logicalrelid::text, shardid
 ;
 
-ALTER VIEW citus.citus_shards SET SCHEMA pg_catalog;
 GRANT SELECT ON pg_catalog.citus_shards TO public;