Add pg_partition_tree to display information about partitions

This new function is useful to display a full tree of partitions with a
partitioned table given in output, and avoids the need of any complex
WITH RECURSIVE query when looking at partition trees which are
deep multiple levels.

It returns a set of records, one for each partition, containing the
partition's name, its immediate parent's name, a boolean value telling
if the relation is a leaf in the tree and an integer telling its level
in the partition tree with given table considered as root, beginning at
zero for the root, and incrementing by one each time the scan goes one
level down.

Author: Amit Langote
Reviewed-by: Jesper Pedersen, Michael Paquier, Robert Haas
Discussion: https://postgr.es/m/8d00e51a-9a51-ad02-d53e-ba6bf50b2e52@lab.ntt.co.jp

1 files changed, 154 insertions, 0 deletions

diff --git a/src/backend/utils/adt/partitionfuncs.c b/src/backend/utils/adt/partitionfuncs.c
new file mode 100644
index 00000000000..8f9218ad0aa
--- /dev/null
+++ b/src/backend/utils/adt/partitionfuncs.c
@@ -0,0 +1,154 @@
+/*-------------------------------------------------------------------------
+ *
+ * partitionfuncs.c
+ *	  Functions for accessing partition-related metadata
+ *
+ * Portions Copyright (c) 1996-2018, PostgreSQL Global Development Group
+ * Portions Copyright (c) 1994, Regents of the University of California
+ *
+ *
+ * IDENTIFICATION
+ *	  src/backend/utils/adt/partitionfuncs.c
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#include "postgres.h"
+
+#include "access/htup_details.h"
+#include "catalog/partition.h"
+#include "catalog/pg_class.h"
+#include "catalog/pg_inherits.h"
+#include "catalog/pg_type.h"
+#include "funcapi.h"
+#include "utils/fmgrprotos.h"
+#include "utils/lsyscache.h"
+
+
+/*
+ * pg_partition_tree
+ *
+ * Produce a view with one row per member of a partition tree, beginning
+ * from the top-most parent given by the caller.  This gives information
+ * about each partition, its immediate partitioned parent, if it is
+ * a leaf partition and its level in the hierarchy.
+ */
+Datum
+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;
+
+	/* 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)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("\"%s\" is not a table, a foreign table, or an index",
+						get_rel_name(rootrelid))));
+
+	/* stuff done only on the first call of the function */
+	if (SRF_IS_FIRSTCALL())
+	{
+		MemoryContext oldcxt;
+		TupleDesc	tupdesc;
+		List	   *partitions;
+
+		/* create a function context for cross-call persistence */
+		funcctx = SRF_FIRSTCALL_INIT();
+
+		/* switch to memory context appropriate for multiple function calls */
+		oldcxt = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx);
+
+		/*
+		 * Find all members of inheritance set.  We only need AccessShareLock
+		 * on the children for the partition information lookup.
+		 */
+		partitions = find_all_inheritors(rootrelid, AccessShareLock, NULL);
+
+		tupdesc = CreateTemplateTupleDesc(PG_PARTITION_TREE_COLS, false);
+		TupleDescInitEntry(tupdesc, (AttrNumber) 1, "relid",
+						   REGCLASSOID, -1, 0);
+		TupleDescInitEntry(tupdesc, (AttrNumber) 2, "parentid",
+						   REGCLASSOID, -1, 0);
+		TupleDescInitEntry(tupdesc, (AttrNumber) 3, "isleaf",
+						   BOOLOID, -1, 0);
+		TupleDescInitEntry(tupdesc, (AttrNumber) 4, "level",
+						   INT4OID, -1, 0);
+
+		funcctx->tuple_desc = BlessTupleDesc(tupdesc);
+
+		/* allocate memory for user context */
+		next = (ListCell **) palloc(sizeof(ListCell *));
+		*next = list_head(partitions);
+		funcctx->user_fctx = (void *) next;
+
+		MemoryContextSwitchTo(oldcxt);
+	}
+
+	/* stuff done on every call of the function */
+	funcctx = SRF_PERCALL_SETUP();
+	next = (ListCell **) funcctx->user_fctx;
+
+	if (*next != NULL)
+	{
+		Datum		result;
+		Datum		values[PG_PARTITION_TREE_COLS];
+		bool		nulls[PG_PARTITION_TREE_COLS];
+		HeapTuple	tuple;
+		Oid			parentid = InvalidOid;
+		Oid			relid = lfirst_oid(*next);
+		char		relkind = get_rel_relkind(relid);
+		int			level = 0;
+		List	   *ancestors = get_partition_ancestors(lfirst_oid(*next));
+		ListCell   *lc;
+
+		/*
+		 * Form tuple with appropriate data.
+		 */
+		MemSet(nulls, 0, sizeof(nulls));
+		MemSet(values, 0, sizeof(values));
+
+		/* relid */
+		values[0] = ObjectIdGetDatum(relid);
+
+		/* parentid */
+		if (ancestors != NIL)
+			parentid = linitial_oid(ancestors);
+		if (OidIsValid(parentid))
+			values[1] = ObjectIdGetDatum(parentid);
+		else
+			nulls[1] = true;
+
+		/* isleaf */
+		values[2] = BoolGetDatum(relkind != RELKIND_PARTITIONED_TABLE &&
+								 relkind != RELKIND_PARTITIONED_INDEX);
+
+		/* level */
+		if (relid != rootrelid)
+		{
+			foreach(lc, ancestors)
+			{
+				level++;
+				if (lfirst_oid(lc) == rootrelid)
+					break;
+			}
+		}
+		values[3] = Int32GetDatum(level);
+
+		*next = lnext(*next);
+
+		tuple = heap_form_tuple(funcctx->tuple_desc, values, nulls);
+		result = HeapTupleGetDatum(tuple);
+		SRF_RETURN_NEXT(funcctx, result);
+	}
+
+	/* done when there are no more elements left */
+	SRF_RETURN_DONE(funcctx);
+}