ZODB export with support for cross-database references.

multizexp.py

"""ZODB export with support for cross-database references."""
# Inital version by David Glick
# http://glicksoftware.com/blog/recombining-zodb-storages
#
# ZODB 3.10 and Blob support by Thomas Buchberger

from ZODB.blob import Blob
from ZODB.utils import p64, u64, cp
from ZODB.ExportImport import export_end_marker
from ZODB.ExportImport import blob_begin_marker
from ZODB.DemoStorage import DemoStorage
import logging
import cPickle
import cStringIO
import os

logger = logging.getLogger('multiexport')


def export_zexp(self, fname):
    context = self
    f = open(fname, 'wb')
    f.write('ZEXP')
    for oid, p, blobfilename in flatten_multidatabase(context):
        f.writelines((oid, p64(len(p)), p))
        if blobfilename:
            f.write(blob_begin_marker)
            f.write(p64(os.stat(blobfilename).st_size))
            blobfile = open(blobfilename, "rb")
            cp(blobfile, f)
            blobfile.close()

    f.write(export_end_marker)
    f.close()


def flatten_multidatabase(context):
    """Walk a multidatabase and yield rewritten pickles with oids for a single
       database"""
    base_oid = context._p_oid
    base_conn = context._p_jar
    dbs = base_conn.connections

    dummy_storage = DemoStorage()

    oids = [(base_conn._db.database_name, base_oid)]
    done_oids = {}
    # table to keep track of mapping old oids to new oids
    ooid_to_oid = {oids[0]: dummy_storage.new_oid()}
    while oids:
        # loop while references remain to objects we haven't exported yet
        (dbname, ooid) = oids.pop(0)
        if (dbname, ooid) in done_oids:
            continue
        done_oids[(dbname, ooid)] = True

        db = dbs[dbname]
        try:
            # get pickle
            p, serial = db._storage.load(ooid)
        except:
            logger.debug("broken reference for db %s, oid %s",
                         (dbname, repr(ooid)),
                         exc_info=True)
        else:
            def persistent_load(ref):
                """Remap a persistent id to a new ID and create a ghost for it.
                   This is called by the unpickler for each reference found.
                """

                # resolve the reference to a database name and oid
                if isinstance(ref, tuple):
                    rdbname, roid = (dbname, ref[0])
                elif isinstance(ref, str):
                    rdbname, roid = (dbname, ref)
                else:
                    try:
                        ref_type, args = ref
                    except ValueError:
                        # weakref
                        return
                    else:
                        if ref_type in ('m', 'n'):
                            rdbname, roid = (args[0], args[1])
                        else:
                            return

                # traverse Products.ZODBMountpoint mountpoints to the mounted
                # location
                rdb = dbs[rdbname]
                p, serial = rdb._storage.load(roid)
                klass = p.split()[0]
                if 'MountedObject' in klass:
                    mountpoint = rdb.get(roid)
                    # get the object with the root as a parent, then unwrap,
                    # since there's no API to get the unwrapped object
                    mounted = mountpoint._getOrOpenObject(app).aq_base
                    rdbname = mounted._p_jar._db.database_name
                    roid = mounted._p_oid

                if roid:
                    print '%s:%s -> %s:%s' % (dbname, u64(ooid),
                                              rdbname, u64(roid))
                    oids.append((rdbname, roid))

                try:
                    oid = ooid_to_oid[(rdbname, roid)]
                except KeyError:
                    # generate a new oid and associate it with this old db/oid
                    ooid_to_oid[(rdbname, roid)] = oid = dummy_storage.new_oid()
                return Ghost(oid)

            if not isinstance(db._reader.getGhost(p), Blob):
                blobfilename = None
            else:
                blobfilename = db._storage.loadBlob(ooid, serial)

            # do the repickling dance to rewrite references

            pfile = cStringIO.StringIO(p)
            unpickler = cPickle.Unpickler(pfile)
            unpickler.persistent_load = persistent_load

            newp = cStringIO.StringIO()
            pickler = cPickle.Pickler(newp, 1)
            pickler.persistent_id = persistent_id

            pickler.dump(unpickler.load())
            pickler.dump(unpickler.load())
            p = newp.getvalue()

            yield ooid_to_oid[(dbname, ooid)], p, blobfilename


class Ghost(object):
    __slots__ = ("oid",)

    def __init__(self, oid):
        self.oid = oid


def persistent_id(obj):
    if isinstance(obj, Ghost):
        return obj.oid


mysite = app.unrestrictedTraverse('/Plone')
export_zexp(mysite, '/tmp/mysite.zexp')