From eca18b0894488c223a18e795eebb442ac825717d Mon Sep 17 00:00:00 2001
From: Ariel Mashraki <7413593+a8m@users.noreply.github.com>
Date: Sun, 20 Dec 2020 10:33:34 +0200
Subject: [PATCH] dialect/sql: add check clause for column builder (#1065)

---
 dialect/sql/builder.go      | 13 ++++++++++++-
 dialect/sql/builder_test.go |  5 ++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/dialect/sql/builder.go b/dialect/sql/builder.go
index 5f16d545c..f714e4e37 100644
--- a/dialect/sql/builder.go
+++ b/dialect/sql/builder.go
@@ -37,6 +37,7 @@ type ColumnBuilder struct {
 	attr   string             // extra attributes.
 	modify bool               // modify existing.
 	fk     *ForeignKeyBuilder // foreign-key constraint.
+	check  func(*Builder)     // column checks.
 }
 
 // Column returns a new ColumnBuilder with the given name.
@@ -66,12 +67,18 @@ func (c *ColumnBuilder) Constraint(fk *ForeignKeyBuilder) *ColumnBuilder {
 	return c
 }
 
+// Check adds a CHECK clause to the ADD COLUMN statement.
+func (c *ColumnBuilder) Check(check func(*Builder)) *ColumnBuilder {
+	c.check = check
+	return c
+}
+
 // Query returns query representation of a Column.
 func (c *ColumnBuilder) Query() (string, []interface{}) {
 	c.Ident(c.name)
 	if c.typ != "" {
 		if c.postgres() && c.modify {
-			c.Pad().WriteString("TYPE")
+			c.WriteString(" TYPE")
 		}
 		c.Pad().WriteString(c.typ)
 	}
@@ -85,6 +92,10 @@ func (c *ColumnBuilder) Query() (string, []interface{}) {
 			c.Pad().WriteString(action)
 		}
 	}
+	if c.check != nil {
+		c.WriteString(" CHECK ")
+		c.Nested(c.check)
+	}
 	return c.String(), c.args
 }
 
diff --git a/dialect/sql/builder_test.go b/dialect/sql/builder_test.go
index 84e255089..3eecc1dd5 100644
--- a/dialect/sql/builder_test.go
+++ b/dialect/sql/builder_test.go
@@ -78,11 +78,14 @@ func TestBuilder(t *testing.T) {
 				Columns(
 					Column("id").Type("int").Attr("auto_increment"),
 					Column("card_id").Type("int"),
+					Column("doc").Type("longtext").Check(func(b *Builder) {
+						b.WriteString("JSON_VALID(").Ident("doc").WriteByte(')')
+					}),
 				).
 				PrimaryKey("id", "name").
 				ForeignKeys(ForeignKey().Columns("card_id").
 					Reference(Reference().Table("cards").Columns("id")).OnDelete("SET NULL")),
-			wantQuery: "CREATE TABLE IF NOT EXISTS `users`(`id` int auto_increment, `card_id` int, PRIMARY KEY(`id`, `name`), FOREIGN KEY(`card_id`) REFERENCES `cards`(`id`) ON DELETE SET NULL)",
+			wantQuery: "CREATE TABLE IF NOT EXISTS `users`(`id` int auto_increment, `card_id` int, `doc` longtext CHECK (JSON_VALID(`doc`)), PRIMARY KEY(`id`, `name`), FOREIGN KEY(`card_id`) REFERENCES `cards`(`id`) ON DELETE SET NULL)",
 		},
 		{
 			input: Dialect(dialect.Postgres).CreateTable("users").