Handle escape characters correctly

roksela · roksela · commit 155f197e46a6 · 2017-11-21T17:30:13.000+01:00
diff --git a/postgresql_csv_loader/csv_loader.py b/postgresql_csv_loader/csv_loader.py
@@ -15,8 +15,9 @@ class CsvLoader(object):
 
     DEFAULT_DELIMITER = ','
     DEFAULT_QUOTE_CHAR = '"'
-    DEFAULT_ESCAPE_CHAR = '"'
+    DEFAULT_ESCAPE_CHAR = None
     DEFAULT_TABLE_PREFIX = "csv_"
+    DEFAULT_DOUBLE_QUOTE = True
     DEFAULT_DATA_TYPE = "varchar"
 
     CREATE_STMT = "CREATE TABLE {} ({});"
@@ -59,10 +60,13 @@ def load_data(self, file_path, delimiter=DEFAULT_DELIMITER, quote_char=DEFAULT_Q
         :param create_table: if True, table will be created
         :param encoding file encoding
         """
+        # doublequote=True by default
+        # don't define escape char if it's the same as quote char
+        escape_char = None if (escape_char == quote_char) else escape_char
+
         original_headers = self._read_headers(file_path, delimiter, quote_char, escape_char, encoding)
         headers = self._normalize_headers(original_headers)
         table_name = self._generate_table_name(file_path)
-
         logging.getLogger('CsvLoader').info('Connecting to database "{}"...'.format(self._database_name))
         connection = connect(dbname=self._database_name, user=self._user, password= self._password,
                              host=self._database_host, port=self._database_port)
@@ -148,8 +152,10 @@ def _copy_from_csv(self, connection, file_path, table_name, headers, delimiter,
         """
         columns = ['"{}"'.format(column) for column in headers]
         columns_def = ",".join(columns)
+
+        copy_from_escape_char = escape_char or quote_char  # use quote if escape is None
         command = self.COPY_STMT.format(table_name, columns_def, delimiter,
-                                        quote_char, escape_char)
+                                        quote_char, copy_from_escape_char)
         # https://www.postgresql.org/docs/current/static/sql-copy.html
 
         cursor = connection.cursor()
diff --git a/test/resources/quoted_headers.csv b/test/resources/quoted_headers.csv
@@ -0,0 +1,6 @@
+"Respondent","Professional","Country"
+"1","Student","United States ""America"""
+2,Student,United Kingdom
+3,Professional developer,United Kingdom
+4,Professional developer,United States
+5,Professional developer,Switzerland
diff --git a/test/test_csv_loader.py b/test/test_csv_loader.py
@@ -19,16 +19,19 @@ class TestCsvLoader(unittest.TestCase):
     CSV_FILENAME_3 = "resources/illegal_column_names.csv"
     CSV_FILENAME_4 = "resources/polish_characters.csv"
     CSV_FILENAME_5 = "resources/weird_format.csv"
+    CSV_FILENAME_6 = "resources/quoted_headers.csv"
     CSV_1_RECORD_COUNT = 30
     CSV_2_RECORD_COUNT = 5
     CSV_3_RECORD_COUNT = 5
     CSV_4_RECORD_COUNT = 1
     CSV_5_RECORD_COUNT = 1
+    CSV_6_RECORD_COUNT = 5
     TABLE_NAME_1 = "csv_stackoverflow_survey_results_public_sample"
     TABLE_NAME_2 = "csv_simple_table"
     TABLE_NAME_3 = "csv_illegal_column_names"
     TABLE_NAME_4 = "csv_polish_characters"
     TABLE_NAME_5 = "csv_weird_format"
+    TABLE_NAME_6 = "csv_quoted_headers"
 
     SELECT_COUNT_STMT = "SELECT count(*) from {};"
     DROP_STMT = "DROP TABLE {};"
@@ -135,6 +138,15 @@ def test_load_weird_format(self):
         self._drop(self.TABLE_NAME_5)
         self.assertEqual(result, self.CSV_5_RECORD_COUNT)
 
+    def test_headers_quotes(self):
+        loader = self._get_loader()
+        loader.load_data(self.CSV_FILENAME_6)
+
+        result = self._check_count(self.TABLE_NAME_6)
+        self._drop(self.TABLE_NAME_6)
+
+        self.assertEqual(result, self.CSV_6_RECORD_COUNT)
+
     def _get_loader(self):
         return CsvLoader(self.database_host, self.database_port, self.database_name, self.database_user)
 
