🩹 Ensure the query splitter handle comments

So that when someone use anyquery run [id], the query splitter for the shell doesn't mess with comments

Author: julien040

controller/shell.go

@@ -146,7 +146,7 @@ func (p *shell) AddMiddleware(m middleware) {
 //
 // If a middleware returns false, the pipeline stops.
 // The pipeline will always run the middlewares in the order they were added
-func (p *shell) Run(rawQuery string) bool {
+func (p *shell) Run(rawQuery string, args ...interface{}) bool {
 	queries := splitMultipleQuery(rawQuery)
 
 	if len(queries) == 0 {
@@ -163,6 +163,7 @@ func (p *shell) Run(rawQuery string) bool {
 			SQLQuery: query,
 			Config:   p.Config,
 			DB:       p.DB,
+			Args:     args,
 		}
 
 		// If the query is .read, we read the file
@@ -407,7 +408,7 @@ func splitMultipleQuery(sqlQuery string) []string {
 	// isDotCommand => true if the query is a dot command or a slash command
 	// isEscaped => true if the character is escaped (in a string '' or "")
 	// mustResetWord => true if we find a query that can be appened to queries
-	var isDotCommand, isEscapedSimpleQuote, isEscapedDoubleQuote, mustResetWord bool
+	var isDotCommand, isEscapedSimpleQuote, isEscapedDoubleQuote, isMultiLineComment, isSingleLineComment, mustResetWord bool
 
 	queries := make([]string, 0)
 	// The index character of the current query
@@ -431,6 +432,25 @@ func splitMultipleQuery(sqlQuery string) []string {
 				isDotCommand = true
 			}
 
+		case '/':
+			if len(sqlQuery) > i+1 && sqlQuery[i+1] == '*' {
+				isMultiLineComment = true
+			}
+
+		case '*':
+			// If we find a star and we are in a slash command
+			// we can assume that it is the end of the slash command
+			if isMultiLineComment && len(sqlQuery) > i+1 &&
+				sqlQuery[i+1] == '/' {
+				isMultiLineComment = false
+			}
+		case '-':
+			// If we find a dash and we are at the beginning of the query
+			// we can assume that it is a single line comment
+			if i == 0 && len(sqlQuery) > i+1 && sqlQuery[i+1] == '-' {
+				isSingleLineComment = true
+			}
+
 		case '.':
 			// If we find a dot and we are at the beginning of the query
 			// we can assume that it is a dot command
@@ -454,7 +474,7 @@ func splitMultipleQuery(sqlQuery string) []string {
 			//
 			// It could be simplified as mustResetWord = condition
 			// but I prefer to keep it like this for readability
-			if !isEscapedSimpleQuote && !isEscapedDoubleQuote {
+			if !isEscapedSimpleQuote && !isEscapedDoubleQuote && !isMultiLineComment && !isSingleLineComment {
 				mustResetWord = true
 			}
 
@@ -467,17 +487,21 @@ func splitMultipleQuery(sqlQuery string) []string {
 
 			// If we find a newline, we're not in a string
 			// and it's a dot command, it means that the query is finished
-			if !isEscapedSimpleQuote && !isEscapedDoubleQuote && isDotCommand {
+			if !isEscapedSimpleQuote && !isEscapedDoubleQuote && !isMultiLineComment && !isSingleLineComment && isDotCommand {
 				mustResetWord = true
 			}
 
+			if isSingleLineComment {
+				isSingleLineComment = false
+			}
+
 		}
 
 		// We append the character to the query
 		// unless it is a newline and it is a dot command
 		//
 		// Or it's a ; and we are not in a string
-		if !((c == '\n' && isDotCommand) || (c == ';' && !isEscapedSimpleQuote && !isEscapedDoubleQuote)) {
+		if !((c == '\n' && isDotCommand) || (c == ';' && !isEscapedSimpleQuote && !isEscapedDoubleQuote && !isMultiLineComment && !isSingleLineComment)) {
 			tempQuery.WriteRune(c)
 		}
 		currentPosition++
@@ -495,6 +519,8 @@ func splitMultipleQuery(sqlQuery string) []string {
 			mustResetWord = false
 			isEscapedSimpleQuote = false
 			isEscapedDoubleQuote = false
+			isMultiLineComment = false
+			isSingleLineComment = false
 		}
 
 	}

controller/shell_test.go

@@ -64,6 +64,26 @@ func TestQuerySplitter(t *testing.T) {
 			query:    "    .mode\n\nSELECT * FROM table;    .exit  ",
 			expected: []string{".mode", "SELECT * FROM table", ".exit"},
 		},
+		{
+			name: "a query with a comment and a sql command",
+			query: `-- This is a; comment
+SELECT * FROM table`,
+			expected: []string{"-- This is a; comment\nSELECT * FROM table"},
+		},
+		{
+			name: "a query with a multi-line comment and a sql command",
+			query: `/*
+This is a multi-line comment; with a semi-colon
+Hey
+*/
+SELECT * FROM table; .exit`,
+			expected: []string{"/*\nThis is a multi-line comment; with a semi-colon\nHey\n*/\nSELECT * FROM table", ".exit"},
+		},
+		{
+			name:     "a query with a multi-line comment in the middle of a sql command",
+			query:    `SELECT *  /* This is a 'multi-line' comment; with "a" semi-colon*/ FROM table; .exit`,
+			expected: []string{"SELECT *  /* This is a 'multi-line' comment; with \"a\" semi-colon*/ FROM table", ".exit"},
+		},
 	}
 
 	for _, tt := range test {