Repository: ariga/atlas
Pull request #630 (2022-03-06 20:28:34)

sql: use smarter approach to ensure expressions are wrapped


Patch:
diff --git a/sql/internal/sqlx/sqlx.go b/sql/internal/sqlx/sqlx.go
index 657bb8a2..d49f1ab6 100644
--- a/sql/internal/sqlx/sqlx.go
+++ b/sql/internal/sqlx/sqlx.go
@@ -341,3 +341,40 @@ func DefaultValue(c *schema.Column) (string, bool) {
 		panic(fmt.Sprintf("unexpected default value type: %T", x))
 	}
 }
+
+// MayWrap ensures the given string is wrapped with parentheses.
+// Used by the different drivers to turn strings valid expressions.
+func MayWrap(s string) string {
+	n := len(s) - 1
+	if len(s) < 2 || s[0] != '(' || s[n] != ')' || !balanced(s[1:n]) {
+		return "(" + s + ")"
+	}
+	return s
+}
+
+func balanced(s string) bool {
+	var d int
+	for i := 0; i < len(s) && d >= 0; i++ {
+	Top:
+		switch s[i] {
+		case '(':
+			d++
+		case ')':
+			d--
+		// String or identifier.
+		case '\'', '"', '`':
+			for j := i + 1; j < len(s); j++ {
+				switch s[j] {
+				case '\\':
+					j++
+				case s[i]:
+					i = j
+					break Top
+				}
+			}
+			// Unexpected EOS.
+			return false
+		}
+	}
+	return d == 0
+}
diff --git a/sql/mysql/migrate.go b/sql/mysql/migrate.go
index 4720a77e..85d71fca 100644
--- a/sql/mysql/migrate.go
+++ b/sql/mysql/migrate.go
@@ -515,7 +515,7 @@ func indexParts(b *sqlx.Builder, parts []*schema.IndexPart) {
 			case part.C != nil:
 				b.Ident(part.C.Name)
 			case part.X != nil:
-				b.WriteString(withParens(part.X.(*schema.RawExpr).X))
+				b.WriteString(sqlx.MayWrap(part.X.(*schema.RawExpr).X))
 			}
 			if s := (&SubPart{}); sqlx.Has(parts[i].Attrs, s) {
 				b.WriteString(fmt.Sprintf("(%d)", s.Len))
@@ -688,7 +688,7 @@ func (s *state) check(b *sqlx.Builder, c *schema.Check) {
 	if c.Name != "" {
 		b.P("CONSTRAINT").Ident(c.Name)
 	}
-	b.P("CHECK", withParens(c.Expr))
+	b.P("CHECK", sqlx.MayWrap(c.Expr))
 	if s.supportsEnforceCheck() && sqlx.Has(c.Attrs, &Enforced{}) {
 		b.P("ENFORCED")
 	}
@@ -711,11 +711,3 @@ func quote(s string) string {
 	}
 	return strconv.Quote(s)
 }
-
-// withParens ensures expressions wrapped with parens.
-func withParens(x string) string {
-	if t := strings.TrimSpace(x); !strings.HasPrefix(t, "(") || !strings.HasSuffix(t, ")") {
-		x = "(" + t + ")"
-	}
-	return x
-}
diff --git a/sql/postgres/migrate.go b/sql/postgres/migrate.go
index 4c5bbed0..0436a2a7 100644
--- a/sql/postgres/migrate.go
+++ b/sql/postgres/migrate.go
@@ -613,7 +613,7 @@ func (s *state) indexParts(b *sqlx.Builder, parts []*schema.IndexPart) {
 			case part.C != nil:
 				b.Ident(part.C.Name)
 			case part.X != nil:
-				b.WriteString(withParens(part.X.(*schema.RawExpr).X))
+				b.WriteString(sqlx.MayWrap(part.X.(*schema.RawExpr).X))
 			}
 			s.partAttrs(b, parts[i])
 		})
