Merge 671c807 into f5c19b4

setup.py

@@ -30,9 +30,15 @@ def run(self):
                 'sqlite': {
                     'NAME': 'sqlite',
                     'ENGINE': 'django.db.backends.sqlite3'
+                },
+                'secondary': {
+                    'NAME': 'test_secondary',
+                    'USER': 'postgres',
+                    'ENGINE': 'django.db.backends.postgresql_psycopg2'
                 }
             },
             INSTALLED_APPS=("tests",),
+            DATABASE_ROUTERS=['tests.router.CustomRouter']
         )
         django.setup()
         call_command('test', 'tests.tests')

tests/data/secondary_db.csv

@@ -0,0 +1,4 @@
+TEXT
+SECONDARY TEXT 1
+SECONDARY TEXT 2
+SECONDARY TEXT 3

tests/models.py

@@ -90,3 +90,8 @@ def pre_insert(self, cursor):
 
     def post_insert(self, cursor):
         self.ran_post_insert = True
+
+
+class SecondaryMockObject(models.Model):
+    text = models.CharField(max_length=500)
+    objects = CopyManager()

tests/router.py

@@ -0,0 +1,10 @@
+class CustomRouter(object):
+    def db_for_read(self, model, **hints):
+        if model.__name__ == 'SecondaryMockObject':
+            return 'secondary'
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model.__name__ == 'SecondaryMockObject':
+            return 'secondary'
+        return None

tests/tests.py

@@ -7,7 +7,8 @@
     ExtendedMockObject,
     LimitedMockObject,
     OverloadMockObject,
-    HookedCopyMapping
+    HookedCopyMapping,
+    SecondaryMockObject,
 )
 from django.db.models import Count
 from postgres_copy import CopyMapping
@@ -25,12 +26,14 @@ def setUp(self):
         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')
+        self.secondarydb_path = os.path.join(self.data_dir, 'secondary_db.csv')
 
     def tearDown(self):
         MockObject.objects.all().delete()
         ExtendedMockObject.objects.all().delete()
         LimitedMockObject.objects.all().delete()
         OverloadMockObject.objects.all().delete()
+        SecondaryMockObject.objects.all().delete()
 
 
 class PostgresCopyToTest(BaseTest):
@@ -47,6 +50,9 @@ def tearDown(self):
     def _load_objects(self, file_path, mapping=dict(name='NAME', number='NUMBER', dt='DATE')):
         MockObject.objects.from_csv(file_path, mapping)
 
+    def _load_secondary_objects(self, file_path, mapping=dict(text='TEXT')):
+        SecondaryMockObject.objects.from_csv(file_path, mapping)
+
     def test_export(self):
         self._load_objects(self.name_path)
         MockObject.objects.to_csv(self.export_path)
@@ -99,6 +105,28 @@ def test_extra(self):
         for row in reader:
             self.assertTrue('lower' in row)
 
+    def test_export_multi_db(self):
+        self._load_objects(self.name_path)
+        self._load_secondary_objects(self.secondarydb_path)
+
+        MockObject.objects.to_csv(self.export_path)
+        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]
+        )
+
+        SecondaryMockObject.objects.to_csv(self.export_path)
+        self.assertTrue(os.path.exists(self.export_path))
+        reader = csv.DictReader(open(self.export_path, 'r'))
+        items = [i['text'] for i in reader]
+        self.assertEqual(len(items), 3)
+        self.assertEqual(
+            ['SECONDARY TEXT 1', 'SECONDARY TEXT 2', 'SECONDARY TEXT 3'],
+            items
+        )
+
 
 class PostgresCopyTest(BaseTest):