Merge branch 'igncampa-missing_kwargs'

docs/index.rst

@@ -133,7 +133,8 @@ How to export data
     class Command(BaseCommand):
 
         def handle(self, *args, **kwargs):
-            # All this method needs is the path to your CSV
+            # All this method needs is the path to your CSV.
+            # (If you don't provide one, the method will return the CSV as a string.)
             Person.objects.to_csv('/path/to/my/export.csv')
 
 Run your exporter and that's it.
@@ -460,21 +461,34 @@ The ``to_csv`` manager method only requires one argument, the path to where the
 =================  =========================================================
 Argument           Description
 =================  =========================================================
-``csv_path``       The path to the delimited data source file
-                   (e.g., a CSV)
+``csv_path``       The path to a file to write out the CSV. Optional.
+                   If you don't provide one, the comma-delimited data is
+                   returned as a string.
 
-``fields``         Strings corresponding to
-                   the model fields to be exported. All fields on the model
-                   are exported by default. Fields on related models can be
-                   included with Django's double underscore notation.
+``fields``         Strings corresponding to the model fields to be exported.
+                   All fields on the model are exported by default. Fields
+                   on related models can be included with Django's double
+                   underscore notation. Optional.
 
 ``delimiter``      String that will be used as a delimiter for the CSV
-                   file.
+                   file. Optional.
 
-``header``         Boolean determines if the header should be exported
+``header``         Boolean determines if the header should be exported.
+                   Optional.
 
 ``null``           String to populate exported null values with. Default
-                   is an empty string.
+                   is an empty string. Optional.
+
+``encoding``       The character encoding that should be used for the file
+                   being written. Optional.
+
+``escape``         The escape character to be used. Optional.
+
+``quote``          The quote character to be used. Optional.
+
+``force_quote``    Force fields to be quoted in the CSV. Default is None.
+                   A field name or list of field names can be submitted.
+                   Pass in True or "*" to quote all fields. Optional.
 =================  =========================================================
 
 Reducing the exported fields

postgres_copy/__init__.py

@@ -3,7 +3,7 @@
 from .copy_from import CopyMapping
 from .copy_to import SQLCopyToCompiler, CopyToQuery
 from .managers import CopyManager, CopyQuerySet
-__version__ = '2.3.0'
+__version__ = '2.3.1'
 
 
 __all__ = (

postgres_copy/copy_from.py

@@ -402,7 +402,7 @@ def drop(self, cursor):
         cursor:
           A cursor object on the db
         """
-        logger.debug("Running INSERT command")
+        logger.debug("Running DROP command")
         drop_sql = self.prep_drop()
         logger.debug(drop_sql)
         cursor.execute(drop_sql)

postgres_copy/copy_to.py

@@ -5,6 +5,7 @@
 """
 from __future__ import unicode_literals
 import logging
+from io import BytesIO
 from django.db import connections
 from psycopg2.extensions import adapt
 from django.db.models.sql.query import Query
@@ -26,13 +27,14 @@ def setup_query(self):
             for field in self.query.copy_to_fields:
                 # raises error if field is not available
                 expression = self.query.resolve_ref(field)
-                if field in self.query.annotations:
-                    selection = (expression, self.compile(expression), field)
-                else:
-                    selection = (expression, self.compile(expression), None)
+                selection = (
+                    expression,
+                    self.compile(expression),
+                    field if field in self.query.annotations else None,
+                )
                 self.select.append(selection)
 
-    def execute_sql(self, csv_path):
+    def execute_sql(self, csv_path=None):
         """
         Run the COPY TO query.
         """
@@ -42,23 +44,38 @@ def execute_sql(self, csv_path):
         params = self.as_sql()[1]
         adapted_params = tuple(adapt(p) for p in params)
 
-        # open file for writing
         # use stdout to avoid file permission issues
-        with open(csv_path, 'wb') as stdout:
-            with connections[self.using].cursor() as c:
-                # 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_to_sql.format(
-                    select_sql,
-                    self.query.copy_to_delimiter,
-                    self.query.copy_to_header,
-                    self.query.copy_to_null_string
-                )
-                # then execute
-                logger.debug(copy_to_sql)
+        with connections[self.using].cursor() as c:
+            # compile the SELECT query
+            select_sql = self.as_sql()[0] % adapted_params
+            # then the COPY TO query
+            copy_to_sql = "COPY ({}) TO STDOUT {} CSV"
+            copy_to_sql = copy_to_sql.format(select_sql, self.query.copy_to_delimiter)
+            # Optional extras
+            options_list = [
+                self.query.copy_to_header,
+                self.query.copy_to_null_string,
+                self.query.copy_to_quote_char,
+                self.query.copy_to_force_quote,
+                self.query.copy_to_encoding,
+                self.query.copy_to_escape
+            ]
+            options_sql = " ".join([o for o in options_list if o]).strip()
+            if options_sql:
+                copy_to_sql = copy_to_sql + " " + options_sql
+            # then execute
+            logger.debug(copy_to_sql)
+
+            # If there is, write it out there.
+            if csv_path:
+                with open(csv_path, 'wb') as stdout:
+                    c.cursor.copy_expert(copy_to_sql, stdout)
+                    return
+            # If there's no csv_path, return the output as a string.
+            else:
+                stdout = BytesIO()
                 c.cursor.copy_expert(copy_to_sql, stdout)
+                return stdout.getvalue()
 
 
 class CopyToQuery(Query):

