dialect/entsql: add support for table checks in schema/migration

dialect/sql/schema/mysql.go

@@ -192,6 +192,7 @@ func (d *MySQL) tBuilder(t *Table) *sql.TableBuilder {
 		if opts := t.Annotation.Options; opts != "" {
 			b.Options(opts)
 		}
+		addChecks(b, t.Annotation)
 	}
 	return b
 }

dialect/sql/schema/postgres.go

@@ -304,6 +304,9 @@ func (d *Postgres) tBuilder(t *Table) *sql.TableBuilder {
 	for _, pk := range t.PrimaryKey {
 		b.PrimaryKey(pk.Name)
 	}
+	if t.Annotation != nil {
+		addChecks(b, t.Annotation)
+	}
 	return b
 }
 

dialect/sql/schema/postgres_test.go

@@ -12,6 +12,7 @@ import (
 	"testing"
 
 	"entgo.io/ent/dialect"
+	"entgo.io/ent/dialect/entsql"
 	"entgo.io/ent/dialect/sql"
 	"entgo.io/ent/schema/field"
 
@@ -65,12 +66,18 @@ func TestPostgres_Create(t *testing.T) {
 						{Name: "price", Type: field.TypeFloat64, SchemaType: map[string]string{dialect.Postgres: "numeric(5,2)"}},
 						{Name: "strings", Type: field.TypeOther, SchemaType: map[string]string{dialect.Postgres: "text[]"}, Nullable: true},
 					},
+					Annotation: &entsql.Annotation{
+						Check: "price > 0",
+						Checks: map[string]string{
+							"valid_name": "name <> ''",
+						},
+					},
 				},
 			},
 			before: func(mock pgMock) {
 				mock.start("120000")
 				mock.tableExists("users", false)
-				mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "users"("id" uuid NOT NULL DEFAULT uuid_generate_v4(), "name" varchar NULL COLLATE "he_IL", "age" bigint NOT NULL, "doc" jsonb NULL, "enums" varchar NOT NULL DEFAULT 'a', "price" numeric(5,2) NOT NULL, "strings" text[] NULL, PRIMARY KEY("id"))`)).
+				mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "users"("id" uuid NOT NULL DEFAULT uuid_generate_v4(), "name" varchar NULL COLLATE "he_IL", "age" bigint NOT NULL, "doc" jsonb NULL, "enums" varchar NOT NULL DEFAULT 'a', "price" numeric(5,2) NOT NULL, "strings" text[] NULL, PRIMARY KEY("id"), CHECK (price > 0), CONSTRAINT "valid_name" CHECK (name <> ''))`)).
 					WillReturnResult(sqlmock.NewResult(0, 1))
 				mock.ExpectCommit()
 			},

dialect/sql/schema/schema.go

@@ -7,6 +7,7 @@ package schema
 
 import (
 	"fmt"
+	"sort"
 	"strconv"
 	"strings"
 
@@ -542,3 +543,33 @@ func compare(v1, v2 int) int {
 	}
 	return 1
 }
+
+// addChecks appends the CHECK clauses from the entsql.Annotation.
+func addChecks(t *sql.TableBuilder, ant *entsql.Annotation) {
+	if check := ant.Check; check != "" {
+		t.Checks(func(b *sql.Builder) {
+			b.WriteString("CHECK " + checkExpr(check))
+		})
+	}
+	if checks := ant.Checks; len(ant.Checks) > 0 {
+		names := make([]string, 0, len(checks))
+		for name := range checks {
+			names = append(names, name)
+		}
+		sort.Strings(names)
+		for _, name := range names {
+			t.Checks(func(b *sql.Builder) {
+				b.WriteString("CONSTRAINT ").Ident(name).WriteString(" CHECK " + checkExpr(checks[name]))
+			})
+		}
+	}
+}
+
+// checkExpr formats the CHECK expression.
+func checkExpr(expr string) string {
+	expr = strings.TrimSpace(expr)
+	if !strings.HasPrefix(expr, "(") && !strings.HasSuffix(expr, ")") {
+		expr = "(" + expr + ")"
+	}
+	return expr
+}

dialect/sql/schema/sqlite.go

@@ -72,6 +72,9 @@ func (d *SQLite) tBuilder(t *Table) *sql.TableBuilder {
 	for _, c := range t.Columns {
 		b.Column(d.addColumn(c))
 	}
+	if t.Annotation != nil {
+		addChecks(b, t.Annotation)
+	}
 	// Unlike in MySQL, we're not able to add foreign-key constraints to table
 	// after it was created, and adding them to the `CREATE TABLE` statement is
 	// not always valid (because circular foreign-keys situation is possible).