From aab5a93200a8d24a10b9c351ac3199cd33e43da3 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Wed, 27 Jul 2016 16:00:09 +0900
Subject: [PATCH 6/9] Teach a few places to use partition check quals.

For example, if a row is inserted directly into a partition we should make
sure that it does not violate its bounds.  So teach copy.c and execMain.c
to apply "partition check constraint".

Also, for constraint exclusion to work with partitioned tables, teach the
optimizer to include check constraint expressions derived from partition bound
bound info in the list of predicates it uses to perform the task.
---
 src/backend/commands/copy.c            |    2 +-
 src/backend/executor/execMain.c        |   76 ++++++++++++++++++++++++++++++-
 src/backend/executor/nodeModifyTable.c |    4 +-
 src/backend/optimizer/util/plancat.c   |   20 ++++++++
 src/include/nodes/execnodes.h          |    4 ++
 src/test/regress/expected/insert.out   |   76 ++++++++++++++++++++++++++++++++
 src/test/regress/expected/update.out   |   27 +++++++++++
 src/test/regress/sql/insert.sql        |   56 +++++++++++++++++++++++
 src/test/regress/sql/update.sql        |   21 +++++++++
 9 files changed, 280 insertions(+), 6 deletions(-)

diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c
index 5120a4d..3f5c4e6 100644
--- a/src/backend/commands/copy.c
+++ b/src/backend/commands/copy.c
@@ -2464,7 +2464,7 @@ CopyFrom(CopyState cstate)
 		if (!skip_tuple)
 		{
 			/* Check the constraints of the tuple */
-			if (cstate->rel->rd_att->constr)
+			if (cstate->rel->rd_att->constr || resultRelInfo->ri_PartitionCheck)
 				ExecConstraints(resultRelInfo, slot, estate);
 
 			if (useHeapMultiInsert)
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index 9773272..3d82658 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -42,6 +42,7 @@
 #include "access/transam.h"
 #include "access/xact.h"
 #include "catalog/namespace.h"
+#include "catalog/partition.h"
 #include "commands/matview.h"
 #include "commands/trigger.h"
 #include "executor/execdebug.h"
@@ -1251,6 +1252,8 @@ InitResultRelInfo(ResultRelInfo *resultRelInfo,
 	resultRelInfo->ri_ConstraintExprs = NULL;
 	resultRelInfo->ri_junkFilter = NULL;
 	resultRelInfo->ri_projectReturning = NULL;
+	resultRelInfo->ri_PartitionCheck =
+						RelationGetPartitionCheckQual(resultRelationDesc);
 }
 
 /*
@@ -1692,6 +1695,50 @@ ExecRelCheck(ResultRelInfo *resultRelInfo,
 	return NULL;
 }
 
+/*
+ * ExecPartitionCheck --- check that tuple meets the partition boundary
+ * specification.
+ *
+ * Note: This is called, *iff* resultRelInfo is the main target table.
+ */
+static bool
+ExecPartitionCheck(ResultRelInfo *resultRelInfo, TupleTableSlot *slot,
+				   EState *estate)
+{
+	ExprContext *econtext;
+
+	/*
+	 * If first time through, build expression state tree for the partition
+	 * check expression.  Keep it in the per-query memory context so they'll
+	 * survive throughout the query.
+	 */
+	if (resultRelInfo->ri_PartitionCheckExpr == NULL)
+	{
+		List *qual = resultRelInfo->ri_PartitionCheck;
+
+		resultRelInfo->ri_PartitionCheckExpr = (List *)
+									ExecPrepareExpr((Expr *) qual, estate);
+	}
+
+	/*
+	 * We will use the EState's per-tuple context for evaluating constraint
+	 * expressions (creating it if it's not already there).
+	 */
+	econtext = GetPerTupleExprContext(estate);
+
+	/* Arrange for econtext's scan tuple to be the tuple under test */
+	econtext->ecxt_scantuple = slot;
+
+	/*
+	 * NOTE: SQL specifies that a NULL result from a constraint expression
+	 * is not to be treated as a failure.  Therefore, tell ExecQual to
+	 * return TRUE for NULL.
+	 *
+	 * XXX - although, it's unlikely that NULL would result.
+	 */
+	return ExecQual(resultRelInfo->ri_PartitionCheckExpr, econtext, true);
+}
+
 void
 ExecConstraints(ResultRelInfo *resultRelInfo,
 				TupleTableSlot *slot, EState *estate)
@@ -1703,9 +1750,9 @@ ExecConstraints(ResultRelInfo *resultRelInfo,
 	Bitmapset  *insertedCols;
 	Bitmapset  *updatedCols;
 
-	Assert(constr);
+	Assert(constr || resultRelInfo->ri_PartitionCheck);
 
-	if (constr->has_not_null)
+	if (constr && constr->has_not_null)
 	{
 		int			natts = tupdesc->natts;
 		int			attrChk;
@@ -1736,7 +1783,7 @@ ExecConstraints(ResultRelInfo *resultRelInfo,
 		}
 	}
 
-	if (constr->num_check > 0)
+	if (constr && constr->num_check > 0)
 	{
 		const char *failed;
 
@@ -1760,6 +1807,29 @@ ExecConstraints(ResultRelInfo *resultRelInfo,
 					 errtableconstraint(rel, failed)));
 		}
 	}
