Merge a500e4b into 4a16cb3

palewire · Mar 27, 2018 · a947727 · a947727
2 parents 4a16cb3 + a500e4b
commit a947727
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 5 deletions.
diff --git a/postgres_copy/copy_to.py b/postgres_copy/copy_to.py
@@ -49,12 +49,17 @@ def execute_sql(self, csv_path=None):
             # compile the SELECT query
             select_sql = self.as_sql()[0] % adapted_params
             # then the COPY TO query
-            copy_to_sql = "COPY ({}) TO STDOUT DELIMITER '{}' CSV {} {}"
+            copy_to_sql = \
+                "COPY ({0}) TO STDOUT DELIMITER '{1}' CSV {2} {3} {4} {5} {6} {7}"
             copy_to_sql = copy_to_sql.format(
-                select_sql,
-                self.query.copy_to_delimiter,
-                self.query.copy_to_header,
-                self.query.copy_to_null_string
+                select_sql,                      # 0
+                self.query.copy_to_delimiter,    # 1
+                self.query.copy_to_header,       # 2
+                self.query.copy_to_null_string,  # 3
+                self.query.copy_to_quote_char,   # 4
+                self.query.copy_to_force_quote,  # 5
+                self.query.copy_to_encoding,     # 6
+                self.query.copy_to_escape,       # 7
             )
             # then execute
             logger.debug(copy_to_sql)

diff --git a/postgres_copy/managers.py b/postgres_copy/managers.py
@@ -172,6 +172,29 @@ def to_csv(self, csv_path=None, *fields, **kwargs):
         null_string = kwargs.get('null', None)
         query.copy_to_null_string = "" if null_string is None else "NULL '{}'".format(null_string)
 
+        # Quote character
+        quote_char = kwargs.get('quote', None)
+        query.copy_to_quote_char = "QUOTE '{}'".format(quote_char) if quote_char else ""
+
+        # Force quote on columns
+        force_quote = kwargs.get('force_quote', None)
+        if force_quote:
+            if type(force_quote) == list:
+                query.copy_to_force_quote = \
+                    "FORCE QUOTE {}".format(", ".join(column for column in force_quote))
+            else:
+                query.copy_to_force_quote = "FORCE QUOTE {}".format(force_quote)
+        else:
+            query.copy_to_force_quote = ""
+
+        # Encoding
+        set_encoding = kwargs.get('encoding', None)
+        query.copy_to_encoding = "ENCODING '{}'".format(set_encoding) if set_encoding else ""
+
+        # Escape character
+        escape_char = kwargs.get('escape', None)
+        query.copy_to_escape = "ESCAPE '{}'".format(escape_char) if escape_char else ""
+
         # Run the query
         compiler = query.get_compiler(self.db, connection=connection)
         data = compiler.execute_sql(csv_path)