@@ -767,7 +767,7 @@ func check(b *sqlx.Builder, c *schema.Check) {
 	if c.Name != "" {
 		b.P("CONSTRAINT").Ident(c.Name)
 	}
-	b.P("CHECK", withParens(c.Expr))
+	b.P("CHECK", sqlx.MayWrap(c.Expr))
 	if sqlx.Has(c.Attrs, &NoInherit{}) {
 		b.P("NO INHERIT")
 	}
@@ -779,11 +779,3 @@ func quote(s string) string {
 	}
 	return "'" + strings.ReplaceAll(s, "'", "''") + "'"
 }
-
-// withParens ensures expressions wrapped with parens.
-func withParens(x string) string {
-	if t := strings.TrimSpace(x); !strings.HasPrefix(t, "(") || !strings.HasSuffix(t, ")") {
-		x = "(" + t + ")"
-	}
-	return x
-}


Test patch:
diff --git a/sql/internal/sqlx/sqlx_test.go b/sql/internal/sqlx/sqlx_test.go
index 7d2d8a7a..0b7922d6 100644
--- a/sql/internal/sqlx/sqlx_test.go
+++ b/sql/internal/sqlx/sqlx_test.go
@@ -5,6 +5,7 @@
 package sqlx
 
 import (
+	"strconv"
 	"testing"
 
 	"ariga.io/atlas/sql/schema"
@@ -39,3 +40,38 @@ func TestBuilder(t *testing.T) {
 		})
 	require.Equal(t, `CREATE TABLE "users" ("a" int NOT NULL, "b" int NOT NULL, "c" int NOT NULL, PRIMARY KEY ("a", "b", "c"))`, b.String())
 }
+
+func TestMayWrap(t *testing.T) {
+	tests := []struct {
+		input   string
+		wrapped bool
+	}{
+		{"", true},
+		{"()", false},
+		{"('text')", false},
+		{"('(')", false},
+		{`('(\\')`, false},
+		{`('\')(')`, false},
+		{`(a) in (b)`, true},
+		{`a in (b)`, true},
+		{`("\\\\(((('")`, false},
+		{`('(')||(')')`, true},
+		// Test examples from SQLite.
+		{"b || 'x'", true},
+		{"a+1", true},
+		{"substr(x, 2)", true},
+		{"(json_extract(x, '$.a'), json_extract(x, '$.b'))", false},
+		{"(substr(a, 2) COLLATE NOCASE, b)", false},
+		{"(a,b+random())", false},
+	}
+	for i, tt := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			expect := tt.input
+			if tt.wrapped {
+				expect = "(" + expect + ")"
+			}
+			require.Equal(t, expect, MayWrap(tt.input))
+
+		})
+	}
+}
diff --git a/sql/postgres/migrate_test.go b/sql/postgres/migrate_test.go
index 346ef00e..87d23ea0 100644
--- a/sql/postgres/migrate_test.go
+++ b/sql/postgres/migrate_test.go
@@ -165,6 +165,7 @@ func TestPlanChanges(t *testing.T) {
 							&schema.Comment{},
 							&schema.Check{Name: "id_nonzero", Expr: `("id" > 0)`},
 							&schema.Check{Name: "text_len", Expr: `(length("text") > 0)`, Attrs: []schema.Attr{&NoInherit{}}},
+							&schema.Check{Name: "a_in_b", Expr: `(a) in (b)`},
 						},
 					},
 				},
@@ -172,7 +173,7 @@ func TestPlanChanges(t *testing.T) {
 			plan: &migrate.Plan{
 				Reversible:    true,
 				Transactional: true,
-				Changes:       []*migrate.Change{{Cmd: `CREATE TABLE "posts" ("id" integer NOT NULL GENERATED BY DEFAULT AS IDENTITY, "text" text NULL, CONSTRAINT "id_nonzero" CHECK ("id" > 0), CONSTRAINT "text_len" CHECK (length("text") > 0) NO INHERIT)`, Reverse: `DROP TABLE "posts"`}},
+				Changes:       []*migrate.Change{{Cmd: `CREATE TABLE "posts" ("id" integer NOT NULL GENERATED BY DEFAULT AS IDENTITY, "text" text NULL, CONSTRAINT "id_nonzero" CHECK ("id" > 0), CONSTRAINT "text_len" CHECK (length("text") > 0) NO INHERIT, CONSTRAINT "a_in_b" CHECK ((a) in (b)))`, Reverse: `DROP TABLE "posts"`}},
 			},
 		},
 		{