Add pg_partition_root to display top-most parent of a partition tree

This is useful when looking at partition trees with multiple layers, and
combined with pg_partition_tree, it provides the possibility to show up
an entire tree by just knowing one member at any level.

Author: Michael Paquier
Reviewed-by: Álvaro Herrera, Amit Langote
Discussion: https://postgr.es/m/20181207014015.GP2407@paquier.xyz

5 files changed, 142 insertions, 11 deletions

diff --git a/src/backend/utils/adt/partitionfuncs.c b/src/backend/utils/adt/partitionfuncs.c
index 5cdf4a4524e..ffd66b64394 100644
--- a/src/backend/utils/adt/partitionfuncs.c
+++ b/src/backend/utils/adt/partitionfuncs.c
@@ -25,6 +25,33 @@
 #include "utils/lsyscache.h"
 #include "utils/syscache.h"
 
+/*
+ * Checks if a given relation can be part of a partition tree.  Returns
+ * false if the relation cannot be processed, in which case it is up to
+ * the caller to decide what to do, by either raising an error or doing
+ * something else.
+ */
+static bool
+check_rel_can_be_partition(Oid relid)
+{
+	char		relkind;
+
+	/* Check if relation exists */
+	if (!SearchSysCacheExists1(RELOID, ObjectIdGetDatum(relid)))
+		return false;
+
+	relkind = get_rel_relkind(relid);
+
+	/* Only allow relation types that can appear in partition trees. */
+	if (relkind != RELKIND_RELATION &&
+		relkind != RELKIND_FOREIGN_TABLE &&
+		relkind != RELKIND_INDEX &&
+		relkind != RELKIND_PARTITIONED_TABLE &&
+		relkind != RELKIND_PARTITIONED_INDEX)
+		return false;
+
+	return true;
+}
 
 /*
  * pg_partition_tree
@@ -39,19 +66,10 @@ pg_partition_tree(PG_FUNCTION_ARGS)
 {
 #define PG_PARTITION_TREE_COLS	4
 	Oid			rootrelid = PG_GETARG_OID(0);
-	char		relkind = get_rel_relkind(rootrelid);
 	FuncCallContext *funcctx;
 	ListCell  **next;
 
-	if (!SearchSysCacheExists1(RELOID, ObjectIdGetDatum(rootrelid)))
-		PG_RETURN_NULL();
-
-	/* Return NULL for relation types that cannot appear in partition trees */
-	if (relkind != RELKIND_RELATION &&
-		relkind != RELKIND_FOREIGN_TABLE &&
-		relkind != RELKIND_INDEX &&
-		relkind != RELKIND_PARTITIONED_TABLE &&
-		relkind != RELKIND_PARTITIONED_INDEX)
+	if (!check_rel_can_be_partition(rootrelid))
 		PG_RETURN_NULL();
 
 	/* stuff done only on the first call of the function */
@@ -153,3 +171,40 @@ pg_partition_tree(PG_FUNCTION_ARGS)
 	/* done when there are no more elements left */
 	SRF_RETURN_DONE(funcctx);
 }
+
+/*
+ * pg_partition_root
+ *
+ * Returns the top-most parent of the partition tree to which a given
+ * relation belongs, or NULL if it's not (or cannot be) part of any
+ * partition tree.
+ */
+Datum
+pg_partition_root(PG_FUNCTION_ARGS)
+{
+	Oid			relid = PG_GETARG_OID(0);
+	Oid			rootrelid;
+	List	   *ancestors;
+
+	if (!check_rel_can_be_partition(relid))
+		PG_RETURN_NULL();
+
+	/*
+	 * If the relation is not a partition (it may be the partition parent),
+	 * return itself as a result.
+	 */
+	if (!get_rel_relispartition(relid))
+		PG_RETURN_OID(relid);
+
+	/* Fetch the top-most parent */
+	ancestors = get_partition_ancestors(relid);
+	rootrelid = llast_oid(ancestors);
+	list_free(ancestors);
+
+	/*
+	 * "rootrelid" must contain a valid OID, given that the input relation is
+	 * a valid partition tree member as checked above.
+	 */
+	Assert(OidIsValid(rootrelid));
+	PG_RETURN_OID(rootrelid);
+}
diff --git a/src/include/catalog/catversion.h b/src/include/catalog/catversion.h
index df505668776..5eb4f6f172a 100644
--- a/src/include/catalog/catversion.h
+++ b/src/include/catalog/catversion.h
@@ -53,6 +53,6 @@
  */
 
 /*							yyyymmddN */
-#define CATALOG_VERSION_NO	201902071
+#define CATALOG_VERSION_NO	201902081
 
 #endif
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index b8de13f03b9..93e3e16f01c 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -10509,4 +10509,9 @@
   proargnames => '{rootrelid,relid,parentrelid,isleaf,level}',
   prosrc => 'pg_partition_tree' },
 
+# function to get the top-most partition root parent
+{ oid => '3424', descr => 'get top-most partition root parent',
+  proname => 'pg_partition_root', prorettype => 'regclass',
+  proargtypes => 'regclass', prosrc => 'pg_partition_root' },
+
 ]
diff --git a/src/test/regress/expected/partition_info.out b/src/test/regress/expected/partition_info.out
index 3e15e02f8da..a884df976fd 100644
--- a/src/test/regress/expected/partition_info.out
+++ b/src/test/regress/expected/partition_info.out
@@ -12,6 +12,18 @@ SELECT * FROM pg_partition_tree(0);
        |             |        |      
 (1 row)
 
+SELECT pg_partition_root(NULL);
+ pg_partition_root 
+-------------------
+ 
+(1 row)
+
+SELECT pg_partition_root(0);
+ pg_partition_root 
+-------------------
+ 
+(1 row)
+
 -- Test table partition trees
 CREATE TABLE ptif_test (a int, b int) PARTITION BY range (a);
 CREATE TABLE ptif_test0 PARTITION OF ptif_test
