Repository: ariga/atlas
Pull request #2691 (2024-04-17 08:57:48)

sql/postgres: bind checks with their columns


Patch:
diff --git a/sql/postgres/inspect.go b/sql/postgres/inspect.go
index 58ad8580..09b49164 100644
--- a/sql/postgres/inspect.go
+++ b/sql/postgres/inspect.go
@@ -701,20 +701,22 @@ func (i *inspect) addChecks(s *schema.Schema, rows *sql.Rows) error {
 		if !ok {
 			return fmt.Errorf("table %q was not found in schema", table)
 		}
-		if _, ok := t.Column(column); !ok {
+		c, ok := t.Column(column)
+		if !ok {
 			return fmt.Errorf("postgres: column %q was not found for check %q", column, name)
 		}
-		check, ok := names[tc{t: table, n: name}]
+		ck, ok := names[tc{t: table, n: name}]
 		if !ok {
-			check = &schema.Check{Name: name, Expr: clause, Attrs: []schema.Attr{&CheckColumns{}}}
+			ck = &schema.Check{Name: name, Expr: clause, Attrs: []schema.Attr{&CheckColumns{}}}
 			if noInherit {
-				check.Attrs = append(check.Attrs, &NoInherit{})
+				ck.AddAttrs(&NoInherit{})
 			}
-			names[tc{t: table, n: name}] = check
-			t.Attrs = append(t.Attrs, check)
+			names[tc{t: table, n: name}] = ck
+			t.AddAttrs(ck)
 		}
-		c := check.Attrs[0].(*CheckColumns)
-		c.Columns = append(c.Columns, column)
+		c.AddAttrs(ck)
+		attr := ck.Attrs[0].(*CheckColumns)
+		attr.Columns = append(attr.Columns, column)
 	}
 	return nil
 }
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 14335c44..57723e0e 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -970,7 +970,7 @@ func (s *state) column(b *sqlx.Builder, c *schema.Column) error {
 	s.columnDefault(b, c)
 	for _, attr := range c.Attrs {
 		switch a := attr.(type) {
-		case *schema.Comment:
+		case *schema.Comment, *schema.Check:
 		case *schema.Collation:
 			b.P("COLLATE").Ident(a.V)
 		case *Identity, *schema.GeneratedExpr:


Test patch:
diff --git a/sql/postgres/inspect_test.go b/sql/postgres/inspect_test.go
index 7c38f9d5..c9994783 100644
--- a/sql/postgres/inspect_test.go
+++ b/sql/postgres/inspect_test.go
@@ -316,18 +316,19 @@ users        | users_check1       | (((c2 + c1) + c3) > 10) | c3          | {2,1
 				require.NoError(err)
 				require.Equal("users", t.Name)
 				require.Equal("public", t.Schema.Name)
-				require.EqualValues([]*schema.Column{
-					{Name: "c1", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}},
-					{Name: "c2", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}},
-					{Name: "c3", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}},
-				}, t.Columns)
-				require.EqualValues([]schema.Attr{
+				checks := []schema.Attr{
 					&schema.Check{Name: "boring", Expr: "(c1 > 1)", Attrs: []schema.Attr{&CheckColumns{Columns: []string{"c1"}}, &NoInherit{}}},
 					&schema.Check{Name: "users_c2_check", Expr: "(c2 > 0)", Attrs: []schema.Attr{&CheckColumns{Columns: []string{"c2"}}}},
 					&schema.Check{Name: "users_c2_check1", Expr: "(c2 > 0)", Attrs: []schema.Attr{&CheckColumns{Columns: []string{"c2"}}}},
 					&schema.Check{Name: "users_check", Expr: "((c2 + c1) > 2)", Attrs: []schema.Attr{&CheckColumns{Columns: []string{"c2", "c1"}}}},
 					&schema.Check{Name: "users_check1", Expr: "(((c2 + c1) + c3) > 10)", Attrs: []schema.Attr{&CheckColumns{Columns: []string{"c2", "c1", "c3"}}}},
-				}, t.Attrs)
+				}
+				require.EqualValues([]*schema.Column{
+					{Name: "c1", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{checks[0], checks[3], checks[4]}},
+					{Name: "c2", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{checks[1], checks[2], checks[3], checks[4]}},
+					{Name: "c3", Type: &schema.ColumnType{Raw: "integer", Type: &schema.IntegerType{T: "integer"}}, Attrs: []schema.Attr{checks[4]}},
+				}, t.Columns)
+				require.EqualValues(checks, t.Attrs)
 			},
 		},
 	}