From 4ef9bfe458019800c4aee08c1322652597147513 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Tue, 16 Jan 2018 19:02:13 +0900
Subject: [PATCH v2] Add assorted partition reporting functions

---
 doc/src/sgml/func.sgml                     |  34 ++++++++
 src/backend/catalog/partition.c            | 129 ++++++++++++++++++++++++++++-
 src/backend/utils/cache/lsyscache.c        |  22 +++++
 src/include/catalog/partition.h            |   1 +
 src/include/catalog/pg_proc.dat            |  18 ++++
 src/include/utils/lsyscache.h              |   1 +
 src/test/regress/expected/create_table.out |  60 ++++++++++++++
 src/test/regress/sql/create_table.sql      |  24 ++++++
 8 files changed, 285 insertions(+), 4 deletions(-)

diff --git a/doc/src/sgml/func.sgml b/doc/src/sgml/func.sgml
index 5dce8ef178..df621d1e17 100644
--- a/doc/src/sgml/func.sgml
+++ b/doc/src/sgml/func.sgml
@@ -19995,6 +19995,40 @@ postgres=# SELECT * FROM pg_walfile_name_offset(pg_stop_backup());
     The function returns the number of new collation objects it created.
    </para>
 
+   <table id="functions-info-partition">
+    <title>Partitioning Information Functions</title>
+    <tgroup cols="3">
+     <thead>
+      <row><entry>Name</entry> <entry>Return Type</entry> <entry>Description</entry></row>
+     </thead>
+
+     <tbody>
+      <row>
+       <entry><literal><function>pg_partition_root_parent(<parameter>regclass</parameter>)</function></literal></entry>
+       <entry><type>regclass</type></entry>
+       <entry>get root table of partition tree of which the table is part</entry>
+      </row>
+      <row>
+       <entry><literal><function>pg_partition_parent(<parameter>regclass</parameter>)</function></literal></entry>
+       <entry><type>regclass</type></entry>
+       <entry>get parent table if the table is a partition, <literal>NULL</literal> otherwise</entry>
+      </row>
+      <row>
+       <entry><literal><function>pg_partition_tree_tables(<parameter>regclass</parameter>)</function></literal></entry>
+       <entry><type>setof regclass</type></entry>
+       <entry>get all tables in partition tree under given root table</entry>
+      </row>
+     </tbody>
+    </tgroup>
+   </table>
+
+   <para>
+    If the table passed to <function>pg_partition_root_parent</function> is not
+    a partition, the same table is returned as the result.  Result of
+    <function>pg_partition_tree_tables</function> also contains the table
+    that's passed to it as the first row.
+   </para>
+
   </sect2>
 
   <sect2 id="functions-admin-index">
diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 558022647c..5b3e8d52c5 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -23,13 +23,16 @@
 #include "catalog/partition.h"
 #include "catalog/pg_inherits.h"
 #include "catalog/pg_partitioned_table.h"
+#include "funcapi.h"
 #include "nodes/makefuncs.h"
 #include "optimizer/clauses.h"
 #include "optimizer/prep.h"
 #include "optimizer/var.h"
 #include "partitioning/partbounds.h"
 #include "rewrite/rewriteManip.h"
+#include "utils/builtins.h"
 #include "utils/fmgroids.h"
+#include "utils/lsyscache.h"
 #include "utils/partcache.h"
 #include "utils/rel.h"
 #include "utils/syscache.h"