@@ -66,6 +78,20 @@ SELECT relid, parentrelid, level, isleaf
  ptif_test01 | ptif_test0  |     0 | t
 (1 row)
 
+-- List all members using pg_partition_root with leaf table reference
+SELECT relid, parentrelid, level, isleaf
+  FROM pg_partition_tree(pg_partition_root('ptif_test01')) p
+  JOIN pg_class c ON (p.relid = c.oid);
+    relid    | parentrelid | level | isleaf 
+-------------+-------------+-------+--------
+ ptif_test   |             |     0 | f
+ ptif_test0  | ptif_test   |     1 | f
+ ptif_test1  | ptif_test   |     1 | f
+ ptif_test2  | ptif_test   |     1 | t
+ ptif_test01 | ptif_test0  |     2 | t
+ ptif_test11 | ptif_test1  |     2 | t
+(6 rows)
+
 -- List all indexes members of the tree
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_test_index');
@@ -98,6 +124,20 @@ SELECT relid, parentrelid, level, isleaf
  ptif_test01_index | ptif_test0_index |     0 | t
 (1 row)
 
+-- List all members using pg_partition_root with leaf index reference
+SELECT relid, parentrelid, level, isleaf
+  FROM pg_partition_tree(pg_partition_root('ptif_test01_index')) p
+  JOIN pg_class c ON (p.relid = c.oid);
+       relid       |   parentrelid    | level | isleaf 
+-------------------+------------------+-------+--------
+ ptif_test_index   |                  |     0 | f
+ ptif_test0_index  | ptif_test_index  |     1 | f
+ ptif_test1_index  | ptif_test_index  |     1 | f
+ ptif_test2_index  | ptif_test_index  |     1 | t
+ ptif_test01_index | ptif_test0_index |     2 | t
+ ptif_test11_index | ptif_test1_index |     2 | t
+(6 rows)
+
 DROP TABLE ptif_test;
 -- Table that is not part of any partition tree is the only member listed.
 CREATE TABLE ptif_normal_table(a int);
@@ -108,6 +148,12 @@ SELECT relid, parentrelid, level, isleaf
  ptif_normal_table |             |     0 | t
 (1 row)
 
+SELECT pg_partition_root('ptif_normal_table');
+ pg_partition_root 
+-------------------
+ ptif_normal_table
+(1 row)
+
 DROP TABLE ptif_normal_table;
 -- Various partitioning-related functions return NULL if passed relations
 -- of types that cannot be part of a partition tree; for example, views,
@@ -126,5 +172,17 @@ SELECT * FROM pg_partition_tree('ptif_test_matview');
        |             |        |      
 (1 row)
 
+SELECT pg_partition_root('ptif_test_view');
+ pg_partition_root 
+-------------------
+ 
+(1 row)
+
+SELECT pg_partition_root('ptif_test_matview');
+ pg_partition_root 
+-------------------
+ 
+(1 row)
+
 DROP VIEW ptif_test_view;
 DROP MATERIALIZED VIEW ptif_test_matview;
diff --git a/src/test/regress/sql/partition_info.sql b/src/test/regress/sql/partition_info.sql
index 5d6e647146d..119b90afe45 100644
--- a/src/test/regress/sql/partition_info.sql
+++ b/src/test/regress/sql/partition_info.sql
@@ -3,6 +3,8 @@
 --
 SELECT * FROM pg_partition_tree(NULL);
 SELECT * FROM pg_partition_tree(0);
+SELECT pg_partition_root(NULL);
+SELECT pg_partition_root(0);
 
 -- Test table partition trees
 CREATE TABLE ptif_test (a int, b int) PARTITION BY range (a);
@@ -39,6 +41,10 @@ SELECT relid, parentrelid, level, isleaf
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_test01') p
   JOIN pg_class c ON (p.relid = c.oid);
+-- List all members using pg_partition_root with leaf table reference
+SELECT relid, parentrelid, level, isleaf
+  FROM pg_partition_tree(pg_partition_root('ptif_test01')) p
+  JOIN pg_class c ON (p.relid = c.oid);
 
 -- List all indexes members of the tree
 SELECT relid, parentrelid, level, isleaf
@@ -51,6 +57,10 @@ SELECT relid, parentrelid, level, isleaf
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_test01_index') p
   JOIN pg_class c ON (p.relid = c.oid);
+-- List all members using pg_partition_root with leaf index reference
+SELECT relid, parentrelid, level, isleaf
+  FROM pg_partition_tree(pg_partition_root('ptif_test01_index')) p
+  JOIN pg_class c ON (p.relid = c.oid);
 
 DROP TABLE ptif_test;
 
@@ -58,6 +68,7 @@ DROP TABLE ptif_test;
 CREATE TABLE ptif_normal_table(a int);
 SELECT relid, parentrelid, level, isleaf
   FROM pg_partition_tree('ptif_normal_table');
+SELECT pg_partition_root('ptif_normal_table');
 DROP TABLE ptif_normal_table;
 
 -- Various partitioning-related functions return NULL if passed relations
@@ -67,5 +78,7 @@ CREATE VIEW ptif_test_view AS SELECT 1;
 CREATE MATERIALIZED VIEW ptif_test_matview AS SELECT 1;
 SELECT * FROM pg_partition_tree('ptif_test_view');
 SELECT * FROM pg_partition_tree('ptif_test_matview');
+SELECT pg_partition_root('ptif_test_view');
+SELECT pg_partition_root('ptif_test_matview');
 DROP VIEW ptif_test_view;
 DROP MATERIALIZED VIEW ptif_test_matview;