From ff5f3bdedf94e08ff89db93df95c404188092d4d Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Mon, 6 Nov 2017 17:04:55 +0100
Subject: [PATCH v3 5/5] allow indexes on partitioned tables to be unique

---
 src/backend/commands/indexcmds.c           |  74 ++++++++++++++--
 src/backend/parser/parse_utilcmd.c         |  24 -----
 src/test/regress/expected/alter_table.out  |   8 --
 src/test/regress/expected/create_table.out |  12 ---
 src/test/regress/expected/indexing.out     | 135 ++++++++++++++++++++++++++++-
 src/test/regress/sql/alter_table.sql       |   2 -
 src/test/regress/sql/create_table.sql      |   8 --
 src/test/regress/sql/indexing.sql          |  72 ++++++++++++++-
 8 files changed, 269 insertions(+), 66 deletions(-)

diff --git a/src/backend/commands/indexcmds.c b/src/backend/commands/indexcmds.c
index c289f62b9d..446c58368e 100644
--- a/src/backend/commands/indexcmds.c
+++ b/src/backend/commands/indexcmds.c
@@ -425,20 +425,11 @@ DefineIndex(Oid relationId,
 					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 					 errmsg("cannot create index on partitioned table \"%s\" concurrently",
 							RelationGetRelationName(rel))));
-		if (stmt->unique)
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("cannot create unique index on partitioned table \"%s\"",
-							RelationGetRelationName(rel))));
 		if (stmt->excludeOpNames)
 			ereport(ERROR,
 					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 					 errmsg("cannot create exclusion constraints on partitioned table \"%s\"",
 							RelationGetRelationName(rel))));
-		if (is_alter_table)
-			ereport(ERROR,
-					(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-					 errmsg("cannot create constraints on partitioned tables")));
 	}
 
 	/*
@@ -635,6 +626,71 @@ DefineIndex(Oid relationId,
 		index_check_primary_key(rel, indexInfo, is_alter_table);
 
 	/*
+	 * If this table is partitioned and we're creating a unique index or a
+	 * primary key, make sure that the indexed columns are part of the
+	 * partition key.  Otherwise it would be possible to violate uniqueness by
+	 * putting values that ought to be unique in different partitions.
+	 *
+	 * We could lift this limitation if we had global indexes, but those have
+	 * their own problems, so this is a useful feature combination.
+	 *
+	 * XXX in some cases rel->rd_partkey is NULL, which caused this code to
+	 * crash.  In what cases can that happen?
+	 */
+	if (partitioned && (stmt->unique || stmt->primary) &&
+		(rel->rd_partkey != NULL))
+	{
+		int			i;
+		PartitionKey key = rel->rd_partkey;
+
+		/*
+		 * A partitioned table can have unique indexes, as long as all the
+		 * columns in the unique key appear in the partition key.  Because
+		 * partitions are non-overlapping, this guarantees that there is a
+		 * single partition that can contain any given key, and thus the
+		 * uniqueness is guaranteed globally.
+		 */
+
+		for (i = 0; i < indexInfo->ii_NumIndexAttrs; i++)
+		{
+			bool	found;
+			int		j;
+
+			/*
+			 * Expression indexes are not supported yet.  We can probably use
+			 * pull_varattnos() on the expression, and verify that all the
+			 * vars mentioned correspond to columns of the partition key.
+			 */
+			if (indexInfo->ii_KeyAttrNumbers[i] == 0)
+				ereport(ERROR,
+					   (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						errmsg("UNIQUE indexes containing expressions are not supported")));
+
+			found = false;
+			for (j = 0; j < key->partnatts; j++)
+			{
+				if (indexInfo->ii_KeyAttrNumbers[i] == key->partattrs[j])
+				{
+					found = true;
+					break;	/* all good */
+				}
+			}
+			if (!found)
+			{
+				char	   *attname;
+
+				attname = NameStr(TupleDescAttr(RelationGetDescr(rel),
+												indexInfo->ii_KeyAttrNumbers[i] - 1)->attname);
+				ereport(ERROR,
+						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+						 errmsg("UNIQUE index on table \"%s\" contains column \"%s\" which is not part of the partition key",
+								RelationGetRelationName(rel), attname)));
+			}
+		}
+	}
+
+
+	/*
 	 * We disallow indexes on system columns other than OID.  They would not
 	 * necessarily get updated correctly, and they don't seem useful anyway.
 	 */
