From 994a1aa6f8142cd2078ac12f625728ce7474159f Mon Sep 17 00:00:00 2001
From: Andrew Farries <andyrb@gmail.com>
Date: Thu, 18 Jan 2024 12:08:53 +0000
Subject: [PATCH] Add `UNIQUE` constraints to `pgroll`'s internal schema
 representation (#242)

Add knowledge of `UNIQUE` constraints defined on a table to `pgroll`'s
internal schema representation.

`UNIQUE` constraints were already present as indexes, but this PR adds
them as their own top-level field in the `Table` schema and includes
information about the columns involved in the constraint.

Part of https://github.com/xataio/pgroll/issues/105
---
 pkg/schema/schema.go    | 11 ++++++++++
 pkg/state/state.go      | 17 +++++++++++++++
 pkg/state/state_test.go | 47 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
index b882f0be..50de06b5 100644
--- a/pkg/schema/schema.go
+++ b/pkg/schema/schema.go
@@ -52,6 +52,9 @@ type Table struct {
 
 	// CheckConstraints is a map of all check constraints defined on the table
 	CheckConstraints map[string]CheckConstraint `json:"checkConstraints"`
+
+	// UniqueConstraints is a map of all unique constraints defined on the table
+	UniqueConstraints map[string]UniqueConstraint `json:"uniqueConstraints"`
 }
 
 type Column struct {
@@ -99,6 +102,14 @@ type CheckConstraint struct {
 	Definition string `json:"definition"`
 }
 
+type UniqueConstraint struct {
+	// Name is the name of the unique constraint in postgres
+	Name string `json:"name"`
+
+	// The columns that the unique constraint is defined on
+	Columns []string `json:"columns"`
+}
+
 func (s *Schema) GetTable(name string) *Table {
 	if s.Tables == nil {
 		return nil
diff --git a/pkg/state/state.go b/pkg/state/state.go
index f05c8c8d..41159229 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -190,6 +190,23 @@ BEGIN
 						AND cc_constraint.contype = 'c'
 						GROUP BY cc_constraint.oid
 					) AS cc_details
+        ),
+				'uniqueConstraints', (
+					SELECT json_object_agg(uc_details.conname, json_build_object(
+						'name', uc_details.conname,
+						'columns', uc_details.columns
+					))
+					FROM (
+						SELECT
+							uc_constraint.conname,
+							array_agg(uc_attr.attname ORDER BY uc_constraint.conkey::int[]) AS columns,
+							pg_get_constraintdef(uc_constraint.oid) AS definition
+						FROM pg_constraint AS uc_constraint
+						INNER JOIN pg_attribute uc_attr ON uc_attr.attrelid = uc_constraint.conrelid AND uc_attr.attnum = ANY(uc_constraint.conkey)
+						WHERE uc_constraint.conrelid = t.oid
+						AND uc_constraint.contype = 'u'
+						GROUP BY uc_constraint.oid
+					) AS uc_details
         ),
 				'foreignKeys', (
 					SELECT json_object_agg(fk_details.conname, json_build_object(
diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go
index 043f824b..c2bbe93f 100644
--- a/pkg/state/state_test.go
+++ b/pkg/state/state_test.go
@@ -107,6 +107,12 @@ func TestReadSchema(t *testing.T) {
 									Name: "id_unique",
 								},
 							},
+							UniqueConstraints: map[string]schema.UniqueConstraint{
+								"id_unique": {
+									Name:    "id_unique",
+									Columns: []string{"id"},
+								},
+							},
 						},
 					},
 				},
@@ -193,6 +199,47 @@ func TestReadSchema(t *testing.T) {
 					},
 				},
 			},
+			{
+				name:       "unique constraint",
+				createStmt: "CREATE TABLE public.table1 (id int PRIMARY KEY, name TEXT, CONSTRAINT name_unique UNIQUE(name) );",
+				wantSchema: &schema.Schema{
+					Name: "public",
+					Tables: map[string]schema.Table{
+						"table1": {
+							Name: "table1",
+							Columns: map[string]schema.Column{
+								"id": {
+									Name:     "id",
+									Type:     "integer",
+									Nullable: false,
+									Unique:   true,
+								},
+								"name": {
+									Name:     "name",
+									Type:     "text",
+									Unique:   true,
+									Nullable: true,
+								},
+							},
+							PrimaryKey: []string{"id"},
+							Indexes: map[string]schema.Index{
+								"table1_pkey": {
+									Name: "table1_pkey",
+								},
+								"name_unique": {
+									Name: "name_unique",
+								},
+							},
+							UniqueConstraints: map[string]schema.UniqueConstraint{
+								"name_unique": {
+									Name:    "name_unique",
+									Columns: []string{"name"},
+								},
+							},
+						},
+					},
+				},
+			},
 		}
 
 		// init the state