Switch to unified view

a/scripts/restore_project.py b/scripts/restore_project.py
1
import os
1
import os
2
import sys
2
import sys
3
import json
3
import struct
4
import logging
4
import logging
5
5
6
from ming.orm import state, session
6
from ming.orm import state, session, mapper, MappedClass
7
from ming.orm.base import instrument, DocumentTracker
7
from ming.orm.base import instrument, DocumentTracker
8
8
9
from pymongo.json_util import default, object_hook
9
from pylons import c
10
from pymongo.bson import BSON
10
from allura import model as M
11
from allura import model as M
11
from allura.command import ReindexCommand
12
from allura.command import ReindexCommand
12
13
13
log = logging.getLogger(__name__)
14
log = logging.getLogger(__name__)
14
15
MONGO_HOME=os.environ.get('MONGO_HOME', '/usr')
16
MONGO_DUMP=os.path.join(MONGO_HOME, 'bin/mongodump')
17
MONGO_RESTORE=os.path.join(MONGO_HOME, 'bin/mongorestore')
18
15
19
def main():
16
def main():
20
    if len(sys.argv) != 4:
17
    if len(sys.argv) != 4:
21
        log.error('Usage: %s <dirname> <new_shortname> <new_unix_group_name>', sys.argv[0])
18
        log.error('Usage: %s <dirname> <new_shortname> <new_unix_group_name>', sys.argv[0])
22
        return 2
19
        return 2
...
...
25
    new_ug_name = sys.argv[3]
22
    new_ug_name = sys.argv[3]
26
    return restore_project(dirname, new_pname, new_ug_name)
23
    return restore_project(dirname, new_pname, new_ug_name)
27
24
28
def restore_project(dirname, new_shortname, new_unix_group_name):
25
def restore_project(dirname, new_shortname, new_unix_group_name):
29
    log.info('Reloading %s into %s', dirname, new_shortname)
26
    log.info('Reloading %s into %s', dirname, new_shortname)
30
    with open(os.path.join(dirname, 'project.json')) as fp:
27
    with open(os.path.join(dirname, 'project.bson')) as fp:
31
        project_doc = json.load(fp, object_hook=object_hook)
28
        project_doc = _read_bson(fp)
32
    project = M.Project.query.get(_id=project_doc['_id'])
29
    project = M.Project.query.get(_id=project_doc['_id'])
33
    st = state(project)
30
    st = state(project)
34
    st.document = instrument(project_doc, DocumentTracker(st))
31
    st.document = instrument(project_doc, DocumentTracker(st))
35
    if project is None:
32
    if project is None:
36
        log.fatal('Project not found')
33
        log.fatal('Project not found')
37
        return 2
34
        return 2
38
    dump_path = os.path.join(dirname, project.database)
39
    with open(os.path.join(dirname, 'project.json')) as fp:
40
        project_doc = json.load(fp, object_hook=object_hook)
41
    st = state(project)
42
    st.document = instrument(project_doc, DocumentTracker(st))
43
    project.shortname = new_shortname
35
    project.shortname = new_shortname
44
    project.set_tool_data('sfx', unix_group_name=new_unix_group_name)
36
    project.set_tool_data('sfx', unix_group_name=new_unix_group_name)
45
    project.database = 'project:' + new_shortname.replace('/', ':').replace('-', '_')
46
    project.deleted = False
37
    project.deleted = False
47
    conn = M.main_doc_session.bind.conn
38
    c.project = project
48
    if project.database in conn.database_names():
39
    for name, cls in MappedClass._registry.iteritems():
49
        raw_input('''Warning: database %s is already populated.  If you do NOT want
40
        if session(cls) is None: continue
50
    to drop the database and create a new one, press Crtl-C NOW!  Otherwise,
41
        m = mapper(cls)
51
    press enter to continue.''' % project.database)
42
        sess = session(cls).impl
52
        conn.drop_database(project.database)
43
        fname = os.path.join(dirname, '%s.bson' % (cls.__mongometa__.name))
53
    os.system('%s --db %s %s' % (
44
        if not os.path.exists(fname): continue
54
            MONGO_RESTORE, project.database, dump_path))
45
        if ('project_id' not in m.property_index
46
            and 'app_config_id' not in m.property_index): continue
47
        with open(fname, 'rb') as fp:
48
            num_objects = 0
49
            while True:
50
                doc = _read_bson(fp)
51
                if doc is None: break
52
                num_objects += 1
53
                sess.insert(m.doc_cls(doc))
54
        log.info('%s: loaded %s objects from %s',
55
                 name, num_objects, fname)
55
    session(project).flush()
56
    session(project).flush()
56
    reindex= ReindexCommand('reindex')
57
    reindex= ReindexCommand('reindex')
57
    reindex.run(['--project', new_shortname])
58
    reindex.run(['--project', new_shortname])
58
    return 0
59
    return 0
59
60
61
def _read_bson(fp):
62
    slen = fp.read(4)
63
    if not slen: return None
64
    bson = BSON(fp.read(struct.unpack('!l', slen)[0]))
65
    return bson.to_dict()
66
60
if __name__ == '__main__':
67
if __name__ == '__main__':
61
    sys.exit(main())
68
    sys.exit(main())