From 31423c4dc43587ef5221bf0ae87484199487da40 Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Mon, 15 May 2017 14:00:54 +0900
Subject: [PATCH 1/2] Fix to make partitioned tables work with event triggers

When creating range partitioned tables, AlterTableInternal is
called which expects the event trigger context info to be set by
EventTriggerAlterTableStart().  Lacking that, a crash occurred
when creating range partitioned tables if event triggers have
been defined in the database.  Instead of fixing it by adding
appropriate event trigger initialization calls in DefineRelation(),
rearrange code such that AlterTableInternal() is no longer directly
used to create NOT NULL constraints on partition key columns.
Now, transformCreateStmt() calls transformPartitionKey() which adds
an AlterTableStmt to add the necessary constraints.  Event triggers
will be taken care of properly when that statement is later
executed by ProcessUtilitySlow().

Reported by Mark Dilger
Report: https://postgr.es/m/4A9B8269-9771-4FB7-BFA3-84249800917D%40gmail.com
---
 src/backend/commands/tablecmds.c            | 31 +------------
 src/backend/parser/parse_utilcmd.c          | 70 ++++++++++++++++++++++++-----
 src/test/regress/expected/event_trigger.out | 32 +++++++++++++
 src/test/regress/sql/event_trigger.sql      | 35 +++++++++++++++
 4 files changed, 128 insertions(+), 40 deletions(-)

diff --git a/src/backend/commands/tablecmds.c b/src/backend/commands/tablecmds.c
index e259378051..e4d2bfb6b0 100644
--- a/src/backend/commands/tablecmds.c
+++ b/src/backend/commands/tablecmds.c
@@ -805,13 +805,11 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 	if (stmt->partspec)
 	{
 		char		strategy;
-		int			partnatts,
-					i;
+		int			partnatts;
 		AttrNumber	partattrs[PARTITION_MAX_KEYS];
 		Oid			partopclass[PARTITION_MAX_KEYS];
 		Oid			partcollation[PARTITION_MAX_KEYS];
 		List	   *partexprs = NIL;
-		List	   *cmds = NIL;
 
 		/*
 		 * We need to transform the raw parsetrees corresponding to partition
@@ -828,33 +826,6 @@ DefineRelation(CreateStmt *stmt, char relkind, Oid ownerId,
 		partnatts = list_length(stmt->partspec->partParams);
 		StorePartitionKey(rel, strategy, partnatts, partattrs, partexprs,
 						  partopclass, partcollation);
-
-		/* Force key columns to be NOT NULL when using range partitioning */
-		if (strategy == PARTITION_STRATEGY_RANGE)
-		{
-			for (i = 0; i < partnatts; i++)
-			{
-				AttrNumber	partattno = partattrs[i];
-				Form_pg_attribute attform = descriptor->attrs[partattno - 1];
-
-				if (partattno != 0 && !attform->attnotnull)
-				{
-					/* Add a subcommand to make this one NOT NULL */
-					AlterTableCmd *cmd = makeNode(AlterTableCmd);
-
-					cmd->subtype = AT_SetNotNull;
-					cmd->name = pstrdup(NameStr(attform->attname));
-					cmds = lappend(cmds, cmd);
-				}
-			}
-
-			/*
-			 * Although, there cannot be any partitions yet, we still need to
-			 * pass true for recurse; ATPrepSetNotNull() complains if we don't
-			 */
-			if (cmds != NIL)
-				AlterTableInternal(RelationGetRelid(rel), cmds, true);
-		}
 	}
 
 	/*
diff --git a/src/backend/parser/parse_utilcmd.c b/src/backend/parser/parse_utilcmd.c
index 882955bb1c..8aad7138ac 100644
--- a/src/backend/parser/parse_utilcmd.c
+++ b/src/backend/parser/parse_utilcmd.c
@@ -90,7 +90,7 @@ typedef struct
 	List	   *alist;			/* "after list" of things to do after creating
 								 * the table */
 	IndexStmt  *pkey;			/* PRIMARY KEY index, if any */
-	bool		ispartitioned;	/* true if table is partitioned */
+	PartitionSpec *partspec;	/* Partition key spec if partitioned table */
 	Node	   *partbound;		/* transformed FOR VALUES */
 } CreateStmtContext;
 
