Imports: manage model limitation (don't create items not in the list)

1 files changed, 84 insertions, 9 deletions

diff --git a/ishtar_common/data_importer.py b/ishtar_common/data_importer.py
index de4883c69..79259b76d 100644
--- a/ishtar_common/data_importer.py
+++ b/ishtar_common/data_importer.py
@@ -29,6 +29,7 @@ import zipfile
 
 from django.conf import settings
 from django.contrib.auth.models import User
+from django.core.exceptions import ImproperlyConfigured
 from django.core.files import File
 from django.db import IntegrityError, DatabaseError, transaction
 from django.template.defaultfilters import slugify
@@ -613,6 +614,8 @@ class Importer(object):
     OBJECT_CLS = None
     IMPORTED_LINE_FIELD = None
     UNICITY_KEYS = []
+    # if set only models inside this list can be created
+    MODEL_CREATION_LIMIT = []
     EXTRA_DEFAULTS = {}
     DEFAULTS = {}
     ERRORS = {
@@ -626,10 +629,19 @@ class Importer(object):
         'no_data': _(u"No data provided"),
         'value_required': _(u"Value is required"),
         'not_enough_cols': _(u"At least %d columns must be filled"),
-        'regex_not_match': _(u"The regexp doesn't match.")
+        'regex_not_match': _(u"The regexp doesn't match."),
+        'improperly_configured': _(
+            u"Force creation is set for model {} but this model is not in the "
+            u"list of model allowed to be created."),
+        'does_not_exist_in_db': _(u"{} with values {} doesn't exist in the "
+          u"database. Create it first or fix your source file"),
     }
 
     def _create_models(self, force=False):
+        """
+        Create a db config from a hardcoded import.
+        Not useful anymore?
+        """
         from ishtar_common import models
         q = models.ImporterType.objects.filter(slug=self.SLUG)
         if not force and (not self.SLUG or q.count()):
@@ -1009,11 +1021,32 @@ class Importer(object):
                     if k not in formater.through_unicity_keys \
                        and k != 'defaults':
                         data['defaults'][k] = data.pop(k)
+            created = False
             if '__force_new' in data:
+                if self.MODEL_CREATION_LIMIT and \
+                                through_cls not in self.MODEL_CREATION_LIMIT:
+                    raise ImproperlyConfigured(
+                        unicode(self.ERRORS[ 'improperly_configured']).format(
+                            through_cls))
                 created = data.pop('__force_new')
                 t_obj = through_cls.objects.create(**data)
             else:
-                t_obj, created = through_cls.objects.get_or_create(**data)
+                if not self.MODEL_CREATION_LIMIT or \
+                        through_cls in self.MODEL_CREATION_LIMIT:
+                    t_obj, created = through_cls.objects.get_or_create(**data)
+                else:
+                    get_data = data.copy()
+                    if 'defaults' in get_data:
+                        get_data.pop('defaults')
+                    try:
+                        t_obj = through_cls.objects.get(**get_data)
+                    except through_cls.DoesNotExist:
+                        values = u", ".join(
+                            [u"{}: {}".format(k, get_data[k]) for k in get_data]
+                        )
+                        raise ImporterError(
+                            unicode(self.ERRORS['does_not_exist_in_db']
+                                    ).format(through_cls, values))
             if not created and 'defaults' in data:
                 for k in data['defaults']:
                     setattr(t_obj, k, data['defaults'][k])
@@ -1247,6 +1280,12 @@ class Importer(object):
                             new_created[attribute].append(key)
                             has_values = bool([1 for k in v if v[k]])
                             if has_values:
+                                if self.MODEL_CREATION_LIMIT and \
+                                        model not in self.MODEL_CREATION_LIMIT:
+                                    raise ImproperlyConfigured(
+                                        unicode(
+                                            self.ERRORS['improperly_configured']
+                                        ).format(model))
                                 v = model.objects.create(**v)
                             else:
                                 continue
@@ -1255,14 +1294,32 @@ class Importer(object):
                             extra_fields = {}
                             # "File" type is a temp object and can be different
                             # for the same filename - it must be treated
-                            # separatly
+                            # separately
                             for field in model._meta.fields:
                                 k = field.name
-                                # attr_class est un attribut de FileField
+                                # attr_class is a FileField attribute
                                 if hasattr(field, 'attr_class') and k in v:
                                     extra_fields[k] = v.pop(k)
-                            v, created = model.objects.get_or_create(
-                                **v)
+                            if not self.MODEL_CREATION_LIMIT or \
+                                    model in self.MODEL_CREATION_LIMIT:
+                                v, created = model.objects.get_or_create(
+                                    **v)
+                            else:
+                                get_v = v.copy()
+                                if 'defaults' in get_v:
+                                    get_v.pop('defaults')
+                                try:
+                                    v = model.objects.get(**get_v)
+                                except model.DoesNotExist:
+                                    values = u", ".join(
+                                        [u"{}: {}".format(k, get_v[k])
+                                         for k in get_v]
+                                    )
+                                    raise ImporterError(
+                                        unicode(
+                                            self.ERRORS[
+                                                'does_not_exist_in_db']
+                                        ).format(model, values))
                             changed = False
                             for k in extra_fields.keys():
                                 if extra_fields[k]:
@@ -1336,6 +1393,7 @@ class Importer(object):
                     'history_modifier': create_dict.pop('history_modifier')
                 })
 
+            created = False
             try:
                 try:
                     dct = create_dict.copy()
@@ -1348,6 +1406,11 @@ class Importer(object):
                             return None, created
                         new_dct = defaults.copy()
                         new_dct.update(dct)
+                        if self.MODEL_CREATION_LIMIT and \
+                                cls not in self.MODEL_CREATION_LIMIT:
+                            raise ImproperlyConfigured(
+                                unicode(self.ERRORS[ 'improperly_configured']
+                                        ).format(cls))
                         obj = cls.objects.create(**new_dct)
                     else:
                         # manage UNICITY_KEYS - only level 1
@@ -1356,9 +1419,21 @@ class Importer(object):
                                 if k not in self.UNICITY_KEYS \
                                    and k != 'defaults':
                                     defaults[k] = dct.pop(k)
-
-                        dct['defaults'] = defaults.copy()
-                        obj, created = cls.objects.get_or_create(**dct)
+                        if not self.MODEL_CREATION_LIMIT or \
+                                cls in self.MODEL_CREATION_LIMIT:
+                            dct['defaults'] = defaults.copy()
+                            obj, created = cls.objects.get_or_create(**dct)
+                        else:
+                            try:
+                                obj = cls.objects.get(**dct)
+                                dct['defaults'] = defaults.copy()
+                            except cls.DoesNotExist:
+                                values = u", ".join(
+                                    [u"{}: {}".format(k, dct[k]) for k in dct]
+                                )
+                                raise ImporterError(
+                                    unicode(self.ERRORS['does_not_exist_in_db']
+                                            ).format(cls, values))
 
                         if not created and not path and self.UNICITY_KEYS:
                             changed = False