diff --git a/src/backend/parser/parse_utilcmd.c b/src/backend/parser/parse_utilcmd.c
index 938e73f8ad..364bd4c81d 100644
--- a/src/backend/parser/parse_utilcmd.c
+++ b/src/backend/parser/parse_utilcmd.c
@@ -692,12 +692,6 @@ transformColumnDefinition(CreateStmtContext *cxt, ColumnDef *column)
 							 errmsg("primary key constraints are not supported on foreign tables"),
 							 parser_errposition(cxt->pstate,
 												constraint->location)));
-				if (cxt->ispartitioned)
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("primary key constraints are not supported on partitioned tables"),
-							 parser_errposition(cxt->pstate,
-												constraint->location)));
 				/* FALL THRU */
 
 			case CONSTR_UNIQUE:
@@ -707,12 +701,6 @@ transformColumnDefinition(CreateStmtContext *cxt, ColumnDef *column)
 							 errmsg("unique constraints are not supported on foreign tables"),
 							 parser_errposition(cxt->pstate,
 												constraint->location)));
-				if (cxt->ispartitioned)
-					ereport(ERROR,
-							(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-							 errmsg("unique constraints are not supported on partitioned tables"),
-							 parser_errposition(cxt->pstate,
-												constraint->location)));
 				if (constraint->keys == NIL)
 					constraint->keys = list_make1(makeString(column->colname));
 				cxt->ixconstraints = lappend(cxt->ixconstraints, constraint);
@@ -809,12 +797,6 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("primary key constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 errmsg("primary key constraints are not supported on partitioned tables"),
-						 parser_errposition(cxt->pstate,
-											constraint->location)));
 			cxt->ixconstraints = lappend(cxt->ixconstraints, constraint);
 			break;
 
@@ -825,12 +807,6 @@ transformTableConstraint(CreateStmtContext *cxt, Constraint *constraint)
 						 errmsg("unique constraints are not supported on foreign tables"),
 						 parser_errposition(cxt->pstate,
 											constraint->location)));
-			if (cxt->ispartitioned)
-				ereport(ERROR,
-						(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
-						 errmsg("unique constraints are not supported on partitioned tables"),
-						 parser_errposition(cxt->pstate,
-											constraint->location)));
 			cxt->ixconstraints = lappend(cxt->ixconstraints, constraint);
 			break;
 
diff --git a/src/test/regress/expected/alter_table.out b/src/test/regress/expected/alter_table.out
index 5a96b0b5e1..e7cf6ae470 100644
--- a/src/test/regress/expected/alter_table.out
+++ b/src/test/regress/expected/alter_table.out
@@ -3229,14 +3229,6 @@ CREATE TABLE partitioned (
 	a int,
 	b int
 ) PARTITION BY RANGE (a, (a+b+1));
-ALTER TABLE partitioned ADD UNIQUE (a);
-ERROR:  unique constraints are not supported on partitioned tables
-LINE 1: ALTER TABLE partitioned ADD UNIQUE (a);
-                                    ^
-ALTER TABLE partitioned ADD PRIMARY KEY (a);
-ERROR:  primary key constraints are not supported on partitioned tables
-LINE 1: ALTER TABLE partitioned ADD PRIMARY KEY (a);
-                                    ^
 ALTER TABLE partitioned ADD FOREIGN KEY (a) REFERENCES blah;
 ERROR:  foreign key constraints are not supported on partitioned tables
 LINE 1: ALTER TABLE partitioned ADD FOREIGN KEY (a) REFERENCES blah;
diff --git a/src/test/regress/expected/create_table.out b/src/test/regress/expected/create_table.out
index 335cd37e18..a89e0c17df 100644
--- a/src/test/regress/expected/create_table.out
+++ b/src/test/regress/expected/create_table.out
@@ -276,12 +276,6 @@ CREATE TABLE partitioned (
 ) PARTITION BY LIST (a1, a2);	-- fail
 ERROR:  cannot use "list" partition strategy with more than one column
 -- unsupported constraint type for partitioned tables
