Changeset 75 for trunk/dbpickle/dbpickle.py

Timestamp:

02/27/07 09:59:46 (3 years ago)

Author:

akaihola

Message:

dbpickle now handles foreign keys correctly.

Files:

• trunk/dbpickle/dbpickle.py (modified) (2 diffs)

trunk/dbpickle/dbpickle.py

@@ -28 +28 @@
 for migrating from one database engine to another.
 
+This version handles ForeignKeys which refer to the related model's
+primary key.  ManyToMany and OneToOne relations are not supported.
+
 Let's suppose your project is using PostgreSQL and the database has
 been populated using the app.
@@ -52 +55 @@
 import cPickle
 import logging
+from django.db import models
 
 def dump(filepath):
-    from django.db import models
+    """
+    Pickle all Django objects from the database and save them into the given file.
+    """
     data = {}
     for model in models.get_models():
         meta = model._meta
         app, model_name = meta.app_label, meta.module_name
-        objects = list(model.objects.all())
-        logging.info('dumping model %s.%s (%d objects)' % (app, model_name, len(objects)))
-        data[app, model_name] = objects
+        for obj in model.objects.all():
+            logging.info('dumping %s.%s %s' % (app, model_name, obj))
+            data[app, model_name, obj._get_pk_val()] = obj
     cPickle.dump(data, file(filepath, 'w'))
+
 
 @transaction.commit_on_success
 def load(filepath):
+    """
+    Unpickle Django objects from the given file and save them into the
+    database.
+    """
+    
+    # unpickle objects from file
     data = cPickle.load(file(filepath))
-    for (app, model), objects in data.items():
-        logging.info('loading model %s.%s' % (app, model))
-        for obj in objects:
-            logging.debug(str(obj))
+
+    # delete objects from all models to be loaded
+    models = set(obj.__class__ for obj in data.itervalues())
+    for model in models:
+        for obj in model._default_manager.all():
+            obj.delete()
+
+    # load all objects
+    while data:
+        key, obj = data.popitem()
+        load_recursive(data, obj)
+
+
+def load_recursive(data, obj):
+    """
+    Save the given object into the database.  If the object has
+    ForeignKey relations to other objects, first make sure they are
+    already saved (and repeat recursively).
+    """
+    
+    for field in obj._meta.fields:
+        if isinstance(field, models.ForeignKey):
+            related_app = field.rel.to._meta.app_label
+            related_model = field.rel.to._meta.module_name
+            related_pk_val = getattr(obj, field.name+'_id')
             try:
-                obj.save()
-            except Exception, e:
-                logging.error('%s while saving %s.%s %s' % (e, app, model, str(obj)))
+                related_obj = data.pop((related_app,
+                                        related_model,
+                                        related_pk_val))
+                load_recursive(data, related_obj)
+            except KeyError:
+                logging.debug('probably loaded already: '
+                              '%(related_app)s.%(related_model)s '
+                              '%(related_pk_val)s' % locals())
+    logging.info('loading %s' % obj)
+    try:
+        obj.save()
+    except Exception, e:
+        logging.error('%s while saving %s' % (e, obj))
+
 
 if __name__ == '__main__':