+
+	if (resultRelInfo->ri_PartitionCheck)
+	{
+		if (!ExecPartitionCheck(resultRelInfo, slot, estate))
+		{
+			char	   *val_desc;
+
+			insertedCols = GetInsertedColumns(resultRelInfo, estate);
+			updatedCols = GetUpdatedColumns(resultRelInfo, estate);
+			modifiedCols = bms_union(insertedCols, updatedCols);
+			val_desc = ExecBuildSlotValueDescription(RelationGetRelid(rel),
+													 slot,
+													 tupdesc,
+													 modifiedCols,
+													 64);
+			ereport(ERROR,
+					(errcode(ERRCODE_CHECK_VIOLATION),
+					 errmsg("new row violates the partition boundary"
+							" specification of \"%s\"",
+							RelationGetRelationName(rel)),
+			  val_desc ? errdetail("Failing row contains %s.", val_desc) : 0));
+		}
+	}
 }
 
 /*
diff --git a/src/backend/executor/nodeModifyTable.c b/src/backend/executor/nodeModifyTable.c
index 5790edc..5b0e8cf 100644
--- a/src/backend/executor/nodeModifyTable.c
+++ b/src/backend/executor/nodeModifyTable.c
@@ -354,7 +354,7 @@ ExecInsert(ModifyTableState *mtstate,
 		/*
 		 * Check the constraints of the tuple
 		 */
-		if (resultRelationDesc->rd_att->constr)
+		if (resultRelationDesc->rd_att->constr || resultRelInfo->ri_PartitionCheck)
 			ExecConstraints(resultRelInfo, slot, estate);
 
 		if (onconflict != ONCONFLICT_NONE && resultRelInfo->ri_NumIndices > 0)
@@ -907,7 +907,7 @@ lreplace:;
 		/*
 		 * Check the constraints of the tuple
 		 */