-CREATE TABLE partitioned (
-	a int PRIMARY KEY
-) PARTITION BY RANGE (a);
-ERROR:  primary key constraints are not supported on partitioned tables
-LINE 2:  a int PRIMARY KEY
-               ^
 CREATE TABLE pkrel (
 	a int PRIMARY KEY
 );
@@ -293,12 +287,6 @@ LINE 2:  a int REFERENCES pkrel(a)
                ^
 DROP TABLE pkrel;
 CREATE TABLE partitioned (
-	a int UNIQUE
-) PARTITION BY RANGE (a);
-ERROR:  unique constraints are not supported on partitioned tables
-LINE 2:  a int UNIQUE
-               ^
-CREATE TABLE partitioned (
 	a int,
 	EXCLUDE USING gist (a WITH &&)
 ) PARTITION BY RANGE (a);
diff --git a/src/test/regress/expected/indexing.out b/src/test/regress/expected/indexing.out
index cde29bebb5..7ec8c5e3f7 100644
--- a/src/test/regress/expected/indexing.out
+++ b/src/test/regress/expected/indexing.out
@@ -25,8 +25,6 @@ drop table idxpart;
 -- Some unsupported cases
 create table idxpart (a int, b int, c text) partition by range (a);
 create table idxpart1 partition of idxpart for values from (0) to (10);
-create unique index on idxpart (a);
-ERROR:  cannot create unique index on partitioned table "idxpart"
 create index concurrently on idxpart (a);
 ERROR:  cannot create index on partitioned table "idxpart" concurrently
 drop table idxpart;
@@ -328,3 +326,136 @@ select attrelid::regclass, attname, attnum from pg_attribute
 (8 rows)
 
 drop table idxpart;
