Merge pull request #45 from rdmurphy/force-null-support

FORCE NOT NULL and FORCE NULL support

docs/index.rst

@@ -117,7 +117,7 @@ Like I said, that's it!
 ``CopyMapping`` API
 -------------------
 
-.. class:: CopyMapping(model, csv_path, mapping[, using=None, delimiter=',', null=None, encoding=None, static_mapping=None])
+.. class:: CopyMapping(model, csv_path, mapping[, using=None, delimiter=',', null=None, force_not_null=None, force_null=None, encoding=None, static_mapping=None])
 
 The following are the arguments and keywords that may be used during
 instantiation of ``CopyMapping`` objects.
@@ -150,6 +150,19 @@ Keyword Arguments
                        The default is an unquoted empty string. This must
                        be a single one-byte character.
 
+``force_not_null``     Specifies which columns that should ignore matches
+                       against the null string. Empty values in these columns
+                       will remain zero-length strings rather than becoming
+                       nulls. The default is None. If passed, this must be
+                       list of column names.
+
+``force_null``         Specifies which columns that should register matches
+                       against the null string, even if it has been quoted.
+                       In the default case where the null string is empty,
+                       this converts a quoted empty string into NULL. The
+                       default is None. If passed, this must be list of
+                       column names.
+
 ``encoding``           Specifies the character set encoding of the strings
                        in the CSV data source.  For example, ``'latin-1'``,
                        ``'utf-8'``, and ``'cp437'`` are all valid encoding

postgres_copy/__init__.py

@@ -24,6 +24,8 @@ def __init__(
         delimiter=',',
         quote_character=None,
         null=None,
+        force_not_null=None,
+        force_null=None,
         encoding=None,
         static_mapping=None
     ):
@@ -38,6 +40,8 @@ def __init__(
         self.quote_character = quote_character
         self.delimiter = delimiter
         self.null = null
+        self.force_not_null = force_not_null
+        self.force_null = force_null
         self.encoding = encoding
         if static_mapping is not None:
             self.static_mapping = OrderedDict(static_mapping)
@@ -202,6 +206,10 @@ def prep_copy(self):
             options['extra_options'] += " DELIMITER '%s'" % self.delimiter
         if self.null is not None:
             options['extra_options'] += " NULL '%s'" % self.null
+        if self.force_not_null is not None:
+            options['extra_options'] += " FORCE NOT NULL %s" % ','.join('"%s"' % s for s in self.force_not_null)
+        if self.force_null is not None:
+            options['extra_options'] += " FORCE NULL %s" % ','.join('"%s"' % s for s in self.force_null)
         if self.encoding:
             options['extra_options'] += " ENCODING '%s'" % self.encoding
         return sql % options

tests/data/blanknulls.csv

@@ -0,0 +1,6 @@
+NAME,NUMBER,DATE,COLOR
+ben,1,2012-01-01,red
+joe,2,2012-01-02,green
+jane,3,2012-01-03,orange
+nullboy,,2012-01-04,
+badboy,x,2012-01-05,blue

tests/models.py

@@ -16,6 +16,20 @@ def copy_name_template(self):
         return 'upper("%(name)s")'
 
 
+class MockBlankObject(models.Model):
+    name = models.CharField(max_length=500)
+    number = MyIntegerField(null=True, db_column='num')
+    dt = models.DateField(null=True)
+    color = models.CharField(max_length=50, blank=True)
+    parent = models.ForeignKey('MockObject', null=True, default=None)
+
+    class Meta:
+        app_label = 'tests'
+
+    def copy_name_template(self):
+        return 'upper("%(name)s")'
+
+
 class ExtendedMockObject(models.Model):
     static_val = models.IntegerField()
     name = models.CharField(max_length=500)

tests/tests.py

@@ -2,6 +2,7 @@
 from datetime import date
 from .models import (
     MockObject,
+    MockBlankObject,
     ExtendedMockObject,
     LimitedMockObject,
     OverloadMockObject,
@@ -19,6 +20,7 @@ def setUp(self):
         self.foreign_path = os.path.join(self.data_dir, 'foreignkeys.csv')
         self.pipe_path = os.path.join(self.data_dir, 'pipes.csv')
         self.quote_path = os.path.join(self.data_dir, 'quote.csv')
+        self.blank_null_path = os.path.join(self.data_dir, 'blanknulls.csv')
         self.null_path = os.path.join(self.data_dir, 'nulls.csv')
         self.backwards_path = os.path.join(self.data_dir, 'backwards.csv')
 
@@ -189,6 +191,38 @@ def test_null_save(self):
             date(2012, 1, 1)
         )
 
+    def test_force_not_null_save(self):
+        c = CopyMapping(
+            MockBlankObject,
+            self.blank_null_path,
+            dict(name='NAME', number='NUMBER', dt='DATE', color='COLOR'),
+            force_not_null=('COLOR',),
+        )
+        c.save()
+        self.assertEqual(MockBlankObject.objects.count(), 5)
+        self.assertEqual(MockBlankObject.objects.get(name='BEN').color, 'red')
+        self.assertEqual(MockBlankObject.objects.get(name='NULLBOY').color, '')
+        self.assertEqual(
+            MockBlankObject.objects.get(name='BEN').dt,
+            date(2012, 1, 1)
+        )
+
+    def test_force_null_save(self):
+        c = CopyMapping(
+            MockObject,
+            self.null_path,
+            dict(name='NAME', number='NUMBER', dt='DATE'),
+            force_null=('NUMBER',),
+        )
+        c.save()
+        self.assertEqual(MockObject.objects.count(), 5)
+        self.assertEqual(MockObject.objects.get(name='BEN').number, 1)
+        self.assertEqual(MockObject.objects.get(name='NULLBOY').number, None)
+        self.assertEqual(
+            MockObject.objects.get(name='BEN').dt,
+            date(2012, 1, 1)
+        )
+
     def test_backwards_save(self):
         c = CopyMapping(
             MockObject,