Merge pull request #2997 from dolthub/angela/escape_comments

Properly escape special characters in comments

Author: angelamayxie

enginetest/queries/alter_table_queries.go

@@ -1029,6 +1029,23 @@ var AlterTableScripts = []ScriptTest{
 			},
 		},
 	},
+	{
+		Name: "alter table comments are escaped",
+		SetUpScript: []string{
+			"create table t (i int);",
+			`alter table t modify column i int comment "newline \n | return \r | backslash \\ | NUL \0 \x00"`,
+			`alter table t add column j int comment "newline \n | return \r | backslash \\ | NUL \0 \x00"`,
+		},
+		Assertions: []ScriptTestAssertion{
+			{
+				Query: "show create table t",
+				Expected: []sql.Row{{
+					"t",
+					"CREATE TABLE `t` (\n  `i` int COMMENT 'newline \\n | return \\r | backslash \\\\ | NUL \\0 x00'," +
+						"\n  `j` int COMMENT 'newline \\n | return \\r | backslash \\\\ | NUL \\0 x00'\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+			},
+		},
+	},
 }
 
 var RenameTableScripts = []ScriptTest{

enginetest/queries/create_table_queries.go

@@ -40,6 +40,42 @@ var CreateTableQueries = []WriteQueryTest{
 		SelectQuery:         "SHOW CREATE TABLE tableWithComment",
 		ExpectedSelect:      []sql.Row{{"tableWithComment", "CREATE TABLE `tableWithComment` (\n  `pk` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin COMMENT='~!@ #$ %^ &* ()'"}},
 	},
+	{
+		WriteQuery:          `create table tableWithComment (pk int) COMMENT "'"`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithComment",
+		ExpectedSelect:      []sql.Row{{"tableWithComment", "CREATE TABLE `tableWithComment` (\n  `pk` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin COMMENT=''''"}},
+	},
+	{
+		WriteQuery:          `create table tableWithComment (pk int) COMMENT "\'"`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithComment",
+		ExpectedSelect:      []sql.Row{{"tableWithComment", "CREATE TABLE `tableWithComment` (\n  `pk` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin COMMENT=''''"}},
+	},
+	{
+		WriteQuery:          `create table tableWithComment (pk int) COMMENT "newline \n | return \r | backslash \\ | NUL \0 \x00"`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithComment",
+		ExpectedSelect:      []sql.Row{{"tableWithComment", "CREATE TABLE `tableWithComment` (\n  `pk` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin COMMENT='newline \\n | return \\r | backslash \\\\ | NUL \\0 x00'"}},
+	},
+	{
+		WriteQuery:          `create table tableWithColumnComment (pk int COMMENT "'")`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithColumnComment",
+		ExpectedSelect:      []sql.Row{{"tableWithColumnComment", "CREATE TABLE `tableWithColumnComment` (\n  `pk` int COMMENT ''''\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+	},
+	{
+		WriteQuery:          `create table tableWithColumnComment (pk int COMMENT "\'")`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithColumnComment",
+		ExpectedSelect:      []sql.Row{{"tableWithColumnComment", "CREATE TABLE `tableWithColumnComment` (\n  `pk` int COMMENT ''''\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+	},
+	{
+		WriteQuery:          `create table tableWithColumnComment (pk int COMMENT "newline \n | return \r | backslash \\ | NUL \0 \x00")`,
+		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
+		SelectQuery:         "SHOW CREATE TABLE tableWithColumnComment",
+		ExpectedSelect:      []sql.Row{{"tableWithColumnComment", "CREATE TABLE `tableWithColumnComment` (\n  `pk` int COMMENT 'newline \\n | return \\r | backslash \\\\ | NUL \\0 x00'\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin"}},
+	},
 	{
 		WriteQuery:          `create table floattypedefs (a float(10), b float(10, 2), c double(10, 2))`,
 		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
@@ -1118,18 +1154,6 @@ var CreateTableAutoIncrementTests = []ScriptTest{
 }
 
 var BrokenCreateTableQueries = []WriteQueryTest{
-	{
-		// TODO: We don't support table comments that contain single quotes due to how table options are parsed.
-		//       Vitess parses them, but turns any double quotes into single quotes and puts all table options back
-		//       into a string that GMS has to reparse. This means we can't tell if the single quote is the end of
-		//       the quoted string, or if it was a single quote inside of double quotes and needs to be escaped.
-		//       To fix this, Vitess should return the parsed table options as structured data, instead of as a
-		//       single string.
-		WriteQuery:          `create table tableWithComment (pk int) COMMENT "'"`,
-		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},
-		SelectQuery:         "SHOW CREATE TABLE tableWithComment",
-		ExpectedSelect:      []sql.Row{{"tableWithComment", "CREATE TABLE `tableWithComment` (\n  `pk` int\n) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_bin COMMENT=''''"}},
-	},
 	{
 		WriteQuery:          `create table t1 (b blob, primary key(b(1)))`,
 		ExpectedWriteResult: []sql.Row{{types.NewOkResult(0)}},

sql/parser.go

@@ -133,6 +133,17 @@ func RemoveSpaceAndDelimiter(query string, d rune) string {
 	})
 }
 
+func EscapeSpecialCharactersInComment(comment string) string {
+	commentString := comment
+	commentString = strings.ReplaceAll(commentString, "'", "''")
+	commentString = strings.ReplaceAll(commentString, "\\", "\\\\")
+	commentString = strings.ReplaceAll(commentString, "\"", "\\\"")
+	commentString = strings.ReplaceAll(commentString, "\n", "\\n")
+	commentString = strings.ReplaceAll(commentString, "\r", "\\r")
+	commentString = strings.ReplaceAll(commentString, "\x00", "\\0")
+	return commentString
+}
+
 type MySqlSchemaFormatter struct{}
 
 var _ SchemaFormatter = &MySqlSchemaFormatter{}
@@ -141,8 +152,7 @@ var _ SchemaFormatter = &MySqlSchemaFormatter{}
 func (m *MySqlSchemaFormatter) GenerateCreateTableStatement(tblName string, colStmts []string, temp, autoInc, tblCharsetName, tblCollName, comment string) string {
 	if comment != "" {
 		// Escape any single quotes in the comment and add the COMMENT keyword
-		comment = strings.ReplaceAll(comment, "'", "''")
-		comment = fmt.Sprintf(" COMMENT='%s'", comment)
+		comment = fmt.Sprintf(" COMMENT='%s'", EscapeSpecialCharactersInComment(comment))
 	}
 
 	if autoInc != "" {
@@ -201,7 +211,7 @@ func (m *MySqlSchemaFormatter) GenerateCreateTableColumnDefinition(col *Column,
 	}
 
 	if col.Comment != "" {
-		stmt = fmt.Sprintf("%s COMMENT '%s'", stmt, col.Comment)
+		stmt = fmt.Sprintf("%s COMMENT '%s'", stmt, EscapeSpecialCharactersInComment(col.Comment))
 	}
 	return stmt
 }