+--
+-- Constraint-related indexes
+--
+-- Verify that it works to add primary key / unique to partitioned tables
+create table idxpart (a int primary key, b int) partition by range (b);
+ERROR:  UNIQUE index on table "idxpart" contains column "a" which is not part of the partition key
+create table idxpart (a int primary key, b int) partition by range (a);
+\d idxpart
+              Table "public.idxpart"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
+ a      | integer |           | not null | 
+ b      | integer |           |          | 
+Partition key: RANGE (a)
+Indexes:
+    "idxpart_pkey" PRIMARY KEY, btree (a)
+
+drop table idxpart;
+create table idxpart (a int unique, b int) partition by range (a);
+\d idxpart
+              Table "public.idxpart"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
+ a      | integer |           |          | 
+ b      | integer |           |          | 
+Partition key: RANGE (a)
+Indexes:
+    "idxpart_a_key" UNIQUE CONSTRAINT, btree (a)
+
+drop table idxpart;
+-- but not other types of index-based constraints
+create table idxpart (a int, exclude (a with = )) partition by range (a);
+ERROR:  exclusion constraints are not supported on partitioned tables
+LINE 1: create table idxpart (a int, exclude (a with = )) partition ...
+                                     ^
+-- It works to add primary keys after the partitioned table is created
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add primary key (a);
+\d idxpart
+              Table "public.idxpart"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
+ a      | integer |           | not null | 
+ b      | integer |           |          | 
+Partition key: RANGE (a)
+Indexes:
+    "idxpart_pkey" PRIMARY KEY, btree (a)
+
+drop table idxpart;
+-- It works to add unique constraints after the partitioned table is created
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add unique (a);
+\d idxpart
+              Table "public.idxpart"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
+ a      | integer |           |          | 
+ b      | integer |           |          | 
+Partition key: RANGE (a)
+Indexes:
+    "idxpart_a_key" UNIQUE CONSTRAINT, btree (a)
+
+drop table idxpart;
+-- Exclusion constraints cannot be added
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add exclude (a with =);
+ERROR:  exclusion constraints are not supported on partitioned tables
+LINE 1: alter table idxpart add exclude (a with =);
+                                ^
+drop table idxpart;
+-- It is an error to add a constraint to columns that are not in the partition
+-- key.
+create table idxpart (a int, b int, primary key (a, b)) partition by range (a);
+ERROR:  UNIQUE index on table "idxpart" contains column "b" which is not part of the partition key
+create table idxpart (a int, b int, primary key (a, b)) partition by range (b);
+ERROR:  UNIQUE index on table "idxpart" contains column "a" which is not part of the partition key
+create table idxpart (a int, b int, unique (a, b)) partition by range (a);
+ERROR:  UNIQUE index on table "idxpart" contains column "b" which is not part of the partition key
+create table idxpart (a int, b int, unique (a, b)) partition by range (b);
+ERROR:  UNIQUE index on table "idxpart" contains column "a" which is not part of the partition key
+-- but using a partial set of columns is okay
+create table idxpart (a int, b int, c int primary key) partition by range (b);
+ERROR:  UNIQUE index on table "idxpart" contains column "c" which is not part of the partition key
+drop table idxpart;
+ERROR:  table "idxpart" does not exist
+create table idxpart (a int, b int, primary key (a)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, primary key (b)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, c int unique) partition by range (b);
+ERROR:  UNIQUE index on table "idxpart" contains column "c" which is not part of the partition key
+drop table idxpart;
+ERROR:  table "idxpart" does not exist
+create table idxpart (a int, b int, unique (a)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, unique (b)) partition by range (a, b);
+drop table idxpart;
+-- When (sub)partitions are created, they also contain the constraint
+create table idxpart (a int, b int, primary key (a, b)) partition by range (a, b);
+create table idxpart1 partition of idxpart for values from (1, 1) to (10, 10);
+create table idxpart2 partition of idxpart for values from (10, 10) to (20, 20)
+  partition by range (b);
+create table idxpart21 partition of idxpart2 for values from (10) to (15);
+create table idxpart22 partition of idxpart2 for values from (15) to (20);
+create table idxpart3 (b int not null, a int not null);
+alter table idxpart attach partition idxpart3 for values from (20, 20) to (30, 30);
+select conname, contype, conrelid::regclass, conindid::regclass, conkey
+  from pg_constraint where conrelid::regclass::text like 'idxpart%'
+  order by conname;
+    conname     | contype | conrelid  |    conindid    | conkey 
+----------------+---------+-----------+----------------+--------
+ idxpart1_pkey  | p       | idxpart1  | idxpart1_pkey  | {1,2}
+ idxpart21_pkey | p       | idxpart21 | idxpart21_pkey | {1,2}
+ idxpart22_pkey | p       | idxpart22 | idxpart22_pkey | {1,2}
+ idxpart2_pkey  | p       | idxpart2  | idxpart2_pkey  | {1,2}
+ idxpart3_pkey  | p       | idxpart3  | idxpart3_pkey  | {2,1}
+ idxpart_pkey   | p       | idxpart   | idxpart_pkey   | {1,2}
+(6 rows)
+
+drop table idxpart;
+create table idxpart (a int primary key, b int) partition by range (a);
+create table idxpart1 partition of idxpart for values from (1) to (10);
+\d idxpart1
+              Table "public.idxpart1"
+ Column |  Type   | Collation | Nullable | Default 
+--------+---------+-----------+----------+---------
+ a      | integer |           | not null | 
+ b      | integer |           |          | 
+Partition of: idxpart FOR VALUES FROM (1) TO (10)
+Indexes:
+    "idxpart1_pkey" PRIMARY KEY, btree (a)
+
+drop table idxpart;
diff --git a/src/test/regress/sql/alter_table.sql b/src/test/regress/sql/alter_table.sql
index 02a33ca7c4..115ed1efad 100644
--- a/src/test/regress/sql/alter_table.sql
+++ b/src/test/regress/sql/alter_table.sql
@@ -2000,8 +2000,6 @@ CREATE TABLE partitioned (
 	a int,
 	b int
 ) PARTITION BY RANGE (a, (a+b+1));
-ALTER TABLE partitioned ADD UNIQUE (a);
-ALTER TABLE partitioned ADD PRIMARY KEY (a);
 ALTER TABLE partitioned ADD FOREIGN KEY (a) REFERENCES blah;
 ALTER TABLE partitioned ADD EXCLUDE USING gist (a WITH &&);
 
diff --git a/src/test/regress/sql/create_table.sql b/src/test/regress/sql/create_table.sql
index b77b476436..673db6b77d 100644
--- a/src/test/regress/sql/create_table.sql
+++ b/src/test/regress/sql/create_table.sql
@@ -294,10 +294,6 @@ CREATE TABLE partitioned (
 ) PARTITION BY LIST (a1, a2);	-- fail
 
 -- unsupported constraint type for partitioned tables
