Fix INSERT statement normalization match backtick table name (#1366)

ClickHouse · Aug 5, 2024 · 06510b2 · 06510b2
1 parent 20fb591
commit 06510b2
Show file tree

Hide file tree

Showing 3 changed files with 96 additions and 2 deletions.
diff --git a/batch.go b/batch.go
@@ -25,18 +25,23 @@ import (
 	"github.com/pkg/errors"
 )
 
-var normalizeInsertQueryMatch = regexp.MustCompile(`(?i)(INSERT\s+INTO\s+([^( ]+)(?:\s*\([^()]*(?:\([^()]*\)[^()]*)*\))?)(?:\s*VALUES)?`)
+var normalizeInsertQueryMatch = regexp.MustCompile(`(?i)(INSERT\s+INTO\s+([^(]+)(?:\s*\([^()]*(?:\([^()]*\)[^()]*)*\))?)(?:\s*VALUES)?`)
+var truncateFormat = regexp.MustCompile(`\sFORMAT\s+[^\s]+`)
+var truncateValues = regexp.MustCompile(`\sVALUES\s.*$`)
 var extractInsertColumnsMatch = regexp.MustCompile(`INSERT INTO .+\s\((?P<Columns>.+)\)$`)
 
 func extractNormalizedInsertQueryAndColumns(query string) (normalizedQuery string, tableName string, columns []string, err error) {
+	query = truncateFormat.ReplaceAllString(query, "")
+	query = truncateValues.ReplaceAllString(query, "")
+
 	matches := normalizeInsertQueryMatch.FindStringSubmatch(query)
 	if len(matches) == 0 {
 		err = errors.Errorf("invalid INSERT query: %s", query)
 		return
 	}
 
 	normalizedQuery = fmt.Sprintf("%s FORMAT Native", matches[1])
-	tableName = matches[2]
+	tableName = strings.TrimSpace(matches[2])
 
 	columns = make([]string, 0)
 	matches = extractInsertColumnsMatch.FindStringSubmatch(matches[1])

diff --git a/batch_test.go b/batch_test.go
@@ -45,6 +45,27 @@ func TestExtractNormalizedInsertQueryAndColumns(t *testing.T) {
 			expectedColumns:         []string{"col1", "col2"},
 			expectedError:           false,
 		},
+		{
+			query:                   "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2`",
+			expectedNormalizedQuery: "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2` FORMAT Native",
+			expectedTableName:       "`_test_1345# $.ДБ`.`2. Таблица №2`",
+			expectedColumns:         []string{},
+			expectedError:           false,
+		},
+		{
+			query:                   "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2` (col1, col2)",
+			expectedNormalizedQuery: "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2` (col1, col2) FORMAT Native",
+			expectedTableName:       "`_test_1345# $.ДБ`.`2. Таблица №2`",
+			expectedColumns:         []string{"col1", "col2"},
+			expectedError:           false,
+		},
+		{
+			query:                   "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2` (col1, col2) VALUES (1, 2)",
+			expectedNormalizedQuery: "INSERT INTO `_test_1345# $.ДБ`.`2. Таблица №2` (col1, col2) FORMAT Native",
+			expectedTableName:       "`_test_1345# $.ДБ`.`2. Таблица №2`",
+			expectedColumns:         []string{"col1", "col2"},
+			expectedError:           false,
+		},
 		{
 			query:                   "INSERT INTO table_name (col1, col2) VALUES (1, 2) FORMAT Native",
 			expectedNormalizedQuery: "INSERT INTO table_name (col1, col2) FORMAT Native",
@@ -66,6 +87,20 @@ func TestExtractNormalizedInsertQueryAndColumns(t *testing.T) {
 			expectedColumns:         []string{},
 			expectedError:           false,
 		},
+		{
+			query:                   "INSERT INTO table_name FORMAT JSONEachRow",
+			expectedNormalizedQuery: "INSERT INTO table_name FORMAT Native",
+			expectedTableName:       "table_name",
+			expectedColumns:         []string{},
+			expectedError:           false,
+		},
+		{
+			query:                   "INSERT INTO `table_name` VALUES (1, 2)",
+			expectedNormalizedQuery: "INSERT INTO `table_name` FORMAT Native",
+			expectedTableName:       "`table_name`",
+			expectedColumns:         []string{},
+			expectedError:           false,
+		},
 		{
 			query:         "SELECT * FROM table_name",
 			expectedError: true,

diff --git a/tests/issues/1345_test.go b/tests/issues/1345_test.go
@@ -0,0 +1,54 @@
+package issues
+
+import (
+	"context"
+	"testing"
+
+	"github.com/ClickHouse/clickhouse-go/v2/tests"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIssue1345(t *testing.T) {
+	ctx := context.Background()
+
+	conn, err := tests.GetConnection("issues", nil, nil, nil)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	require.NoError(t, conn.Exec(ctx, "CREATE DATABASE IF NOT EXISTS `_test_1345#$.ДБ`"))
+	defer conn.Exec(ctx, "DROP TABLE `_test_1345#$.ДБ`")
+
+	require.NoError(t, conn.Exec(ctx, "CREATE  TABLE  IF NOT EXISTS `_test_1345#$.ДБ`.`2. Таблица №2` (i UInt64, s String) ENGINE = Memory()"))
+
+	batch, err := conn.PrepareBatch(ctx, "INSERT INTO `_test_1345#$.ДБ`.`2. Таблица №2`")
+	require.NoError(t, err)
+
+	var (
+		i = uint64(32)
+		s = "b"
+	)
+
+	err = batch.Append(i, s)
+	require.NoError(t, err)
+
+	err = batch.Send()
+	require.NoError(t, err)
+
+	rows, err := conn.Query(ctx, "SELECT * FROM `_test_1345#$.ДБ`.`2. Таблица №2`")
+	require.NoError(t, err)
+
+	require.True(t, rows.Next())
+
+	var (
+		actualInt uint64
+		actualStr string
+	)
+	err = rows.Scan(&actualInt, &actualStr)
+	require.NoError(t, err)
+
+	require.Equal(t, i, actualInt)
+	require.Equal(t, s, actualStr)
+
+	require.NoError(t, rows.Close())
+	require.NoError(t, rows.Err())
+}