From fec4bb427581ad0dea7b3dd86f396eae764b64c0 Mon Sep 17 00:00:00 2001
From: Nick Tobey <nick@dolthub.com>
Date: Fri, 23 Feb 2024 14:37:53 -0800
Subject: [PATCH] Allow importing tables when the columns contain characters
 that require escaping.

---
 .../doltcore/mvdata/engine_table_writer.go    |  4 +--
 .../bats/helper/escaped-characters.csv        |  2 ++
 integration-tests/bats/import-tables.bats     | 27 ++++++++++++++++++-
 3 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 integration-tests/bats/helper/escaped-characters.csv

diff --git a/go/libraries/doltcore/mvdata/engine_table_writer.go b/go/libraries/doltcore/mvdata/engine_table_writer.go
index f2f57c8450..6788f564f3 100644
--- a/go/libraries/doltcore/mvdata/engine_table_writer.go
+++ b/go/libraries/doltcore/mvdata/engine_table_writer.go
@@ -275,7 +275,7 @@ func (s *SqlEngineTableWriter) createTable() error {
 	var pks string
 	var sep string
 	for _, i := range s.tableSchema.PkOrdinals {
-		pks += sep + s.tableSchema.Schema[i].Name
+		pks += sep + sql.QuoteIdentifier(s.tableSchema.Schema[i].Name)
 		sep = ", "
 	}
 	if len(sep) > 0 {
@@ -303,7 +303,7 @@ func (s *SqlEngineTableWriter) getInsertNode(inputChannel chan sql.Row, replace
 	}
 	sep := ""
 	for _, col := range s.rowOperationSchema.Schema {
-		colNames += fmt.Sprintf("%s`%s`", sep, col.Name)
+		colNames += fmt.Sprintf("%s%s", sep, sql.QuoteIdentifier(col.Name))
 		values += fmt.Sprintf("%s1", sep)
 		if update {
 			duplicate += fmt.Sprintf("%s`%s` = VALUES(`%s`)", sep, col.Name, col.Name)
diff --git a/integration-tests/bats/helper/escaped-characters.csv b/integration-tests/bats/helper/escaped-characters.csv
new file mode 100644
index 0000000000..d160632c28
--- /dev/null
+++ b/integration-tests/bats/helper/escaped-characters.csv
@@ -0,0 +1,2 @@
+--Key1,"""Key2",'Key3,`Key4,/Key5
+0, 1, 2, 3, 4
\ No newline at end of file
diff --git a/integration-tests/bats/import-tables.bats b/integration-tests/bats/import-tables.bats
index 822010df77..7d69a66f55 100644
--- a/integration-tests/bats/import-tables.bats
+++ b/integration-tests/bats/import-tables.bats
@@ -21,4 +21,29 @@ teardown() {
     run dolt table import -c -u -r t test.csv
     [ "$status" -eq 1 ]
     [[ "$output" =~ "Must specify exactly one of -c, -u, -a, or -r." ]] || false
-}
\ No newline at end of file
+}
+
+@test "import-tables: import tables where field names need to be escaped" {
+    dolt table import -c t `batshelper escaped-characters.csv`
+
+    run dolt sql -q "show create table t;"
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ '`--Key1` int' ]] || false
+    [[ "$output" =~ '`"Key2` int' ]] || false
+    [[ "$output" =~ '`'"'"'Key3` int' ]] || false
+    [[ "$output" =~ '```Key4` int' ]] || false
+    [[ "$output" =~ '`/Key5` int' ]] || false
+
+    run dolt sql -q "select * from t;"
+    [ "$status" -eq 0 ]
+    [[ "$output" =~ '| --Key1 | "Key2 | '"'"'Key3 | `Key4 | /Key5 |' ]] || false
+    [[ "$output" =~ '| 0      | 1     | 2     | 3     | 4     |' ]] || false
+}
+
+@test "import-tables: import tables where primary key names need to be escaped" {
+    dolt table import -c -pk "--Key1" t `batshelper escaped-characters.csv`
+    dolt table import -r -pk '"Key2' t `batshelper escaped-characters.csv`
+    dolt table import -r -pk "'Key3" t `batshelper escaped-characters.csv`
+    dolt table import -r -pk '`Key4' t `batshelper escaped-characters.csv`
+    dolt table import -r -pk "/Key5" t `batshelper escaped-characters.csv`
+}