Merge pull request #1432 from fyfe/fix/1430

fix import command

modoboa/admin/management/commands/subcommands/_export.py

@@ -6,14 +6,18 @@ from __future__ import unicode_literals
 
 import sys
 
-from backports import csv
-
 from django.core.management.base import BaseCommand
+from django.utils import six
 
 from modoboa.core.extensions import exts_pool
 from modoboa.core.models import User
 from .... import models
 
+if six.PY2:
+    from backports import csv
+else:
+    import csv
+
 
 class ExportCommand(BaseCommand):
     """Command class."""

modoboa/admin/management/commands/subcommands/_import.py

@@ -8,15 +8,23 @@ import io
 import os
 
 import progressbar
-from backports import csv
+from chardet.universaldetector import UniversalDetector
 
 from django.core.management.base import BaseCommand, CommandError
+from django.db import transaction
+from django.utils import six
+from django.utils.translation import ugettext as _
 
 from modoboa.core import models as core_models
 from modoboa.core.extensions import exts_pool
 from modoboa.lib.exceptions import Conflict
 from .... import signals
 
+if six.PY2:
+    from backports import csv
+else:
+    import csv
+
 
 class ImportCommand(BaseCommand):
     """Command class."""
@@ -27,7 +35,7 @@ class ImportCommand(BaseCommand):
     def add_arguments(self, parser):
         """Add extra arguments to command."""
         parser.add_argument(
-            "--sepchar", type=str, default=";",
+            "--sepchar", type=six.text_type, default=";",
             help="Separator used in file.")
         parser.add_argument(
             "--continue-if-exists", action="store_true",
@@ -37,23 +45,25 @@ class ImportCommand(BaseCommand):
             "--crypt-password", action="store_true",
             default=False, help="Encrypt provided passwords.")
         parser.add_argument(
-            "files", type=str, nargs="+", help="CSV files to import.")
+            "files", type=six.text_type, nargs="+", help="CSV files to import.")
 
-    def _import(self, filename, options):
+    def _import(self, filename, options, encoding="utf-8"):
         """Import domains or identities."""
         superadmin = (
             core_models.User.objects.filter(is_superuser=True).first()
         )
-        if not os.path.exists(filename):
+        if not os.path.isfile(filename):
             raise CommandError("File not found")
 
-        num_lines = sum(1 for line in open(filename) if line)
+        num_lines = sum(
+            1 for line in io.open(filename, encoding=encoding) if line
+        )
         pbar = progressbar.ProgressBar(
             widgets=[
                 progressbar.Percentage(), progressbar.Bar(), progressbar.ETA()
             ], maxval=num_lines
         ).start()
-        with io.open(filename, encoding="utf8") as f:
+        with io.open(filename, encoding=encoding, newline="") as f:
             reader = csv.reader(f, delimiter=options["sepchar"])
             i = 0
             for row in reader:
@@ -82,4 +92,39 @@ class ImportCommand(BaseCommand):
         """Command entry point."""
         exts_pool.load_all()
         for filename in options["files"]:
-            self._import(filename, options)
+            try:
+                with transaction.atomic():
+                    self._import(filename, options)
+            except CommandError as exc:
+                raise exc
+            except UnicodeDecodeError:
+                self.stdout.write(self.style.NOTICE(
+                    _("CSV file is not encoded in UTF-8, attempting to guess "
+                      "encoding")
+                ))
+                detector = UniversalDetector()
+                with io.open(filename, "rb") as fp:
+                    for line in fp:
+                        detector.feed(line)
+                        if detector.done:
+                            break
+                    detector.close()
+
+                self.stdout.write(self.style.NOTICE(
+                    _("Reading CSV file using %(encoding)s encoding") %
+                    detector.result
+                ))
+                try:
+                    with transaction.atomic():
+                        self._import(
+                            filename, options,
+                            encoding=detector.result["encoding"]
+                        )
+                except UnicodeDecodeError as exc:
+                    six.raise_from(
+                        CommandError(
+                            _("Unable to decode CSV file using %(encoding)s "
+                              "encoding") % detector.result
+                        ),
+                        exc
+                    )

modoboa/admin/tests/test_data/import_domains.csv

@@ -0,0 +1,2 @@
+domain;domain1.com;1000;100;True
+domain;dómain2.com;2000;200;True

modoboa/admin/tests/test_data/import_domains_duplicates.csv

@@ -0,0 +1,3 @@
+domain;domain1.com;1000;100;True
+domain;dómain2.com;2000;200;True
+domain;dómain2.com;2000;200;True

modoboa/admin/tests/test_import_.py