postgres_copy/managers.py

@@ -147,7 +147,7 @@ def from_csv(self, csv_path, mapping=None, drop_constraints=True, drop_indexes=T
 
         return insert_count
 
-    def to_csv(self, csv_path, *fields, **kwargs):
+    def to_csv(self, csv_path=None, *fields, **kwargs):
         """
         Copy current QuerySet to CSV at provided path.
         """
@@ -162,19 +162,52 @@ def to_csv(self, csv_path, *fields, **kwargs):
         query.copy_to_fields = fields
 
         # Delimiter
-        query.copy_to_delimiter = kwargs.get('delimiter', ',')
+        query.copy_to_delimiter = "DELIMITER '{}'".format(kwargs.get('delimiter', ','))
 
-        # Header?
+        # Header
         with_header = kwargs.get('header', True)
         query.copy_to_header = "HEADER" if with_header else ""
 
         # Null string
         null_string = kwargs.get('null', None)
-        query.copy_to_null_string = "" if null_string is None else "NULL '{}'".format(null_string)
+        query.copy_to_null_string = "NULL '{}'".format(null_string) if null_string else ""
+
+        # 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 it's a list of fields, pass them in with commas
+            if type(force_quote) == list:
+                query.copy_to_force_quote = \
+                    "FORCE QUOTE {}".format(", ".join(column for column in force_quote))
+            # If it's True or a * force quote everything
+            elif force_quote is True or force_quote == "*":
+                query.copy_to_force_quote = "FORCE QUOTE *"
+            # Otherwise, assume it's a string and pass it through
+            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)
-        compiler.execute_sql(csv_path)
+        data = compiler.execute_sql(csv_path)
+
+        # If no csv_path is provided, then the query will come back as a string.
+        if csv_path is None:
+            # So return that.
+            return data
 
 
 CopyManager = models.Manager.from_queryset(CopyQuerySet)

setup.py

@@ -72,7 +72,7 @@ def run(self):
 
 setup(
     name='django-postgres-copy',
-    version='2.3.0',
+    version='2.3.1',
     description="Quickly move comma-delimited data in and out of a Django model using PostgreSQL's COPY command",
     author='Ben Welsh',
     author_email='ben.welsh@gmail.com',

tests/tests.py

@@ -70,6 +70,15 @@ def test_export(self):
             [i['name'] for i in reader]
         )
 
+    def test_export_to_str(self):
+        self._load_objects(self.name_path)
+        export = MockObject.objects.to_csv()
+        self.assertEqual(export, b"""id,name,num,dt,parent_id
+83,BEN,1,2012-01-01,
+84,JOE,2,2012-01-02,
+85,JANE,3,2012-01-03,
+""")
+
     def test_export_header_setting(self):
         self._load_objects(self.name_path)
         MockObject.objects.to_csv(self.export_path)
@@ -121,6 +130,59 @@ def test_export_null_string(self):
             [i['num'] for i in reader]
         )
 
+    def test_export_quote_character_and_force_quoting(self):
+        self._load_objects(self.name_path)
+
+        # Single column being force_quoted with pipes
+        MockObject.objects.to_csv(self.export_path, quote='|', force_quote='NAME')
+        self.assertTrue(os.path.exists(self.export_path))
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        self.assertTrue(
+            ['|BEN|', '|JOE|', '|JANE|'],
+            [i['name'] for i in reader]
+        )
+
+        # Multiple columns passed as a list and force_quoted with pipes
+        MockObject.objects.to_csv(self.export_path, quote='|', force_quote=['NAME', 'DT'])
+        self.assertTrue(os.path.exists(self.export_path))
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        self.assertTrue(
+            [('|BEN|', '|2012-01-01|'), ('|JOE|', '|2012-01-02|'), ('|JANE|', '|2012-01-03|')],
+            [(i['name'], i['dt']) for i in reader]
+        )
+
+        # All columns force_quoted with pipes
+        MockObject.objects.to_csv(self.export_path, quote='|', force_quote=True)
+        self.assertTrue(os.path.exists(self.export_path))
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        reader = next(reader)
+        self.assertTrue(
+            ['|BEN|', '|1|', '|2012-01-01|'],
+            list(reader.values())[1:]
+        )
+
+    def test_export_encoding(self):
+        self._load_objects(self.name_path)
+
+        # Function should pass on valid inputs ('utf-8', 'Unicode', 'LATIN2')
+        # If these don't raise an error, then they passed nicely
+        MockObject.objects.to_csv(self.export_path, encoding='utf-8')
+        MockObject.objects.to_csv(self.export_path, encoding='Unicode')
+        MockObject.objects.to_csv(self.export_path, encoding='LATIN2')
+
+        # Function should fail on known invalid inputs ('ASCII', 'utf-16')
+        self.assertRaises(Exception, MockObject.objects.to_csv(self.export_path), encoding='utf-16')
+        self.assertRaises(Exception, MockObject.objects.to_csv(self.export_path), encoding='ASCII')
+
+    def test_export_escape_character(self):
+        self._load_objects(self.name_path)
+
+        # Function should not fail on known valid inputs
+        MockObject.objects.to_csv(self.export_path, escape='-')
+
+        # Function should fail on known invalid inputs
+        self.assertRaises(Exception, MockObject.objects.to_csv(self.export_path), escape='--')
+
     def test_filter(self):
         self._load_objects(self.name_path)
         MockObject.objects.filter(name="BEN").to_csv(self.export_path)