-		if (resultRelationDesc->rd_att->constr)
+		if (resultRelationDesc->rd_att->constr || resultRelInfo->ri_PartitionCheck)
 			ExecConstraints(resultRelInfo, slot, estate);
 
 		/*
diff --git a/src/backend/optimizer/util/plancat.c b/src/backend/optimizer/util/plancat.c
index 8ecc116..70e7d4f 100644
--- a/src/backend/optimizer/util/plancat.c
+++ b/src/backend/optimizer/util/plancat.c
@@ -27,6 +27,7 @@
 #include "catalog/catalog.h"
 #include "catalog/dependency.h"
 #include "catalog/heap.h"
+#include "catalog/partition.h"
 #include "catalog/pg_am.h"
 #include "foreign/fdwapi.h"
 #include "miscadmin.h"
@@ -1127,6 +1128,7 @@ get_relation_constraints(PlannerInfo *root,
 	Index		varno = rel->relid;
 	Relation	relation;
 	TupleConstr *constr;
+	List		*pcqual;
 
 	/*
 	 * We assume the relation has already been safely locked.
@@ -1212,6 +1214,24 @@ get_relation_constraints(PlannerInfo *root,
 		}
 	}
 
+	/* Append partition predicates, if any */
+	pcqual = RelationGetPartitionCheckQual(relation);
+	if (pcqual)
+	{
+		/*
+		 * Run each expression through const-simplification and
+		 * canonicalization similar to check constraints.
+		 */
+		pcqual = (List *) eval_const_expressions(root, (Node *) pcqual);
+		pcqual = (List *) canonicalize_qual((Expr *) pcqual);
+
+		/* Fix Vars to have the desired varno */
+		if (varno != 1)
+			ChangeVarNodes((Node *) pcqual, 1, varno, 0);
+
+		result = list_concat(result, pcqual);
+	}
+
 	heap_close(relation, NoLock);
 
 	return result;
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index e7fd7bd..31ca9ed 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -319,6 +319,8 @@ typedef struct JunkFilter
  *		projectReturning		for computing a RETURNING list
  *		onConflictSetProj		for computing ON CONFLICT DO UPDATE SET
  *		onConflictSetWhere		list of ON CONFLICT DO UPDATE exprs (qual)
+ *		PartitionCheck			partition check expression
+ *		PartitionCheckExpr		partition check expression state
  * ----------------
  */
 typedef struct ResultRelInfo
@@ -343,6 +345,8 @@ typedef struct ResultRelInfo
 	ProjectionInfo *ri_projectReturning;
 	ProjectionInfo *ri_onConflictSetProj;
 	List	   *ri_onConflictSetWhere;
