Skip to content

Commit

Permalink
Fix column list parsing for multiline INSERT statements (#1373)
Browse files Browse the repository at this point in the history 
* Fix column list parsing for multiline INSERT statements

* Fixed additional bug with lowercase statements
Added few new tests to cover corner cases
  • Loading branch information
@Fiery-Fenix
Fiery-Fenix authored Aug 19, 2024
1 parent 9ef8e96 commit 79d6313
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 2 deletions.
4 changes: 2 additions & 2 deletions batch.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,9 @@ import (
)

var normalizeInsertQueryMatch = regexp.MustCompile(`(?i)(INSERT\s+INTO\s+([^(]+)(?:\s*\([^()]*(?:\([^()]*\)[^()]*)*\))?)(?:\s*VALUES)?`)
var truncateFormat = regexp.MustCompile(`\sFORMAT\s+[^\s]+`)
var truncateFormat = regexp.MustCompile(`(?i)\sFORMAT\s+[^\s]+`)
var truncateValues = regexp.MustCompile(`\sVALUES\s.*$`)
var extractInsertColumnsMatch = regexp.MustCompile(`INSERT INTO .+\s\((?P<Columns>.+)\)$`)
var extractInsertColumnsMatch = regexp.MustCompile(`(?si)INSERT INTO .+\s\((?P<Columns>.+)\)$`)

func extractNormalizedInsertQueryAndColumns(query string) (normalizedQuery string, tableName string, columns []string, err error) {
query = truncateFormat.ReplaceAllString(query, "")
Expand Down
96 changes: 96 additions & 0 deletions batch_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,83 +25,179 @@ import (

func TestExtractNormalizedInsertQueryAndColumns(t *testing.T) {
var testCases = []struct {
name string
query string
expectedNormalizedQuery string
expectedTableName string
expectedColumns []string
expectedError bool
}{
{
name: "Regular insert",
query: "INSERT INTO table_name (col1, col2) VALUES (1, 2)",
expectedNormalizedQuery: "INSERT INTO table_name (col1, col2) FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Lowercase insert",
query: "insert into table_name (col1, col2) values (1, 2)",
expectedNormalizedQuery: "insert into table_name (col1, col2) FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Insert with mixed case, multiline and format specified",
query: `INSERT INTO "db"."table_name" (
col1,
col2
) Values (
1,
2
)
format JSONEachRow`,
expectedNormalizedQuery: `INSERT INTO "db"."table_name" (
col1,
col2
) FORMAT Native`,
expectedTableName: "\"db\".\"table_name\"",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Multiline insert",
query: `INSERT INTO table_name (
col1,
col2
) VALUES (
1,
2
)`,
expectedNormalizedQuery: `INSERT INTO table_name (
col1,
col2
) FORMAT Native`,
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Multiline insert, with columns inline",
query: `INSERT INTO table_name (col1, col2) VALUES (
1,
2
)`,
expectedNormalizedQuery: `INSERT INTO table_name (col1, col2) FORMAT Native`,
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Multiline insert, with values inline",
query: `INSERT INTO table_name (
col1,
col2
) VALUES (1, 2)`,
expectedNormalizedQuery: `INSERT INTO table_name (
col1,
col2
) FORMAT Native`,
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Insert with backtick quoted database and table names",
query: "INSERT INTO `db`.`table_name` (col1, col2) VALUES (1, 2)",
expectedNormalizedQuery: "INSERT INTO `db`.`table_name` (col1, col2) FORMAT Native",
expectedTableName: "`db`.`table_name`",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Insert with double quoted database and table names",
query: "INSERT INTO \"db\".\"table_name\" (col1, col2) VALUES (1, 2)",
expectedNormalizedQuery: "INSERT INTO \"db\".\"table_name\" (col1, col2) FORMAT Native",
expectedTableName: "\"db\".\"table_name\"",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Insert with special characters in database and table names",
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,
},
{
name: "Insert with special characters in database and table names, with columns",
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,
},
{
name: "Insert with special characters in database and table names, with columns and values",
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,
},
{
name: "Insert without database name",
query: "INSERT INTO table_name (col1, col2) VALUES (1, 2) FORMAT Native",
expectedNormalizedQuery: "INSERT INTO table_name (col1, col2) FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{"col1", "col2"},
expectedError: false,
},
{
name: "Insert without columns and values",
query: "INSERT INTO table_name",
expectedNormalizedQuery: "INSERT INTO table_name FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{},
expectedError: false,
},
{
name: "Insert with format",
query: "INSERT INTO table_name FORMAT Native",
expectedNormalizedQuery: "INSERT INTO table_name FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{},
expectedError: false,
},
{
name: "Insert with lowercase format",
query: "INSERT INTO table_name format Native",
expectedNormalizedQuery: "INSERT INTO table_name FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{},
expectedError: false,
},
{
name: "Insert with JSONEachRow format",
query: "INSERT INTO table_name FORMAT JSONEachRow",
expectedNormalizedQuery: "INSERT INTO table_name FORMAT Native",
expectedTableName: "table_name",
expectedColumns: []string{},
expectedError: false,
},
{
name: "Insert with quoted table name only",
query: "INSERT INTO `table_name` VALUES (1, 2)",
expectedNormalizedQuery: "INSERT INTO `table_name` FORMAT Native",
expectedTableName: "`table_name`",
expectedColumns: []string{},
expectedError: false,
},
{
name: "Select, should produce error",
query: "SELECT * FROM table_name",
expectedError: true,
},
Expand Down

0 comments on commit 79d6313

Please sign in to comment.