@@ -44,10 +47,6 @@ static void get_partition_ancestors_worker(Relation inhRel, Oid relid,
  *		Obtain direct parent of given relation
  *
  * Returns inheritance parent of a partition by scanning pg_inherits
- *
- * Note: Because this function assumes that the relation whose OID is passed
- * as an argument will have precisely one parent, it should only be called
- * when it is known that the relation is a partition.
  */
 Oid
 get_partition_parent(Oid relid)
@@ -55,6 +54,9 @@ get_partition_parent(Oid relid)
 	Relation	catalogRelation;
 	Oid			result;
 
+	if (!get_rel_relispartition(relid))
+		return InvalidOid;
+
 	catalogRelation = heap_open(InheritsRelationId, AccessShareLock);
 
 	result = get_partition_parent_worker(catalogRelation, relid);
@@ -71,6 +73,10 @@ get_partition_parent(Oid relid)
  * get_partition_parent_worker
  *		Scan the pg_inherits relation to return the OID of the parent of the
  *		given relation
+ *
+ * Note: Because this function assumes that the relation whose OID is passed
+ * as an argument will have precisely one parent, it should only be called
+ * when it is known that the relation is a partition.
  */
 static Oid
 get_partition_parent_worker(Relation inhRel, Oid relid)
@@ -148,6 +154,27 @@ get_partition_ancestors_worker(Relation inhRel, Oid relid, List **ancestors)
 }
 
 /*
+ * get_partition_root_parent
+ *
+ * Returns root inheritance ancestor of a partition.
+ */
+Oid
+get_partition_root_parent(Oid relid)
+{
+	List	   *ancestors;
+	Oid			result;
+
+	if (!get_rel_relispartition(relid))
+		return InvalidOid;
+
+	ancestors = get_partition_ancestors(relid);
+	result = llast_oid(ancestors);
+	list_free(ancestors);
+
+	return result;
+}
+
+/*
  * map_partition_varattnos - maps varattno of any Vars in expr from the
  * attno's of 'from_rel' to the attno's of 'to_rel' partition, each of which
  * may be either a leaf partition or a partitioned table, but both of which
@@ -357,3 +384,97 @@ get_proposed_default_constraint(List *new_part_constraints)
 
 	return make_ands_implicit(defPartConstraint);
 }
+
+/*
+ * SQL wrapper around get_partition_root_parent().
+ */
+Datum
+pg_partition_root_parent(PG_FUNCTION_ARGS)
+{
+	Oid		reloid = PG_GETARG_OID(0);
+	Oid		rootoid;
+
+	rootoid = get_partition_root_parent(reloid);
+
+	/*
+	 * If the relation is actually a partition, 'rootoid' has been set to the
+	 * OID of the root table in the partition tree.
+	 */
+	if (OidIsValid(rootoid))
+		PG_RETURN_OID(rootoid);
+	/*
+	 * Otherwise, the table's not a partition.  That is, it's either the root
+	 * table in a partition tree or a standalone table that's not part of any
+	 * partition tree.  In any case, return the table OID itself as the
+	 * result.
+	 */
+	else
+		PG_RETURN_OID(reloid);
+}
+
+/*
+ * SQL wrapper around get_partition_parent().
+ */
+Datum
+pg_partition_parent(PG_FUNCTION_ARGS)
+{
+	Oid		reloid = PG_GETARG_OID(0);
+	Oid		parentoid;
+
+	parentoid = get_partition_parent(reloid);
+
+	/*
+	 * If the relation is actually a partition, 'parentoid' has been set to
+	 * the OID of the immediate parent.
+	 */
+	if (OidIsValid(parentoid))
+		PG_RETURN_OID(parentoid);
+	else
+	/* Not a partition, return NULL. */
+		PG_RETURN_NULL();
+}
+
+/*
+ * Returns OIDs of tables in a partition tree.
+ */
+Datum
+pg_partition_tree_tables(PG_FUNCTION_ARGS)
+{
+	FuncCallContext *funcctx;
+	Oid		reloid = PG_GETARG_OID(0);
+	List   *partoids;
+	ListCell  **lc;
+
+	/* stuff done only on the first call of the function */
+	if (SRF_IS_FIRSTCALL())
+	{
+		MemoryContext oldcontext;
+
+		/* create a function context for cross-call persistence */
+		funcctx = SRF_FIRSTCALL_INIT();
+
+		/* switch to memory context appropriate for multiple function calls */
+		oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+
+		partoids = find_all_inheritors(reloid, NoLock, NULL);
+		lc = (ListCell **) palloc(sizeof(ListCell *));
+		*lc = list_head(partoids);
+		funcctx->user_fctx = (void *) lc;
+
+		MemoryContextSwitchTo(oldcontext);
+	}
+
+	/* stuff done on every call of the function */
+	funcctx = SRF_PERCALL_SETUP();
+	lc = (ListCell **) funcctx->user_fctx;
+
+	while (*lc != NULL)
+	{
+		Oid		partoid = lfirst_oid(*lc);
+
+		*lc = lnext(*lc);
+		SRF_RETURN_NEXT(funcctx, ObjectIdGetDatum(partoid));
+	}
+
+	SRF_RETURN_DONE(funcctx);
+}
diff --git a/src/backend/utils/cache/lsyscache.c b/src/backend/utils/cache/lsyscache.c
index bba595ad1d..19262c6c4d 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -1821,6 +1821,28 @@ get_rel_relkind(Oid relid)
 }
 
 /*
+ * get_rel_relispartition
+ *
+ *		Returns the value of pg_class.relispartition for a given relation.
+ */
+char
+get_rel_relispartition(Oid relid)
+{
+	HeapTuple	tp;
+	Form_pg_class reltup;
+	bool	result;
+
+	tp = SearchSysCache1(RELOID, ObjectIdGetDatum(relid));
+	if (!HeapTupleIsValid(tp))
+		elog(ERROR, "cache lookup failed for relation %u", relid);
+	reltup = (Form_pg_class) GETSTRUCT(tp);
+	result = reltup->relispartition;
+	ReleaseSysCache(tp);
+
+	return result;
+}
+
+/*
  * get_rel_tablespace
  *
  *		Returns the pg_tablespace OID associated with a given relation.
diff --git a/src/include/catalog/partition.h b/src/include/catalog/partition.h
index 1f49e5d3a9..a2b11f40bc 100644
--- a/src/include/catalog/partition.h
+++ b/src/include/catalog/partition.h
@@ -32,6 +32,7 @@ typedef struct PartitionDescData
 
 extern Oid	get_partition_parent(Oid relid);
 extern List *get_partition_ancestors(Oid relid);
+extern Oid	get_partition_root_parent(Oid relid);
 extern List *map_partition_varattnos(List *expr, int fromrel_varno,
 						Relation to_rel, Relation from_rel,
 						bool *found_whole_row);
diff --git a/src/include/catalog/pg_proc.dat b/src/include/catalog/pg_proc.dat
index 40d54ed030..b4725b8634 100644
--- a/src/include/catalog/pg_proc.dat
+++ b/src/include/catalog/pg_proc.dat
@@ -10206,4 +10206,22 @@
   proisstrict => 'f', prorettype => 'bool', proargtypes => 'oid int4 int4 any',
   proargmodes => '{i,i,i,v}', prosrc => 'satisfies_hash_partition' },
 
+# function to get the root partition parent
+{ oid => '3423', descr => 'oid of the partition root parent',
+  proname => 'pg_partition_root_parent', prorettype => 'regclass',
+  proargtypes => 'regclass', prosrc => 'pg_partition_root_parent' },
+
+# function to get the partition parent
+{ oid => '3424', descr => 'oid of the partition immediate parent',
+  proname => 'pg_partition_parent', prorettype => 'regclass',
+  proargtypes => 'regclass', prosrc => 'pg_partition_parent' },
+
+# function to get OIDs of all tables in a given partition tree
+{ oid => '3425', descr => 'get OIDs of tables in a partition tree',
+  proname => 'pg_partition_tree_tables', prorettype => '2205',
+  prorows => '100', proretset => 't', proargtypes => 'regclass',
+  proallargtypes => '{regclass,regclass}',
+  proargmodes => '{i,o}',
+  proargnames => '{relid,relid}', prosrc => 'pg_partition_tree_tables' }
+
 ]
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index e55ea4035b..d396d17ff1 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -126,6 +126,7 @@ extern char *get_rel_name(Oid relid);
 extern Oid	get_rel_namespace(Oid relid);
 extern Oid	get_rel_type_id(Oid relid);
 extern char get_rel_relkind(Oid relid);
+extern char get_rel_relispartition(Oid relid);
 extern Oid	get_rel_tablespace(Oid relid);
 extern char get_rel_persistence(Oid relid);
 extern Oid	get_transform_fromsql(Oid typid, Oid langid, List *trftypes);
diff --git a/src/test/regress/expected/create_table.out b/src/test/regress/expected/create_table.out
index 672719e5d5..f702518d4e 100644
--- a/src/test/regress/expected/create_table.out
+++ b/src/test/regress/expected/create_table.out
@@ -909,3 +909,63 @@ ERROR:  cannot create a temporary relation as partition of permanent relation "p
 create temp table temp_part partition of temp_parted default; -- ok
 drop table perm_parted cascade;
 drop table temp_parted cascade;
+-- tests to show partition tree inspection functions
+create table ptif_test (a int, b int) partition by range (a);
+create table ptif_test0 partition of ptif_test for values from (minvalue) to (0) partition by list (b);
+create table ptif_test01 partition of ptif_test0 for values in (1);
+create table ptif_test1 partition of ptif_test for values from (0) to (maxvalue) partition by list (b);
+create table ptif_test11 partition of ptif_test1 for values in (1);
+insert into ptif_test select i, 1 from generate_series(-5, 5) i;
+select pg_partition_parent('ptif_test0') as parent;
+  parent   
+-----------
+ ptif_test
+(1 row)
+
+select pg_partition_parent('ptif_test01') as parent;
+   parent   
+------------
+ ptif_test0
+(1 row)
+
+select pg_partition_root_parent('ptif_test01') as root_parent;
+ root_parent 
+-------------
+ ptif_test
+(1 row)
+
+select	p as relname,
+		pg_partition_parent(p) as parent,
+		pg_partition_root_parent(p) as root_parent
+from	pg_partition_tree_tables('ptif_test') p;
+   relname   |   parent   | root_parent 
+-------------+------------+-------------
+ ptif_test   |            | ptif_test
+ ptif_test0  | ptif_test  | ptif_test
+ ptif_test1  | ptif_test  | ptif_test
+ ptif_test01 | ptif_test0 | ptif_test
+ ptif_test11 | ptif_test1 | ptif_test
+(5 rows)
+
+select	p as relname,
+		pg_partition_parent(p) as parent,
+		pg_partition_root_parent(p) as root_parent,
+		pg_relation_size(p) as size
+from	pg_partition_tree_tables('ptif_test') p;
+   relname   |   parent   | root_parent | size 
+-------------+------------+-------------+------
+ ptif_test   |            | ptif_test   |    0
+ ptif_test0  | ptif_test  | ptif_test   |    0
+ ptif_test1  | ptif_test  | ptif_test   |    0
+ ptif_test01 | ptif_test0 | ptif_test   | 8192
+ ptif_test11 | ptif_test1 | ptif_test   | 8192
+(5 rows)
+
+select	sum(pg_relation_size(p)) as total_size
+from	pg_partition_tree_tables('ptif_test') p;
+ total_size 
+------------
+      16384
+(1 row)
+
+drop table ptif_test;
diff --git a/src/test/regress/sql/create_table.sql b/src/test/regress/sql/create_table.sql
index 78944950fe..d0a083dd69 100644
--- a/src/test/regress/sql/create_table.sql
+++ b/src/test/regress/sql/create_table.sql
@@ -735,3 +735,27 @@ create temp table temp_part partition of perm_parted default; -- error
 create temp table temp_part partition of temp_parted default; -- ok
 drop table perm_parted cascade;
 drop table temp_parted cascade;
+
+-- tests to show partition tree inspection functions
+create table ptif_test (a int, b int) partition by range (a);
+create table ptif_test0 partition of ptif_test for values from (minvalue) to (0) partition by list (b);
+create table ptif_test01 partition of ptif_test0 for values in (1);
+create table ptif_test1 partition of ptif_test for values from (0) to (maxvalue) partition by list (b);
+create table ptif_test11 partition of ptif_test1 for values in (1);
+insert into ptif_test select i, 1 from generate_series(-5, 5) i;
+
+select pg_partition_parent('ptif_test0') as parent;
+select pg_partition_parent('ptif_test01') as parent;
+select pg_partition_root_parent('ptif_test01') as root_parent;
+select	p as relname,
+		pg_partition_parent(p) as parent,
+		pg_partition_root_parent(p) as root_parent
+from	pg_partition_tree_tables('ptif_test') p;
+select	p as relname,
+		pg_partition_parent(p) as parent,
+		pg_partition_root_parent(p) as root_parent,
+		pg_relation_size(p) as size
+from	pg_partition_tree_tables('ptif_test') p;
+select	sum(pg_relation_size(p)) as total_size
+from	pg_partition_tree_tables('ptif_test') p;
+drop table ptif_test;
-- 
2.11.0