+	List	   *ri_PartitionCheck;
+	List	   *ri_PartitionCheckExpr;
 } ResultRelInfo;
 
 /* ----------------
diff --git a/src/test/regress/expected/insert.out b/src/test/regress/expected/insert.out
index 70107b5..89d5760 100644
--- a/src/test/regress/expected/insert.out
+++ b/src/test/regress/expected/insert.out
@@ -160,3 +160,79 @@ Rules:
 drop table inserttest2;
 drop table inserttest;
 drop type insert_test_type;
+-- direct partition inserts should check partition bound constraint
+create table range_parted (
+	a text,
+	b int
+) partition by range (a, b);
+create table part_a_1_a_10 partition of range_parted for values start ('a', 1) end ('a', 10);
+create table part_a_10_a_20 partition of range_parted for values start ('a', 10) end ('a', 20);
+create table part_b_1_b_10 partition of range_parted for values start ('b', 1) end ('b', 10);
+create table part_b_10_b_20 partition of range_parted for values start ('b', 10) end ('b', 20);
+-- fail
+insert into part_a_1_a_10 values ('a', 11);
+ERROR:  new row violates the partition boundary specification of "part_a_1_a_10"
+DETAIL:  Failing row contains (a, 11).
+insert into part_a_1_a_10 values ('b', 1);
+ERROR:  new row violates the partition boundary specification of "part_a_1_a_10"
+DETAIL:  Failing row contains (b, 1).
+-- ok
+insert into part_a_1_a_10 values ('a', 1);
+-- fail
+insert into part_b_10_b_20 values ('b', 21);
+ERROR:  new row violates the partition boundary specification of "part_b_10_b_20"
+DETAIL:  Failing row contains (b, 21).
+insert into part_b_10_b_20 values ('a', 10);
+ERROR:  new row violates the partition boundary specification of "part_b_10_b_20"
+DETAIL:  Failing row contains (a, 10).
+-- ok
+insert into part_b_10_b_20 values ('b', 10);
+-- fail (a is null but a range partition key column should not be null)
+insert into part_b_10_b_20(b) values (10);
+ERROR:  new row violates the partition boundary specification of "part_b_10_b_20"
+DETAIL:  Failing row contains (null, 10).
+create table list_parted (
+	a text,
+	b int
+) partition by list (upper(a));
+create table part_AA_BB partition of list_parted FOR VALUES IN ('AA', 'BB');
+create table part_CC_DD partition of list_parted FOR VALUES IN ('CC', 'DD');
+-- fail
+insert into part_AA_BB values ('cc', 1);
+ERROR:  new row violates the partition boundary specification of "part_aa_bb"
+DETAIL:  Failing row contains (cc, 1).
+insert into part_AA_BB values ('AAa', 1);
+ERROR:  new row violates the partition boundary specification of "part_aa_bb"
+DETAIL:  Failing row contains (AAa, 1).
+-- ok
+insert into part_CC_DD values ('cC', 1);
+-- XXX - fail (a is null but part_AA_BB does not allow nulls in its list of values)
+-- insert into part_AA_BB (b) values (1);
+-- check in case of multi-level partitioned table
+create table part_EE_FF partition of list_parted for values in ('EE', 'FF') partition by range (b);
+create table part_EE_FF_1_10 partition of part_EE_FF for values start (1) end (10);
+create table part_EE_FF_10_20 partition of part_EE_FF for values start (10) end (20);
+-- fail (both its own and all ancestors' partition bound spec applies)
+insert into part_EE_FF_1_10 values ('EE', 11);
+ERROR:  new row violates the partition boundary specification of "part_ee_ff_1_10"
+DETAIL:  Failing row contains (EE, 11).
+insert into part_EE_FF_1_10 values ('cc', 1);
+ERROR:  new row violates the partition boundary specification of "part_ee_ff_1_10"
+DETAIL:  Failing row contains (cc, 1).
+-- ok
+insert into part_EE_FF_1_10 values ('ff', 1);
+insert into part_EE_FF_10_20 values ('ff', 11);
+-- cleanup
+drop table range_parted cascade;
+NOTICE:  drop cascades to 4 other objects
+DETAIL:  drop cascades to table part_a_1_a_10
+drop cascades to table part_a_10_a_20
+drop cascades to table part_b_1_b_10
+drop cascades to table part_b_10_b_20
+drop table list_parted cascade;
+NOTICE:  drop cascades to 5 other objects
+DETAIL:  drop cascades to table part_aa_bb
+drop cascades to table part_cc_dd
+drop cascades to table part_ee_ff
+drop cascades to table part_ee_ff_1_10
+drop cascades to table part_ee_ff_10_20
diff --git a/src/test/regress/expected/update.out b/src/test/regress/expected/update.out
index adc1fd7..df6eb30 100644
--- a/src/test/regress/expected/update.out
+++ b/src/test/regress/expected/update.out
@@ -182,3 +182,30 @@ INSERT INTO upsert_test VALUES (1, 'Bat') ON CONFLICT(a)
 
 DROP TABLE update_test;
 DROP TABLE upsert_test;
+-- update to a partition should check partition bound constraint for the new tuple
+create table range_parted (
+	a text,
+	b int
+) partition by range (a, b);
+create table part_a_1_a_10 partition of range_parted for values start ('a', 1) end ('a', 10);
+create table part_a_10_a_20 partition of range_parted for values start ('a', 10) end ('a', 20);
+create table part_b_1_b_10 partition of range_parted for values start ('b', 1) end ('b', 10);
+create table part_b_10_b_20 partition of range_parted for values start ('b', 10) end ('b', 20);
+insert into part_a_1_a_10 values ('a', 1);
+insert into part_b_10_b_20 values ('b', 10);
+-- fail
+update part_a_1_a_10 set a = 'b' where a = 'a';
+ERROR:  new row violates the partition boundary specification of "part_a_1_a_10"
+DETAIL:  Failing row contains (b, 1).
+update range_parted set b = b - 1 where b = 10;
+ERROR:  new row violates the partition boundary specification of "part_b_10_b_20"
+DETAIL:  Failing row contains (b, 9).
+-- ok
+update range_parted set b = b + 1 where b = 10;
+-- cleanup
+drop table range_parted cascade;
+NOTICE:  drop cascades to 4 other objects
+DETAIL:  drop cascades to table part_a_1_a_10
+drop cascades to table part_a_10_a_20
+drop cascades to table part_b_1_b_10
+drop cascades to table part_b_10_b_20
diff --git a/src/test/regress/sql/insert.sql b/src/test/regress/sql/insert.sql
index 7924d5d..4bf042e 100644
--- a/src/test/regress/sql/insert.sql
+++ b/src/test/regress/sql/insert.sql
@@ -84,3 +84,59 @@ create rule irule3 as on insert to inserttest2 do also
 drop table inserttest2;
 drop table inserttest;
 drop type insert_test_type;
+
+-- direct partition inserts should check partition bound constraint
+create table range_parted (
+	a text,
+	b int
+) partition by range (a, b);
+create table part_a_1_a_10 partition of range_parted for values start ('a', 1) end ('a', 10);
+create table part_a_10_a_20 partition of range_parted for values start ('a', 10) end ('a', 20);
+create table part_b_1_b_10 partition of range_parted for values start ('b', 1) end ('b', 10);
+create table part_b_10_b_20 partition of range_parted for values start ('b', 10) end ('b', 20);
+
+-- fail
+insert into part_a_1_a_10 values ('a', 11);
+insert into part_a_1_a_10 values ('b', 1);
+-- ok
+insert into part_a_1_a_10 values ('a', 1);
+-- fail
+insert into part_b_10_b_20 values ('b', 21);
+insert into part_b_10_b_20 values ('a', 10);
+-- ok
+insert into part_b_10_b_20 values ('b', 10);
+
+-- fail (a is null but a range partition key column should not be null)
+insert into part_b_10_b_20(b) values (10);
+
+create table list_parted (
+	a text,
+	b int
+) partition by list (upper(a));
+create table part_AA_BB partition of list_parted FOR VALUES IN ('AA', 'BB');
+create table part_CC_DD partition of list_parted FOR VALUES IN ('CC', 'DD');
+
+-- fail
+insert into part_AA_BB values ('cc', 1);
+insert into part_AA_BB values ('AAa', 1);
+-- ok
+insert into part_CC_DD values ('cC', 1);
+
+-- XXX - fail (a is null but part_AA_BB does not allow nulls in its list of values)
+-- insert into part_AA_BB (b) values (1);
+
+-- check in case of multi-level partitioned table
+create table part_EE_FF partition of list_parted for values in ('EE', 'FF') partition by range (b);
+create table part_EE_FF_1_10 partition of part_EE_FF for values start (1) end (10);
+create table part_EE_FF_10_20 partition of part_EE_FF for values start (10) end (20);
+
+-- fail (both its own and all ancestors' partition bound spec applies)
+insert into part_EE_FF_1_10 values ('EE', 11);
+insert into part_EE_FF_1_10 values ('cc', 1);
+-- ok
+insert into part_EE_FF_1_10 values ('ff', 1);
+insert into part_EE_FF_10_20 values ('ff', 11);
+
+-- cleanup
+drop table range_parted cascade;
+drop table list_parted cascade;
diff --git a/src/test/regress/sql/update.sql b/src/test/regress/sql/update.sql
index 5637c68..4997877 100644
--- a/src/test/regress/sql/update.sql
+++ b/src/test/regress/sql/update.sql
@@ -96,3 +96,24 @@ INSERT INTO upsert_test VALUES (1, 'Bat') ON CONFLICT(a)
 
 DROP TABLE update_test;
 DROP TABLE upsert_test;
+
+-- update to a partition should check partition bound constraint for the new tuple
+create table range_parted (
+	a text,
+	b int
+) partition by range (a, b);
+create table part_a_1_a_10 partition of range_parted for values start ('a', 1) end ('a', 10);
+create table part_a_10_a_20 partition of range_parted for values start ('a', 10) end ('a', 20);
+create table part_b_1_b_10 partition of range_parted for values start ('b', 1) end ('b', 10);
+create table part_b_10_b_20 partition of range_parted for values start ('b', 10) end ('b', 20);
+insert into part_a_1_a_10 values ('a', 1);
+insert into part_b_10_b_20 values ('b', 10);
+
+-- fail
+update part_a_1_a_10 set a = 'b' where a = 'a';
+update range_parted set b = b - 1 where b = 10;
+-- ok
+update range_parted set b = b + 1 where b = 10;
+
+-- cleanup
+drop table range_parted cascade;
-- 
1.7.1

