From 76be9aa755c0733852074c84e3e09102d1768427 Mon Sep 17 00:00:00 2001
From: Mike Palmiotto <mike.palmiotto@crunchydata.com>
Date: Mon, 8 Jul 2019 12:46:21 +0000
Subject: [PATCH] Flexible partition pruning hook

This hook allows partition pruning to be performed for entire child relations
which are deemed inaccessible by RLS policy prior to those relations being
opened on disk.
---
 src/backend/catalog/pg_inherits.c     |  6 ++++++
 src/backend/optimizer/path/allpaths.c |  7 +++++++
 src/include/partitioning/partprune.h  | 18 ++++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/src/backend/catalog/pg_inherits.c b/src/backend/catalog/pg_inherits.c
index 00f7957..c1fa8b0 100644
--- a/src/backend/catalog/pg_inherits.c
+++ b/src/backend/catalog/pg_inherits.c
@@ -24,7 +24,9 @@
 #include "access/table.h"
 #include "catalog/indexing.h"
 #include "catalog/pg_inherits.h"
+#include "optimizer/cost.h"
 #include "parser/parse_type.h"
+#include "partitioning/partprune.h"
 #include "storage/lmgr.h"
 #include "utils/builtins.h"
 #include "utils/fmgroids.h"
@@ -92,6 +94,10 @@ find_inheritance_children(Oid parentrelId, LOCKMODE lockmode)
 	while ((inheritsTuple = systable_getnext(scan)) != NULL)
 	{
 		inhrelid = ((Form_pg_inherits) GETSTRUCT(inheritsTuple))->inhrelid;
+
+		if (!InvokePartitionChildAccessHook(inhrelid))
+			continue;
+
 		if (numoids >= maxoids)
 		{
 			maxoids *= 2;
diff --git a/src/backend/optimizer/path/allpaths.c b/src/backend/optimizer/path/allpaths.c
index b772348..bb81b1b 100644
--- a/src/backend/optimizer/path/allpaths.c
+++ b/src/backend/optimizer/path/allpaths.c
@@ -1038,6 +1038,13 @@ set_append_rel_size(PlannerInfo *root, RelOptInfo *rel,
 			continue;
 		}
 
+		if (!InvokePartitionChildAccessHook(childRTE->relid))
+		{
+			/* Implement custom partition pruning filter*/
+			set_dummy_rel_pathlist(childrel);
+			continue;
+		}
+
 		/*
 		 * Constraint exclusion failed, so copy the parent's join quals and
 		 * targetlist to the child, with appropriate variable substitutions.
diff --git a/src/include/partitioning/partprune.h b/src/include/partitioning/partprune.h
index b3e9268..3138ff9 100644
--- a/src/include/partitioning/partprune.h
+++ b/src/include/partitioning/partprune.h
@@ -15,6 +15,7 @@
 #define PARTPRUNE_H
 
 #include "nodes/execnodes.h"
+#include "optimizer/cost.h"
 #include "partitioning/partdefs.h"
 
 struct PlannerInfo;				/* avoid including pathnodes.h here */
@@ -68,6 +69,23 @@ typedef struct PartitionPruneContext
 #define PruneCxtStateIdx(partnatts, step_id, keyno) \
 	((partnatts) * (step_id) + (keyno))
 
+/* Custom partition child access hook. Provides further partition pruning given
+ * child OID.
+ */
+typedef bool (*partitionChildAccess_hook_type) (Oid childOID);
+PGDLLIMPORT partitionChildAccess_hook_type partitionChildAccess_hook;
+
+/* Macro to use partitionChildAccess_hook. Handles NULL-checking. */
+static inline bool InvokePartitionChildAccessHook (Oid childOID)
+{
+	if (partitionChildAccess_hook && enable_partition_pruning && childOID)
+	{
+		return (*partitionChildAccess_hook) (childOID);
+	}
+
+	return true;
+}
+
 extern PartitionPruneInfo *make_partition_pruneinfo(struct PlannerInfo *root,
 													struct RelOptInfo *parentrel,
 													List *subpaths,
-- 
1.8.3.1