From c2055ffc5235b28d37dbb89a1818e8f7df50f8da Mon Sep 17 00:00:00 2001
From: Oliver Tan <otan@cockroachlabs.com>
Date: Fri, 22 Jan 2021 11:13:57 +1100
Subject: [PATCH] sql: teach ADD UNIQUE CONSTRAINT to respect PARTITION ALL BY

Release note (sql change): Fix-up ALTER TABLE ... ADD CONSTRAINT ...
UNIQUE to partition correctly under a PARTITION ALL BY table.
---
 .../testdata/logic_test/partitioning_implicit | 37 ++++++++++++-------
 pkg/sql/alter_table.go                        | 35 +++++-------------
 2 files changed, 34 insertions(+), 38 deletions(-)

diff --git a/pkg/ccl/logictestccl/testdata/logic_test/partitioning_implicit b/pkg/ccl/logictestccl/testdata/logic_test/partitioning_implicit
index 31b8537e6306..ffcb2a0793d4 100644
--- a/pkg/ccl/logictestccl/testdata/logic_test/partitioning_implicit
+++ b/pkg/ccl/logictestccl/testdata/logic_test/partitioning_implicit
@@ -221,10 +221,10 @@ CREATE TABLE public.t (
   pk int PRIMARY KEY,
   pk2 int NOT NULL,
   partition_by int,
-  a int,
-  b int,
-  c int,
-  d int,
+  a int NOT NULL,
+  b int NOT NULL,
+  c int NOT NULL,
+  d int NOT NULL,
   INDEX (a),
   UNIQUE (b),
   INDEX (partition_by, c),
@@ -242,6 +242,9 @@ CREATE INDEX created_idx ON t(c) PARTITION BY LIST (d) (
 statement ok
 CREATE INDEX created_idx ON t(c)
 
+statement ok
+ALTER TABLE t ADD CONSTRAINT unique_c_d UNIQUE(c, d)
+
 query T
 SELECT create_statement FROM [SHOW CREATE TABLE t]
 ----
@@ -249,15 +252,16 @@ CREATE TABLE public.t (
   pk INT8 NOT NULL,
   pk2 INT8 NOT NULL,
   partition_by INT8 NULL,
-  a INT8 NULL,
-  b INT8 NULL,
-  c INT8 NULL,
-  d INT8 NULL,
+  a INT8 NOT NULL,
+  b INT8 NOT NULL,
+  c INT8 NOT NULL,
+  d INT8 NOT NULL,
   CONSTRAINT "primary" PRIMARY KEY (pk ASC),
   INDEX t_a_idx (a ASC),
   UNIQUE INDEX t_b_key (b ASC),
   INDEX t_partition_by_c_idx (partition_by ASC, c ASC),
   INDEX created_idx (c ASC),
+  UNIQUE INDEX unique_c_d (c ASC, d ASC),
   FAMILY fam_0_pk_pk2_partition_by_a_b_c_d (pk, pk2, partition_by, a, b, c, d)
 ) PARTITION ALL BY LIST (partition_by) (
   PARTITION one VALUES IN ((1)),
@@ -281,6 +285,9 @@ t_b_key               b             false
 t_b_key               partition_by  true
 t_partition_by_c_idx  c             false
 t_partition_by_c_idx  partition_by  false
+unique_c_d            c             false
+unique_c_d            d             false
+unique_c_d            partition_by  true
 
 statement ok
 ALTER TABLE t ALTER PRIMARY KEY USING COLUMNS (pk2)
@@ -292,16 +299,17 @@ CREATE TABLE public.t (
   pk INT8 NOT NULL,
   pk2 INT8 NOT NULL,
   partition_by INT8 NULL,
-  a INT8 NULL,
-  b INT8 NULL,
-  c INT8 NULL,
-  d INT8 NULL,
+  a INT8 NOT NULL,
+  b INT8 NOT NULL,
+  c INT8 NOT NULL,
+  d INT8 NOT NULL,
   CONSTRAINT "primary" PRIMARY KEY (pk2 ASC),
   UNIQUE INDEX t_pk_key (pk ASC),
   INDEX t_a_idx (a ASC),
   UNIQUE INDEX t_b_key (b ASC),
   INDEX t_partition_by_c_idx (partition_by ASC, c ASC),
   INDEX created_idx (c ASC),
+  UNIQUE INDEX unique_c_d (c ASC, d ASC),
   FAMILY fam_0_pk_pk2_partition_by_a_b_c_d (pk, pk2, partition_by, a, b, c, d)
 ) PARTITION ALL BY LIST (partition_by) (
   PARTITION one VALUES IN ((1)),
@@ -327,11 +335,14 @@ t_partition_by_c_idx  c             false
 t_partition_by_c_idx  partition_by  false
 t_pk_key              partition_by  true
 t_pk_key              pk            false
+unique_c_d            c             false
+unique_c_d            d             false
+unique_c_d            partition_by  true
 
 statement ok
 DROP TABLE t
 
-# Tests for PARTITION ALL BY RANGE
+# Tests for PARTITION ALL BY RANGE.
 statement ok
 CREATE TABLE public.t (
   pk int PRIMARY KEY,
diff --git a/pkg/sql/alter_table.go b/pkg/sql/alter_table.go
index 9cdae7c79eef..0a520cf97dab 100644
--- a/pkg/sql/alter_table.go
+++ b/pkg/sql/alter_table.go
@@ -220,31 +220,16 @@ func (n *alterTableNode) startExec(params runParams) error {
 				if err := idx.FillColumns(d.Columns); err != nil {
 					return err
 				}
-				if d.PartitionByIndex.ContainsPartitions() {
-					var numImplicitColumns int
-					var err error
-					idx, numImplicitColumns, err = detectImplicitPartitionColumns(
-						params.EvalContext(),
-						n.tableDesc,
-						idx,
-						d.PartitionByIndex.PartitionBy,
-					)
-					if err != nil {
-						return err
-					}
-					partitioning, err := CreatePartitioning(
-						params.ctx,
-						params.p.ExecCfg().Settings,
-						params.EvalContext(),
-						n.tableDesc,
-						&idx,
-						numImplicitColumns,
-						d.PartitionByIndex.PartitionBy,
-					)
-					if err != nil {
-						return err
-					}
-					idx.Partitioning = partitioning
+
+				var err error
+				idx, err = params.p.configureIndexDescForNewIndexPartitioning(
+					params.ctx,
+					n.tableDesc,
+					idx,
+					d.PartitionByIndex,
+				)
+				if err != nil {
+					return err
 				}
 				foundIndex, err := n.tableDesc.FindIndexWithName(string(d.Name))
 				if err == nil {