From d739d6d82329c62db2f70e85e1e4d664dace3499 Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <7413593+a8m@users.noreply.github.com>
Date: Mon, 30 Dec 2024 22:47:17 +0200
Subject: [PATCH] sql/internal/sqlx: compare check constraint names (#3289)

---
 .../testdata/postgres/table-checks.txt        | 25 +++++++++-
 sql/internal/sqlx/diff.go                     | 46 +++++++++++++++++--
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/internal/integration/testdata/postgres/table-checks.txt b/internal/integration/testdata/postgres/table-checks.txt
index e38f3e0bfc6..b37024c31c0 100644
--- a/internal/integration/testdata/postgres/table-checks.txt
+++ b/internal/integration/testdata/postgres/table-checks.txt
@@ -15,6 +15,11 @@ cmpmig 1 migration.v2.sql
 atlas migrate diff v2-check --to file://schema.v2.sql --dev-url URL --format '{{ sql . "  " }}'
 stdout 'The migration directory is synced with the desired state, no changes to be made'
 
+atlas migrate diff v3 --to file://schema.v3.sql --dev-url URL --format '{{ sql . "  " }}'
+cmpmig 2 migration.v3.sql
+atlas migrate diff v3-check --to file://schema.v3.sql --dev-url URL --format '{{ sql . "  " }}'
+stdout 'The migration directory is synced with the desired state, no changes to be made'
+
 -- schema.sql --
 create table t1 (
   a int constraint c1 check (a > 0),
@@ -118,4 +123,22 @@ create table t2 (
 -- Modify "t1" table
 ALTER TABLE "t1" DROP CONSTRAINT "c1", ADD CONSTRAINT "c1" CHECK (a > 1), DROP CONSTRAINT "c2", ADD CONSTRAINT "c2" CHECK (b > 1);
 -- Modify "t2" table
-ALTER TABLE "t2" DROP CONSTRAINT "c1", ADD CONSTRAINT "c1" CHECK (a > 1), DROP CONSTRAINT "c4", ADD CONSTRAINT "c4" CHECK (c > 1);
\ No newline at end of file
+ALTER TABLE "t2" DROP CONSTRAINT "c1", ADD CONSTRAINT "c1" CHECK (a > 1), DROP CONSTRAINT "c4", ADD CONSTRAINT "c4" CHECK (c > 1);
+-- schema.v3.sql --
+create table t1 (
+  a int constraint c1 check (a > 1),
+  b int constraint c2 check (b > 1),
+  -- Rename constraint.
+  constraint c4 check (a < b)
+);
+create table t2 (
+  a int constraint c1 check (a > 1),
+  c int constraint c4 check (c > 1),
+  -- Rename constraint.
+  constraint c6 check (a < c)
+);
+-- migration.v3.sql --
+-- Modify "t1" table
+ALTER TABLE "t1" DROP CONSTRAINT "c3", ADD CONSTRAINT "c4" CHECK (a < b);
+-- Modify "t2" table
+ALTER TABLE "t2" DROP CONSTRAINT "c5", ADD CONSTRAINT "c6" CHECK (a < c);
\ No newline at end of file
diff --git a/sql/internal/sqlx/diff.go b/sql/internal/sqlx/diff.go
index 34a3a4eda3a..b200fd83d6d 100644
--- a/sql/internal/sqlx/diff.go
+++ b/sql/internal/sqlx/diff.go
@@ -7,6 +7,7 @@ package sqlx
 import (
 	"fmt"
 	"reflect"
+	"slices"
 	"sort"
 	"strconv"
 	"strings"
@@ -810,9 +811,9 @@ func CommentDiff(from, to []schema.Attr) schema.Change {
 // if the schema.DiffMode is equal to schema.DiffModeNormalized.
 func CheckDiffMode(from, to *schema.Table, mode schema.DiffMode, compare ...func(c1, c2 *schema.Check) bool) []schema.Change {
 	if !mode.Is(schema.DiffModeNormalized) {
-		return CheckDiff(from, to, compare...)
+		return checksSimilarDiff(from, to, compare...)
 	}
-	return CheckDiff(from, to, func(c1, c2 *schema.Check) bool {
+	return ChecksDiff(from, to, func(c1, c2 *schema.Check) bool {
 		if len(compare) == 1 && !compare[0](c1, c2) {
 			return false
 		}
@@ -820,9 +821,44 @@ func CheckDiffMode(from, to *schema.Table, mode schema.DiffMode, compare ...func
 	})
 }
 
-// CheckDiff computes the change diff between the 2 tables. A compare
-// function is provided to check if a Check object was modified.
-func CheckDiff(from, to *schema.Table, compare ...func(c1, c2 *schema.Check) bool) []schema.Change {
+// ChecksDiff computes the change diff between the 2 tables.
+func ChecksDiff(from, to *schema.Table, compare func(c1, c2 *schema.Check) bool) []schema.Change {
+	var (
+		changes    []schema.Change
+		fromC, toC = checks(from.Attrs), checks(to.Attrs)
+	)
+	for _, c1 := range fromC {
+		idx := slices.IndexFunc(toC, func(c2 *schema.Check) bool {
+			return c1.Name == c2.Name
+		})
+		if idx == -1 {
+			changes = append(changes, &schema.DropCheck{
+				C: c1,
+			})
+		} else if c2 := toC[idx]; !compare(c1, c2) {
+			changes = append(changes, &schema.ModifyCheck{
+				From: c1,
+				To:   c2,
+			})
+		}
+	}
+	for _, c1 := range toC {
+		if !slices.ContainsFunc(fromC, func(c2 *schema.Check) bool {
+			return c1.Name == c2.Name
+		}) {
+			changes = append(changes, &schema.AddCheck{
+				C: c1,
+			})
+		}
+	}
+	return changes
+}
+
+// checksSimilarDiff computes the change diff between the 2 tables.
+// Unlike ChecksDiff, it does not compare the constraint name, but
+// determines if there is any similar constraint by its expression.
+// This is an old implementation that is not used anymore by the CLI.
+func checksSimilarDiff(from, to *schema.Table, compare ...func(c1, c2 *schema.Check) bool) []schema.Change {
 	var changes []schema.Change
 	// Drop or modify checks.
 	for _, c1 := range checks(from.Attrs) {