-CREATE TABLE partitioned (
-	a int PRIMARY KEY
-) PARTITION BY RANGE (a);
-
 CREATE TABLE pkrel (
 	a int PRIMARY KEY
 );
@@ -307,10 +303,6 @@ CREATE TABLE partitioned (
 DROP TABLE pkrel;
 
 CREATE TABLE partitioned (
-	a int UNIQUE
-) PARTITION BY RANGE (a);
-
-CREATE TABLE partitioned (
 	a int,
 	EXCLUDE USING gist (a WITH &&)
 ) PARTITION BY RANGE (a);
diff --git a/src/test/regress/sql/indexing.sql b/src/test/regress/sql/indexing.sql
index 53a28f120f..0831837560 100644
--- a/src/test/regress/sql/indexing.sql
+++ b/src/test/regress/sql/indexing.sql
@@ -14,7 +14,6 @@ drop table idxpart;
 -- Some unsupported cases
 create table idxpart (a int, b int, c text) partition by range (a);
 create table idxpart1 partition of idxpart for values from (0) to (10);
-create unique index on idxpart (a);
 create index concurrently on idxpart (a);
 drop table idxpart;
 
@@ -140,3 +139,74 @@ select attrelid::regclass, attname, attnum from pg_attribute
   where attrelid in ('idxpart'::regclass, 'idxpart1'::regclass) and attnum > 0
   order by attrelid::regclass, attnum;
 drop table idxpart;
+
+--
+-- Constraint-related indexes
+--
+
+-- Verify that it works to add primary key / unique to partitioned tables
+create table idxpart (a int primary key, b int) partition by range (b);
+create table idxpart (a int primary key, b int) partition by range (a);
+\d idxpart
+drop table idxpart;
+create table idxpart (a int unique, b int) partition by range (a);
+\d idxpart
+drop table idxpart;
+-- but not other types of index-based constraints
+create table idxpart (a int, exclude (a with = )) partition by range (a);
+
+-- It works to add primary keys after the partitioned table is created
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add primary key (a);
+\d idxpart
+drop table idxpart;
+
+-- It works to add unique constraints after the partitioned table is created
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add unique (a);
+\d idxpart
+drop table idxpart;
+
+-- Exclusion constraints cannot be added
+create table idxpart (a int, b int) partition by range (a);
+alter table idxpart add exclude (a with =);
+drop table idxpart;
+
+-- It is an error to add a constraint to columns that are not in the partition
+-- key.
+create table idxpart (a int, b int, primary key (a, b)) partition by range (a);
+create table idxpart (a int, b int, primary key (a, b)) partition by range (b);
+create table idxpart (a int, b int, unique (a, b)) partition by range (a);
+create table idxpart (a int, b int, unique (a, b)) partition by range (b);
+-- but using a partial set of columns is okay
+create table idxpart (a int, b int, c int primary key) partition by range (b);
+drop table idxpart;
+create table idxpart (a int, b int, primary key (a)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, primary key (b)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, c int unique) partition by range (b);
+drop table idxpart;
+create table idxpart (a int, b int, unique (a)) partition by range (a, b);
+drop table idxpart;
+create table idxpart (a int, b int, unique (b)) partition by range (a, b);
+drop table idxpart;
+
+-- When (sub)partitions are created, they also contain the constraint
+create table idxpart (a int, b int, primary key (a, b)) partition by range (a, b);
+create table idxpart1 partition of idxpart for values from (1, 1) to (10, 10);
+create table idxpart2 partition of idxpart for values from (10, 10) to (20, 20)
+  partition by range (b);
+create table idxpart21 partition of idxpart2 for values from (10) to (15);
+create table idxpart22 partition of idxpart2 for values from (15) to (20);
+create table idxpart3 (b int not null, a int not null);
+alter table idxpart attach partition idxpart3 for values from (20, 20) to (30, 30);
+select conname, contype, conrelid::regclass, conindid::regclass, conkey
+  from pg_constraint where conrelid::regclass::text like 'idxpart%'
+  order by conname;
+drop table idxpart;
+
+create table idxpart (a int primary key, b int) partition by range (a);
+create table idxpart1 partition of idxpart for values from (1) to (10);
+\d idxpart1
+drop table idxpart;
-- 
2.11.0

