Changeset 83 for trunk/dbpickle/dbpickle.py

Timestamp:

03/02/07 19:59:34 (3 years ago)

Author:

akaihola

Message:

dbpickle now supports many-to-many relations.

Files:

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

trunk/dbpickle/dbpickle.py

@@ -29 +29 @@
 
 This version handles ForeignKeys which refer to the related model's
-primary key.  ManyToMany and OneToOne relations are not supported.
+primary key.  ManyToMany relations are now supported, but OneToOne
+relations are not.
 
 You can also use you own plugin for pre-processing objects before
@@ -76 +77 @@
 def dump(filepath):
     """
-    Pickle all Django objects from the database and save them into the given file.
-    """
-    data = {}
+    Pickle all Django objects from the database and save them into the
+    given file.
+    """
+    objects = {}   # model instances
+    m2m_lists = [] # ManyToMany relations between model instances
     for model in models.get_models():
         meta = model._meta
         app, model_name = meta.app_label, meta.module_name
+
+        # get all many-to-many relation field names for this model
+        m2ms = [m2m.name for m2m in meta.many_to_many]
+        
         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'))
+            pk = obj._get_pk_val()
+            objects[app, model_name, pk] = obj
+            
+            for m2m in m2ms:
+                # store many-to-many related objects for every
+                # many-to-many relation of this object
+                foreign_objs = getattr(obj, m2m).all()
+                logging.info('dumping %s.%s.%s x %d' % (app, model_name, m2m, len(foreign_objs)))
+                m2m_lists.append((obj, m2m, tuple(foreign_objs)))
+
+    # pickle all objects and many-to-many relations on disk
+    cPickle.dump((objects, m2m_lists), file(filepath, 'w'))
 
 
@@ -95 +112 @@
     """
 
+    # load the plugin if specified on the command line
     if pluginpath:
         plugin = load_source('plugin', pluginpath)
@@ -100 +118 @@
         plugin = None
 
+    # get the hook functions from the plugin
     hooks = {}
     for hookname in 'pre_save', 'pre_walk':
         hooks[hookname] = getattr(plugin, '%s_hook' % hookname, lambda obj: False)
     
-    # unpickle objects from file
-    data = cPickle.load(file(filepath))
-
-    # delete objects from all models to be loaded
-    models = set([obj.__class__ for obj in data.itervalues()])
+    # unpickle objects and many-to-many relations from disk
+    objects, m2m_lists = cPickle.load(file(filepath))
+
+    # Find distinct models of all unpickled objects and delete all
+    # objects before loading.  Note that models which have not been
+    # dumped are not emptied.
+    models = set( [obj.__class__ for obj in objects.itervalues()] )
     for model in models:
         for obj in model._default_manager.all():
@@ -114 +135 @@
 
     # load all objects
-    while data:
-        key, obj = data.popitem()
-        load_recursive(data, obj, hooks)
+    while objects:
+        key, obj = objects.popitem()
+        load_recursive(objects, obj, hooks)
+
+    # load all many-to-many relations
+    for obj1, m2m, foreign_objs in m2m_lists:
+        meta1 = obj1._meta
+        for obj2 in foreign_objs:
+            meta2 = obj2._meta
+            logging.info('loading ManyToMany %s.%s.%s -> %s.%s.%s' % (
+                meta1.app_label, meta1.module_name, obj1._get_pk_val(),
+                meta2.app_label, meta2.module_name, obj2._get_pk_val()))
+            getattr(obj1, m2m).add(obj2)
 
 load = transaction.commit_on_success(load)
 
 
-def load_recursive(data, obj, hooks):
+def load_recursive(objects, obj, hooks):
     """
     Save the given object into the database.  If the object has
@@ -128 +159 @@
     """
 
+    meta = obj._meta
+    
     hooks['pre_walk'](obj)
-    for field in obj._meta.fields:
+    for field in meta.fields:
         if isinstance(field, models.ForeignKey):
-            related_app = field.rel.to._meta.app_label
-            related_model = field.rel.to._meta.module_name
+            related_meta = field.rel.to._meta
+            related_app = related_meta.app_label
+            related_model = related_meta.module_name
             related_pk_val = getattr(obj, field.name+'_id')
             try:
-                related_obj = data.pop((related_app,
-                                        related_model,
-                                        related_pk_val))
-                load_recursive(data, related_obj, hooks)
+                related_obj = objects.pop((related_app,
+                                           related_model,
+                                           related_pk_val))
+                load_recursive(objects, related_obj, hooks)
             except KeyError:
                 logging.debug('probably loaded already: '
                               '%(related_app)s.%(related_model)s '
                               '%(related_pk_val)s' % locals())
-    if logging.getLogger().level <= logging.INFO:
-        # Must go through this extra step so str(obj) isn't evaluated
-        # if no logging message is to be printed.  Sometimes --quiet
-        # is needed to avoid error messages when __repr__ or __str__
-        # tries to look up objects not yet loaded.
-        logging.info('loading %s' % obj)
+                                      
+    logging.info('loading %s.%s %s' % (
+        meta.app_label,
+        meta.module_name,
+        obj._get_pk_val()))
     try:
         hooks['pre_save'](obj)