@@ -135,6 +135,7 @@ static void transformConstraintAttrs(CreateStmtContext *cxt,
 static void transformColumnType(CreateStmtContext *cxt, ColumnDef *column);
 static void setSchemaName(char *context_schema, char **stmt_schema_name);
 static void transformPartitionCmd(CreateStmtContext *cxt, PartitionCmd *cmd);
+static void transformPartitionKey(CreateStmtContext *cxt);
 
 
 /*
@@ -235,7 +236,7 @@ transformCreateStmt(CreateStmt *stmt, const char *queryString)
 	cxt.blist = NIL;
 	cxt.alist = NIL;
 	cxt.pkey = NULL;
-	cxt.ispartitioned = stmt->partspec != NULL;
+	cxt.partspec = stmt->partspec;
 
 	/*
 	 * Notice that we allow OIDs here only for plain tables, even though
@@ -345,6 +346,11 @@ transformCreateStmt(CreateStmt *stmt, const char *queryString)
 	transformCheckConstraints(&cxt, true);
 
 	/*
+	 * Postprocess partition key that gives rise to NOT NULL constraints.
+	 */
+	transformPartitionKey(&cxt);
+
+	/*
 	 * Output results.
 	 */
 	stmt->tableElts = cxt.columns;
@@ -684,7 +690,7 @@ transformColumnDefinition(CreateStmtContext *cxt, ColumnDef *column)
 							 errmsg("primary key constraints are not supported on foreign tables"),
 							 parser_errposition(cxt->pstate,
 												constraint->location)));
-				if (cxt->ispartitioned)
+				if (cxt->partspec)
 					ereport(ERROR,
 							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 							 errmsg("primary key constraints are not supported on partitioned tables"),
@@ -699,7 +705,7 @@ transformColumnDefinition(CreateStmtContext *cxt, ColumnDef *column)
 							 errmsg("unique constraints are not supported on foreign tables"),
 							 parser_errposition(cxt->pstate,
 												constraint->location)));
-				if (cxt->ispartitioned)
+				if (cxt->partspec)
 					ereport(ERROR,
 							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 							 errmsg("unique constraints are not supported on partitioned tables"),
@@ -722,7 +728,7 @@ transformColumnDefinition(CreateStmtContext *cxt, ColumnDef *column)
 							 errmsg("foreign key constraints are not supported on foreign tables"),
 							 parser_errposition(cxt->pstate,
 												constraint->location)));
