From 9d634790b7eb67eae5fe465584e020a2cef1a71f Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Wed, 17 Jan 2024 11:37:44 +0000
Subject: [PATCH 1/2] Add `CHECK` constraints to internal schema

Add knowledge of `CHECK` constraints to `pgroll`'s internal schema
representation.
---
 pkg/schema/schema.go | 14 ++++++++++++++
 pkg/state/state.go   | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/pkg/schema/schema.go b/pkg/schema/schema.go
index f88ae265..b882f0be 100644
--- a/pkg/schema/schema.go
+++ b/pkg/schema/schema.go
@@ -49,6 +49,9 @@ type Table struct {
 
 	// ForeignKeys is a map of all foreign keys defined on the table
 	ForeignKeys map[string]ForeignKey `json:"foreignKeys"`
+
+	// CheckConstraints is a map of all check constraints defined on the table
+	CheckConstraints map[string]CheckConstraint `json:"checkConstraints"`
 }
 
 type Column struct {
@@ -85,6 +88,17 @@ type ForeignKey struct {
 	ReferencedColumns []string `json:"referencedColumns"`
 }
 
+type CheckConstraint struct {
+	// Name is the name of the check constraint in postgres
+	Name string `json:"name"`
+
+	// The columns that the check constraint is defined on
+	Columns []string `json:"columns"`
+
+	// The definition of the check constraint
+	Definition string `json:"definition"`
+}
+
 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 882b4ec7..f05c8c8d 100644
--- a/pkg/state/state.go
+++ b/pkg/state/state.go
@@ -173,6 +173,24 @@ BEGIN
 				  FROM pg_index pi 
 				  WHERE pi.indrelid = t.oid::regclass
 				),
+				'checkConstraints', (
+					SELECT json_object_agg(cc_details.conname, json_build_object(
+						'name', cc_details.conname,
+						'columns', cc_details.columns,
+						'definition', cc_details.definition
+					))
+					FROM (
+						SELECT
+							cc_constraint.conname,
+							array_agg(cc_attr.attname ORDER BY cc_constraint.conkey::int[]) AS columns,
+							pg_get_constraintdef(cc_constraint.oid) AS definition
+						FROM pg_constraint AS cc_constraint
+						INNER JOIN pg_attribute cc_attr ON cc_attr.attrelid = cc_constraint.conrelid AND cc_attr.attnum = ANY(cc_constraint.conkey)
+						WHERE cc_constraint.conrelid = t.oid
+						AND cc_constraint.contype = 'c'
+						GROUP BY cc_constraint.oid
+					) AS cc_details
+        ),
 				'foreignKeys', (
 					SELECT json_object_agg(fk_details.conname, json_build_object(
 						'name', fk_details.conname,

From 7e9e4ec8ad52244175334b0aaa1191036debabbb Mon Sep 17 00:00:00 2001
From: Andrew Farries <andrew.farries@xata.io>
Date: Wed, 17 Jan 2024 13:24:23 +0000
Subject: [PATCH 2/2] Add test for db schema with a `CHECK` constraint

---
 pkg/state/state_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go
index a5decd43..043f824b 100644
--- a/pkg/state/state_test.go
+++ b/pkg/state/state_test.go
@@ -155,6 +155,44 @@ func TestReadSchema(t *testing.T) {
 					},
 				},
 			},
+			{
+				name:       "check constraint",
+				createStmt: "CREATE TABLE public.table1 (id int PRIMARY KEY, age INTEGER, CONSTRAINT age_check CHECK (age > 18));",
+				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,
+								},
+								"age": {
+									Name:     "age",
+									Type:     "integer",
+									Nullable: true,
+								},
+							},
+							PrimaryKey: []string{"id"},
+							Indexes: map[string]schema.Index{
+								"table1_pkey": {
+									Name: "table1_pkey",
+								},
+							},
+							CheckConstraints: map[string]schema.CheckConstraint{
+								"age_check": {
+									Name:       "age_check",
+									Columns:    []string{"age"},
+									Definition: "CHECK ((age > 18))",
+								},
+							},
+						},
+					},
+				},
+			},
 		}
 
 		// init the state