@@ -2,9 +2,13 @@
 
 from __future__ import unicode_literals
 
+import os
+
 import dns.resolver
 
 from django.core.files.base import ContentFile
+from django.core.management import call_command
+from django.core.management.base import CommandError
 from django.urls import reverse
 
 from modoboa.core import factories as core_factories
@@ -353,3 +357,63 @@ dómain; dómªin2.com; 1000; 100; True
         self.assertEqual(dom.default_mailbox_quota, 100)
         self.assertTrue(dom.enabled)
         self.assertTrue(admin.is_owner(dom))
+
+    def test_import_command_missing_file(self):
+        with mock.patch("os.path.isfile") as mock_isfile:
+            mock_isfile.return_value = False
+            with self.assertRaises(CommandError) as cm:
+                call_command("modo", "import", "test.csv")
+                ex_message = cm.exception.messages
+                self.assertEqual(ex_message, "File not found")
+
+    def test_import_command(self):
+        test_file = os.path.join(
+            os.path.dirname(__file__),
+            "test_data/import_domains.csv"
+        )
+        call_command("modo", "import", test_file)
+
+        admin = User.objects.get(username="admin")
+
+        dom = Domain.objects.get(name="domain1.com")
+        self.assertEqual(dom.quota, 1000)
+        self.assertEqual(dom.default_mailbox_quota, 100)
+        self.assertTrue(dom.enabled)
+        self.assertTrue(admin.is_owner(dom))
+
+        dom = Domain.objects.get(name="dómain2.com")
+        self.assertEqual(dom.quota, 2000)
+        self.assertEqual(dom.default_mailbox_quota, 200)
+        self.assertTrue(dom.enabled)
+        self.assertTrue(admin.is_owner(dom))
+
+    def test_import_command_duplicates(self):
+        test_file = os.path.join(
+            os.path.dirname(__file__),
+            "test_data/import_domains_duplicates.csv"
+        )
+        with self.assertRaises(CommandError) as cm:
+            call_command("modo", "import", test_file)
+            ex_message = cm.exception.messages
+            self.assertTrue(ex_message.startswith("Object already exists: "))
+
+    def test_import_command_duplicates_continue(self):
+        test_file = os.path.join(
+            os.path.dirname(__file__),
+            "test_data/import_domains_duplicates.csv"
+        )
+        call_command("modo", "import", "--continue-if-exists", test_file)
+
+        admin = User.objects.get(username="admin")
+
+        dom = Domain.objects.get(name="domain1.com")
+        self.assertEqual(dom.quota, 1000)
+        self.assertEqual(dom.default_mailbox_quota, 100)
+        self.assertTrue(dom.enabled)
+        self.assertTrue(admin.is_owner(dom))
+
+        dom = Domain.objects.get(name="dómain2.com")
+        self.assertEqual(dom.quota, 2000)
+        self.assertEqual(dom.default_mailbox_quota, 200)
+        self.assertTrue(dom.enabled)
+        self.assertTrue(admin.is_owner(dom))

modoboa/admin/views/export.py

@@ -4,7 +4,6 @@
 
 from __future__ import unicode_literals
 
-from backports import csv
 from rfc6266 import build_header
 
 from django.contrib.auth.decorators import (
@@ -19,6 +18,11 @@ from django.utils.translation import ugettext as _
 from ..forms import ExportDomainsForm, ExportIdentitiesForm
 from ..lib import get_domains, get_identities
 
+if six.PY2:
+    from backports import csv
+else:
+    import csv
+
 
 def _export(content, filename):
     """Export a csv file's content

modoboa/admin/views/import_.py

@@ -6,7 +6,6 @@ from __future__ import unicode_literals
 
 import io
 
-from backports import csv
 from reversion import revisions as reversion
 
 from django.contrib.auth.decorators import (
@@ -15,6 +14,7 @@ from django.contrib.auth.decorators import (
 from django.db import transaction
 from django.shortcuts import render
 from django.urls import reverse
+from django.utils import six
 from django.utils.encoding import smart_text
 from django.utils.translation import ugettext as _
 
@@ -22,6 +22,11 @@ from modoboa.lib.exceptions import Conflict, ModoboaException
 from .. import signals
 from ..forms import ImportDataForm, ImportIdentitiesForm
 
+if six.PY2:
+    from backports import csv
+else:
+    import csv
+
 
 @reversion.create_revision()
 def importdata(request, formclass=ImportDataForm):

requirements.txt

@@ -26,5 +26,5 @@ pytz
 requests
 rfc6266
 lxml
-backports.csv
+backports.csv; python_version < '3'
 chardet