format: escape values enclosed by single quotes (schemalex#97)

format/format.go

@@ -92,9 +92,7 @@ func formatTableOption(ctx *fmtCtx, option model.TableOption) error {
 	buf.WriteString(option.Key())
 	buf.WriteString(" = ")
 	if option.NeedQuotes() {
-		buf.WriteByte('\'')
-		buf.WriteString(option.Value())
-		buf.WriteByte('\'')
+		buf.WriteString(util.Singlequote(option.Value()))
 	} else {
 		buf.WriteString(option.Value())
 	}
@@ -216,19 +214,15 @@ func formatTableColumn(ctx *fmtCtx, col model.TableColumn) error {
 	case model.ColumnTypeEnum:
 		buf.WriteString(" (")
 		for enumValue := range col.EnumValues() {
-			buf.WriteByte('\'')
-			buf.WriteString(enumValue)
-			buf.WriteByte('\'')
+			buf.WriteString(util.Singlequote(enumValue))
 			buf.WriteByte(',')
 		}
 		buf.Truncate(buf.Len() - 1)
 		buf.WriteByte(')')
 	case model.ColumnTypeSet:
 		buf.WriteString(" (")
 		for setValue := range col.SetValues() {
-			buf.WriteByte('\'')
-			buf.WriteString(setValue)
-			buf.WriteByte('\'')
+			buf.WriteString(util.Singlequote(setValue))
 			buf.WriteByte(',')
 		}
 		buf.Truncate(buf.Len() - 1)
@@ -286,9 +280,7 @@ func formatTableColumn(ctx *fmtCtx, col model.TableColumn) error {
 	if col.HasDefault() {
 		buf.WriteString(" DEFAULT ")
 		if col.IsQuotedDefault() {
-			buf.WriteByte('\'')
-			buf.WriteString(col.Default())
-			buf.WriteByte('\'')
+			buf.WriteString(util.Singlequote(col.Default()))
 		} else {
 			buf.WriteString(col.Default())
 		}
@@ -309,9 +301,8 @@ func formatTableColumn(ctx *fmtCtx, col model.TableColumn) error {
 	}
 
 	if col.HasComment() {
-		buf.WriteString(" COMMENT '")
-		buf.WriteString(col.Comment())
-		buf.WriteByte('\'')
+		buf.WriteString(" COMMENT ")
+		buf.WriteString(util.Singlequote(col.Comment()))
 	}
 
 	if _, err := buf.WriteTo(ctx.dst); err != nil {

internal/util/util.go

@@ -1,7 +1,32 @@
 package util
 
+import (
+	"strings"
+)
+
 // Backquote surrounds the given string in backquotes
 func Backquote(s string) string {
 	// XXX Does this require escaping
 	return "`" + s + "`"
 }
+
+// Singlequote surrounds the given string in singlequotes
+func Singlequote(s string) string {
+	b := strings.Builder{}
+	b.Grow(len(s) + 2)
+	b.WriteRune('\'')
+	r := strings.NewReader(s)
+	for {
+		c, _, err := r.ReadRune()
+		if err != nil {
+			break
+		}
+		switch c {
+		case '\'', '\\':
+			b.WriteRune('\\')
+		}
+		b.WriteRune(c)
+	}
+	b.WriteRune('\'')
+	return b.String()
+}

parser_test.go

@@ -285,6 +285,40 @@ primary key (id, c)
 		Input:  "CREATE TABLE foo (id INT(10) NOT NULL) ENGINE = InnoDB, DEFAULT CHARACTER SET = utf8mb4 \n/**/ ;",
 		Expect: "CREATE TABLE `foo` (\n`id` INT (10) NOT NULL\n) ENGINE = InnoDB, DEFAULT CHARACTER SET = utf8mb4",
 	})
+
+	parse("EscapedDefaultValues", &Spec{
+		Input: `create table hoge (
+a varchar(20) default "hoge""s",
+b varchar(20) default 'hoge\'s',
+c varchar(20) default 'hoge''s'
+);
+`,
+		Expect: "CREATE TABLE `hoge` (\n`a` VARCHAR (20) DEFAULT 'hoge\"s',\n`b` VARCHAR (20) DEFAULT 'hoge\\'s',\n`c` VARCHAR (20) DEFAULT 'hoge\\'s'\n)",
+	})
+
+	parse("EscapedTableComment", &Spec{
+		Input:  "create table hoge (id bigint unsigned not null auto_increment) ENGINE=InnoDB AUTO_INCREMENT 10 DEFAULT CHARACTER SET = utf8 COMMENT = 'hoge''s comment';",
+		Expect: "CREATE TABLE `hoge` (\n`id` BIGINT (20) UNSIGNED NOT NULL AUTO_INCREMENT\n) ENGINE = InnoDB, AUTO_INCREMENT = 10, DEFAULT CHARACTER SET = utf8, COMMENT = 'hoge\\'s comment'",
+	})
+
+	parse("EscapedColumnOptionCommentPrimaryKey1", &Spec{
+		// see https://github.com/schemalex/schemalex/pull/40
+		Input:  "CREATE TABLE `test` (\n`id` int(11) PRIMARY KEY COMMENT 'aaa\\'' NOT NULL,\nhoge int default 1 not null COMMENT 'bbb''' UNIQUE\n);",
+		Expect: "CREATE TABLE `test` (\n`id` INT (11) NOT NULL COMMENT 'aaa\\'',\n`hoge` INT (11) NOT NULL DEFAULT 1 COMMENT 'bbb\\'',\nPRIMARY KEY (`id`),\nUNIQUE INDEX `hoge` (`hoge`)\n)",
+	})
+	parse("EscapedColumnOptionCommentPrimaryKey2", &Spec{
+		// see https://github.com/schemalex/schemalex/pull/40
+		Input:  "CREATE TABLE `test` (\n`id` int(11) COMMENT 'aaa\\'' PRIMARY KEY NOT NULL,\nhoge int default 1 UNIQUE not null COMMENT 'bbb'''\n);",
+		Expect: "CREATE TABLE `test` (\n`id` INT (11) NOT NULL COMMENT 'aaa\\'',\n`hoge` INT (11) NOT NULL DEFAULT 1 COMMENT 'bbb\\'',\nPRIMARY KEY (`id`),\nUNIQUE INDEX `hoge` (`hoge`)\n)",
+	})
+	parse("EscapedEnum", &Spec{
+		Input:  "CREATE TABLE `test` (\n`status` ENUM('on\\'', 'off''') NOT NULL DEFAULT 'off'\n);",
+		Expect: "CREATE TABLE `test` (\n`status` ENUM ('on\\'','off\\'') NOT NULL DEFAULT 'off'\n)",
+	})
+	parse("EscapedSet", &Spec{
+		Input:  "CREATE TABLE `test` (\n`status` SET('foo\\'', 'bar''', 'baz') NOT NULL DEFAULT 'foo,baz'\n);",
+		Expect: "CREATE TABLE `test` (\n`status` SET ('foo\\'','bar\\'','baz') NOT NULL DEFAULT 'foo,baz'\n)",
+	})
 }
 
 func testParse(t *testing.T, spec *Spec) {