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

---
 src/backend/catalog/partition.c     | 117 +++++++++++++++++++++++++++++++++++-
 src/backend/utils/cache/lsyscache.c |  22 +++++++
 src/include/catalog/partition.h     |   1 +
 src/include/catalog/pg_proc.h       |  12 ++++
 src/include/utils/lsyscache.h       |   1 +
 5 files changed, 152 insertions(+), 1 deletion(-)

diff --git a/src/backend/catalog/partition.c b/src/backend/catalog/partition.c
index 8adc4ee977..ac92bbfa71 100644
--- a/src/backend/catalog/partition.c
+++ b/src/backend/catalog/partition.c
@@ -32,6 +32,7 @@
 #include "catalog/pg_type.h"
 #include "commands/tablecmds.h"
 #include "executor/executor.h"
+#include "funcapi.h"
 #include "miscadmin.h"
 #include "nodes/makefuncs.h"
 #include "nodes/nodeFuncs.h"
@@ -181,6 +182,7 @@ static int partition_bound_bsearch(PartitionKey key,
 static int	get_partition_bound_num_indexes(PartitionBoundInfo b);
 static int	get_greatest_modulus(PartitionBoundInfo b);
 static uint64 compute_hash_value(PartitionKey key, Datum *values, bool *isnull);
+static Oid get_partition_parent_internal(Oid relid, bool recurse_to_root);
 
 /* SQL-callable function for use in hash partition CHECK constraints */
 PG_FUNCTION_INFO_V1(satisfies_hash_partition);
@@ -1362,7 +1364,7 @@ check_default_allows_bound(Relation parent, Relation default_rel,
 /*
  * get_partition_parent
  *
- * Returns inheritance parent of a partition by scanning pg_inherits
+ * Returns inheritance parent of a partition.
  *
  * 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
@@ -1371,6 +1373,37 @@ check_default_allows_bound(Relation parent, Relation default_rel,
 Oid
 get_partition_parent(Oid relid)
 {
+	if (!get_rel_relispartition(relid))
+		return InvalidOid;
+
+	return get_partition_parent_internal(relid, false);
+}
+
+/*
+ * get_partition_root_parent
+ *
+ * Returns root inheritance ancestor of a partition.
+ */
+Oid
+get_partition_root_parent(Oid relid)
+{
+	if (!get_rel_relispartition(relid))
+		return InvalidOid;
+
+	return get_partition_parent_internal(relid, true);
+}
+
+/*
+ * get_partition_parent_internal
+ *
+ * Returns inheritance parent of a partition by scanning pg_inherits.
+ * If recurse_to_root, it will check if the parent itself is a partition and
+ * if so, it will recurse to find its parent and so on until root parent is
+ * found.
+ */
+static Oid
+get_partition_parent_internal(Oid relid, bool recurse_to_root)
+{
 	Form_pg_inherits form;
 	Relation	catalogRelation;
 	SysScanDesc scan;
@@ -1402,6 +1435,9 @@ get_partition_parent(Oid relid)
 	systable_endscan(scan);
 	heap_close(catalogRelation, AccessShareLock);
 
+	if (recurse_to_root && get_rel_relispartition(result))
+		result = get_partition_parent_internal(result, recurse_to_root);
+
 	return result;
 }
 
@@ -3396,3 +3432,82 @@ satisfies_hash_partition(PG_FUNCTION_ARGS)
 
 	PG_RETURN_BOOL(rowHash % modulus == remainder);
 }
+
+/*
+ * SQL wrapper around get_partition_root_parent() in
+ * src/backend/catalog/partition.c.
+ */
+Datum
+pg_partition_root(PG_FUNCTION_ARGS)
+{
+	Oid		reloid = PG_GETARG_OID(0);
+	Oid		rootoid;
+
+	rootoid = get_partition_root_parent(reloid);
+	if (OidIsValid(rootoid))
+		PG_RETURN_OID(rootoid);
+	else
+		PG_RETURN_OID(reloid);
+}
+
+/*
+ * SQL wrapper around get_partition_parent() in
+ * src/backend/catalog/partition.c.
+ */
+Datum
+pg_partition_parent(PG_FUNCTION_ARGS)
+{
+	Oid		reloid = PG_GETARG_OID(0);
+	Oid		parentoid;
+
+	parentoid = get_partition_parent(reloid);
+	if (OidIsValid(parentoid))
+		PG_RETURN_OID(parentoid);
+	else
+		PG_RETURN_NULL();
+}
+
+/*
+ * Returns Oids of tables in a publication.
+ */
+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 e8aa179347..92353a6004 100644
--- a/src/backend/utils/cache/lsyscache.c
+++ b/src/backend/utils/cache/lsyscache.c
@@ -1836,6 +1836,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 2faf0ca26e..287642b01b 100644
--- a/src/include/catalog/partition.h
+++ b/src/include/catalog/partition.h
@@ -52,6 +52,7 @@ extern PartitionBoundInfo partition_bounds_copy(PartitionBoundInfo src,
 extern void check_new_partition_bound(char *relname, Relation parent,
 						  PartitionBoundSpec *spec);
 extern Oid	get_partition_parent(Oid relid);
+extern Oid	get_partition_root_parent(Oid relid);
 extern List *get_qual_from_partbound(Relation rel, Relation parent,
 						PartitionBoundSpec *spec);
 extern List *map_partition_varattnos(List *expr, int fromrel_varno,
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index f01648c961..64942b310c 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -5533,6 +5533,18 @@ DESCR("list of files in the WAL directory");
 DATA(insert OID = 5028 ( satisfies_hash_partition PGNSP PGUID 12 1 0 2276 0 f f f f f f i s 4 0 16 "26 23 23 2276" _null_ "{i,i,i,v}" _null_ _null_ _null_ satisfies_hash_partition _null_ _null_ _null_ ));
 DESCR("hash partition CHECK constraint");
 
+/* function to get the root partition parent */
+DATA(insert OID = 3281 (  pg_partition_root PGNSP PGUID 12 10 0 0 0 f f f f t f s s 1 0 2205 "2205" _null_ _null_ _null_ _null_ _null_ pg_partition_root _null_ _null_ _null_ ));
+DESCR("oid of the partition root parent");
+
+/* function to get the partition parent */
+DATA(insert OID = 3556 (  pg_partition_parent PGNSP PGUID 12 10 0 0 0 f f f f t f s s 1 0 2205 "2205" _null_ _null_ _null_ _null_ _null_ pg_partition_parent _null_ _null_ _null_ ));
+DESCR("oid of the partition immediate parent");
+
+/* function to get OIDs of all tables in a given partition tree */
+DATA(insert OID = 3696 ( pg_partition_tree_tables	PGNSP PGUID 12 1 1000 0 0 f f f f t t s s 1 0 2205 "2205" "{2205,2205}" "{i,o}" "{relid,relid}" _null_ _null_ pg_partition_tree_tables _null_ _null_ _null_ ));
+DESCR("get OIDs of tables in a partition tree");
+
 /*
  * Symbolic values for provolatile column: these indicate whether the result
  * of a function is dependent *only* on the values of its explicit arguments,
diff --git a/src/include/utils/lsyscache.h b/src/include/utils/lsyscache.h
index 9731e6f7ae..1000d9fd13 100644
--- a/src/include/utils/lsyscache.h
+++ b/src/include/utils/lsyscache.h
@@ -127,6 +127,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);
-- 
2.11.0

