Migrate legacy tables without a pk

legacy/scripts/migration.py

@@ -1,14 +1,13 @@
 from django.apps.config import AppConfig
 from django.core.exceptions import ObjectDoesNotExist
-from django.db import connection, models
+from django.db import connections, models
 from django.db.models.base import ModelBase
+from model_mommy import mommy
 
 from field_mappings import field_mappings
-from migration_base import legacy_app
+from migration_base import legacy_app, appconfs
 from parlamentares.models import Parlamentar
 
-from model_mommy import mommy
-
 
 def info(msg):
     print 'INFO: ' + msg
@@ -36,7 +35,7 @@ def none_to_false(value):
     return bool(value)
 
 
-def migrate(obj, count_limit=None):
+def migrate(obj=appconfs, count_limit=None):
     # warning: model/app migration order is of utmost importance
 
     to_delete = []
@@ -59,12 +58,23 @@ def _do_migrate(obj, to_delete, count_limit=None):
         raise TypeError('Parameter must be a Model, AppConfig or a sequence of them')
 
 
-def exec_sql(sql):
+def exec_sql(sql, db='default'):
-    cursor = connection.cursor()
+    cursor = connections[db].cursor()
     cursor.execute(sql)
     return cursor
 
 
+def iter_sql_records(sql, db):
+    class Record(object):
+        pass
+    cursor = exec_sql(sql, db)
+    fieldnames = [name[0] for name in cursor.description]
+    for row in cursor.fetchall():
+        record = Record()
+        record.__dict__.update(zip(fieldnames, row))
+        yield record
+
+
 def save_with_id(new, id):
     sequence_name = '%s_id_seq' % type(new)._meta.db_table
     cursor = exec_sql('SELECT last_value from %s;' % sequence_name)
@@ -93,42 +103,56 @@ def migrate_model(model, to_delete, count_limit=None):
 
     legacy_model = legacy_app.get_model(model.__name__)
     old_pk_name = legacy_model._meta.pk.name
+
+    # setup migration strategy for tables with or without a pk
     if old_pk_name == 'id':
         # There is no pk in the legacy table
-        pass
+
-        # ...
+        def get_old_pk(old):
-        # ...
+            return '-- WITHOUT PK --'
+
+        def save(new, id):
+            new.save()
+
+        old_records = iter_sql_records(
+            'select * from ' + legacy_model._meta.db_table, 'legacy')
     else:
-        for old in legacy_model.objects.all().order_by(old_pk_name)[:count_limit]:
+        def get_old_pk(old):
-            old_pk = getattr(old, old_pk_name)
+            return getattr(old, old_pk_name)
-            new = model()
+        save = save_with_id
-            for new_field, old_field in field_mappings[model].items():
+        old_records = legacy_model.objects.all().order_by(old_pk_name)[:count_limit]
-                value = getattr(old, old_field)
+
-                model_field = model._meta.get_field(new_field)
+    # convert old records to new ones
-                transform = special_transforms.get(model_field)
+    for old in old_records:
-                if transform:
+        old_pk = get_old_pk(old)
-                    value = transform(value)
+        new = model()
-                else:
+        for new_field, old_field in field_mappings[model].items():
-                    # check for a relation
+            value = getattr(old, old_field)
-                    if isinstance(model_field, models.ForeignKey) and value is not None:
+            model_field = model._meta.get_field(new_field)
-                        try:
+            transform = special_transforms.get(model_field)
-                            value = model_field.related_model.objects.get(id=value)
+            if transform:
-                        except ObjectDoesNotExist:
+                value = transform(value)
-                            msg = 'FK [%s (%s) : %s] not found for value %s' % (
+            else:
-                                model.__name__, old_pk, model_field.name, value)
+                # check for a relation
-                            if value == 0:
+                if isinstance(model_field, models.ForeignKey) and value is not None:
-                                # we interpret FK == 0 as actually FK == NONE
+                    try:
-                                value = None
+                        value = model_field.related_model.objects.get(id=value)
-                                warn(msg + ' => NONE for zero value')
+                    except ObjectDoesNotExist:
-                            else:
+                        msg = 'FK [%s (%s) : %s] not found for value %s' % (
-                                value = make_stub(model_field.related_model, value)
+                            model.__name__, old_pk, model_field.name, value)
-                                warn(msg + ' => STUB CREATED')
+                        if value == 0:
+                            # we interpret FK == 0 as actually FK == NONE
+                            value = None
+                            warn(msg + ' => NONE for zero value')
                         else:
-                            assert value
+                            value = make_stub(model_field.related_model, value)
-                setattr(new, new_field, value)
+                            warn(msg + ' => STUB CREATED')
-            save_with_id(new, old_pk)
+                    else:
-            if getattr(old, 'ind_excluido', False):
+                        assert value
-                to_delete.append(new)
+            setattr(new, new_field, value)
+        save(new, old_pk)
+        if getattr(old, 'ind_excluido', False):
+            to_delete.append(new)
 
 
 def get_ind_excluido(obj):