-				if (cxt->ispartitioned)
+				if (cxt->partspec)
 					ereport(ERROR,
 							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 							 errmsg("foreign key constraints are not supported on partitioned tables"),
@@ -801,7 +807,7 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("primary key constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
+			if (cxt->partspec)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("primary key constraints are not supported on partitioned tables"),
@@ -817,7 +823,7 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("unique constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
+			if (cxt->partspec)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("unique constraints are not supported on partitioned tables"),
@@ -833,7 +839,7 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("exclusion constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
+			if (cxt->partspec)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("exclusion constraints are not supported on partitioned tables"),
@@ -853,7 +859,7 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("foreign key constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
+			if (cxt->partspec)
 				ereport(ERROR,
 						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 						 errmsg("foreign key constraints are not supported on partitioned tables"),
@@ -2103,6 +2109,48 @@ transformIndexConstraint(Constraint *constraint, CreateStmtContext *cxt)
 }
 
 /*
+ * transformPartitionKey
+ *		Check if we need to emit NOT NULL constraints for partition keys
+ *
+ * When using range partitioning, we need to add NOT NULL constraints to
+ * the simple columns in the partition key.
+ */
+static void
+transformPartitionKey(CreateStmtContext *cxt)
+{
+	PartitionSpec *spec = cxt->partspec;
+	List	 *cmds = NIL;
+	ListCell *lc;
+	AlterTableStmt *alterstmt;
+
+	if (spec == NULL || strcmp(spec->strategy, "range") != 0)
+		return;
+
+	foreach (lc, spec->partParams)
+	{
+		PartitionElem   *pelem = lfirst(lc);
+
+		/* Only consider simple columns, not expressions. */
+		if (pelem->name)
+		{
+			/* Add a subcommand to make this one NOT NULL */
+			AlterTableCmd *cmd = makeNode(AlterTableCmd);
+
+			cmd->subtype = AT_SetNotNull;
+			cmd->name = pstrdup(pelem->name);
+			cmds = lappend(cmds, cmd);
+		}
+	}
+
+	alterstmt = makeNode(AlterTableStmt);
+	alterstmt->relation = cxt->relation;
+	alterstmt->cmds = cmds;
+	alterstmt->relkind = OBJECT_TABLE;
+
+	cxt->alist = lappend(cxt->alist, alterstmt);
+}
+
+/*
  * transformCheckConstraints
  *		handle CHECK constraints
  *
@@ -2681,7 +2729,9 @@ transformAlterTableStmt(Oid relid, AlterTableStmt *stmt,
 	cxt.blist = NIL;
 	cxt.alist = NIL;
 	cxt.pkey = NULL;
-	cxt.ispartitioned = (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE);
+	/* We don't really need the contents of PartitionSpec here. */
+	cxt.partspec = (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+						? makeNode(PartitionSpec) : NULL;
 	cxt.partbound = NULL;
 
 	/*
diff --git a/src/test/regress/expected/event_trigger.out b/src/test/regress/expected/event_trigger.out
index 906dcb8b31..d5021b648e 100644
--- a/src/test/regress/expected/event_trigger.out
+++ b/src/test/regress/expected/event_trigger.out
@@ -454,3 +454,35 @@ NOTICE:  DROP POLICY - ddl_command_end
 DROP EVENT TRIGGER start_rls_command;
 DROP EVENT TRIGGER end_rls_command;
 DROP EVENT TRIGGER sql_drop_command;
+-- check events triggers work correctly in the case of creating partitioned
+-- tables, especially range-partitioned ones which invoke
+-- AlterTableInternal().
+CREATE OR REPLACE FUNCTION create_partitioned_table_before_proc() RETURNS
+event_trigger AS $$
+BEGIN
+	RAISE WARNING 'preparing to create table: % %', tg_event, tg_tag;
+END;
+$$ LANGUAGE plpgsql;
+CREATE EVENT TRIGGER create_partitioned_table_before_trigger
+	ON ddl_command_start
+	WHEN TAG IN ('CREATE TABLE')
+	EXECUTE PROCEDURE create_partitioned_table_before_proc();
+CREATE OR REPLACE FUNCTION create_partitioned_table_after_proc()
+RETURNS event_trigger AS $$
+BEGIN
+	RAISE WARNING 'finished creating table: % %', tg_event, tg_tag;
+END;
+$$ LANGUAGE plpgsql;
+CREATE EVENT TRIGGER create_partitioned_table_after_trigger
+	ON ddl_command_end
+	WHEN TAG IN ('CREATE TABLE')
+	EXECUTE PROCEDURE create_partitioned_table_after_proc();
+-- create a range-partitioned table
+CREATE TABLE range_partitioned (a int) PARTITION BY RANGE (a);
+WARNING:  preparing to create table: ddl_command_start CREATE TABLE
+WARNING:  finished creating table: ddl_command_end CREATE TABLE
+DROP TABLE range_partitioned;
+DROP EVENT TRIGGER create_partitioned_table_before_trigger;
+DROP EVENT TRIGGER create_partitioned_table_after_trigger;
+DROP FUNCTION create_partitioned_table_before_proc();
+DROP FUNCTION create_partitioned_table_after_proc();
diff --git a/src/test/regress/sql/event_trigger.sql b/src/test/regress/sql/event_trigger.sql
index b65bf3ec66..1491fd7f07 100644
--- a/src/test/regress/sql/event_trigger.sql
+++ b/src/test/regress/sql/event_trigger.sql
@@ -375,3 +375,38 @@ DROP POLICY p2 ON event_trigger_test;
 DROP EVENT TRIGGER start_rls_command;
 DROP EVENT TRIGGER end_rls_command;
 DROP EVENT TRIGGER sql_drop_command;
+
+-- check events triggers work correctly in the case of creating partitioned
+-- tables, especially range-partitioned ones which invoke
+-- AlterTableInternal().
+CREATE OR REPLACE FUNCTION create_partitioned_table_before_proc() RETURNS
+event_trigger AS $$
+BEGIN
+	RAISE WARNING 'preparing to create table: % %', tg_event, tg_tag;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE EVENT TRIGGER create_partitioned_table_before_trigger
+	ON ddl_command_start
+	WHEN TAG IN ('CREATE TABLE')
+	EXECUTE PROCEDURE create_partitioned_table_before_proc();
+
+CREATE OR REPLACE FUNCTION create_partitioned_table_after_proc()
+RETURNS event_trigger AS $$
+BEGIN
+	RAISE WARNING 'finished creating table: % %', tg_event, tg_tag;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE EVENT TRIGGER create_partitioned_table_after_trigger
+	ON ddl_command_end
+	WHEN TAG IN ('CREATE TABLE')
+	EXECUTE PROCEDURE create_partitioned_table_after_proc();
+
+-- create a range-partitioned table
+CREATE TABLE range_partitioned (a int) PARTITION BY RANGE (a);
+DROP TABLE range_partitioned;
+DROP EVENT TRIGGER create_partitioned_table_before_trigger;
+DROP EVENT TRIGGER create_partitioned_table_after_trigger;
+DROP FUNCTION create_partitioned_table_before_proc();
+DROP FUNCTION create_partitioned_table_after_proc();
-- 
2.11.0

