← Back to branch summary

~ubuntu-branches/ubuntu/trusty/mongodb/trusty-proposed

« back to all changes in this revision

Viewing changes to tools/restore.cpp

  • Committer: Bazaar Package Importer
  • Author(s): Antonin Kral
  • Date: 2010-01-29 19:48:45 UTC
  • Revision ID: james.westby@ubuntu.com-20100129194845-8wbmkf626fwcavc9
Tags: upstream-1.3.1
Import upstream version 1.3.1

Show diffs side-by-side

added added

removed removed

Lines of Context:
tools/restore.cpp
 
1
// restore.cpp
 
2
 
 
3
/**
 
4
*    Copyright (C) 2008 10gen Inc.
 
5
*
 
6
*    This program is free software: you can redistribute it and/or  modify
 
7
*    it under the terms of the GNU Affero General Public License, version 3,
 
8
*    as published by the Free Software Foundation.
 
9
*
 
10
*    This program is distributed in the hope that it will be useful,
 
11
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
 
12
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
13
*    GNU Affero General Public License for more details.
 
14
*
 
15
*    You should have received a copy of the GNU Affero General Public License
 
16
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
17
*/
 
18
 
 
19
#include "../stdafx.h"
 
20
#include "../client/dbclient.h"
 
21
#include "../util/mmap.h"
 
22
#include "tool.h"
 
23
 
 
24
#include <boost/program_options.hpp>
 
25
 
 
26
#include <fcntl.h>
 
27
 
 
28
using namespace mongo;
 
29
 
 
30
namespace po = boost::program_options;
 
31
 
 
32
class Restore : public Tool {
 
33
public:
 
34
    Restore() : Tool( "restore" , "" , "" ){
 
35
        add_hidden_options()
 
36
            ("dir", po::value<string>()->default_value("dump"), "directory to restore from")
 
37
            ;
 
38
        addPositionArg("dir", 1);
 
39
    }
 
40
 
 
41
    virtual void printExtraHelp(ostream& out) {
 
42
        out << "usage: " << _name << " [options] [directory or filename to restore from]" << endl;
 
43
    }
 
44
 
 
45
    int run(){
 
46
        auth();
 
47
        path root = getParam("dir");
 
48
 
 
49
        /* If _db is not "" then the user specified a db name to restore as.
 
50
         *
 
51
         * In that case we better be given either a root directory that
 
52
         * contains only .bson files or a single .bson file  (a db).
 
53
         *
 
54
         * In the case where a collection name is specified we better be
 
55
         * given either a root directory that contains only a single
 
56
         * .bson file, or a single .bson file itself (a collection).
 
57
         */
 
58
        drillDown(root, _db != "", _coll != "");
 
59
        return EXIT_CLEAN;
 
60
    }
 
61
 
 
62
    void drillDown( path root, bool use_db = false, bool use_coll = false ) {
 
63
        log(2) << "drillDown: " << root.string() << endl;
 
64
 
 
65
        if ( is_directory( root ) ) {
 
66
            directory_iterator end;
 
67
            directory_iterator i(root);
 
68
            while ( i != end ) {
 
69
                path p = *i;
 
70
                i++;
 
71
 
 
72
                if (use_db) {
 
73
                    if (is_directory(p)) {
 
74
                        cerr << "ERROR: root directory must be a dump of a single database" << endl;
 
75
                        cerr << "       when specifying a db name with --db" << endl;
 
76
                        printHelp(cout);
 
77
                        return;
 
78
                    }
 
79
                }
 
80
 
 
81
                if (use_coll) {
 
82
                    if (is_directory(p) || i != end) {
 
83
                        cerr << "ERROR: root directory must be a dump of a single collection" << endl;
 
84
                        cerr << "       when specifying a collection name with --collection" << endl;
 
85
                        printHelp(cout);
 
86
                        return;
 
87
                    }
 
88
                }
 
89
 
 
90
                drillDown(p, use_db, use_coll);
 
91
            }
 
92
            return;
 
93
        }
 
94
 
 
95
        if ( ! ( endsWith( root.string().c_str() , ".bson" ) ||
 
96
                 endsWith( root.string().c_str() , ".bin" ) ) ) {
 
97
            cerr << "don't know what to do with [" << root.string() << "]" << endl;
 
98
            return;
 
99
        }
 
100
 
 
101
        out() << root.string() << endl;
 
102
 
 
103
        string ns;
 
104
        if (use_db) {
 
105
            ns += _db;
 
106
        } else {
 
107
            string dir = root.branch_path().string();
 
108
            if ( dir.find( "/" ) == string::npos )
 
109
                ns += dir;
 
110
            else
 
111
                ns += dir.substr( dir.find_last_of( "/" ) + 1 );
 
112
        }
 
113
 
 
114
        if (use_coll) {
 
115
            ns += "." + _coll;
 
116
        } else {
 
117
            string l = root.leaf();
 
118
            l = l.substr( 0 , l.find_last_of( "." ) );
 
119
            ns += "." + l;
 
120
        }
 
121
 
 
122
        long long fileLength = file_size( root );
 
123
 
 
124
        if ( fileLength == 0 ) {
 
125
            out() << "file " + root.native_file_string() + " empty, skipping" << endl;
 
126
            return;
 
127
        }
 
128
 
 
129
        out() << "\t going into namespace [" << ns << "]" << endl;
 
130
 
 
131
        string fileString = root.string();
 
132
        ifstream file( fileString.c_str() , ios_base::in | ios_base::binary);
 
133
        if ( ! file.is_open() ){
 
134
            log() << "error opening file: " << fileString << endl;
 
135
            return;
 
136
        }
 
137
 
 
138
        log(1) << "\t file size: " << fileLength << endl;
 
139
 
 
140
        long long read = 0;
 
141
        long long num = 0;
 
142
 
 
143
        const int BUF_SIZE = 1024 * 1024 * 5;
 
144
        char * buf = (char*)malloc( BUF_SIZE );
 
145
 
 
146
        ProgressMeter m( fileLength );
 
147
 
 
148
        while ( read < fileLength ) {
 
149
            file.read( buf , 4 );
 
150
            int size = ((int*)buf)[0];
 
151
            if ( size >= BUF_SIZE ){
 
152
                cerr << "got an object of size: " << size << "  terminating..." << endl;
 
153
            }
 
154
            uassert( 10264 ,  "invalid object size" , size < BUF_SIZE );
 
155
 
 
156
            file.read( buf + 4 , size - 4 );
 
157
 
 
158
            BSONObj o( buf );
 
159
            conn().insert( ns.c_str() , o );
 
160
 
 
161
            read += o.objsize();
 
162
            num++;
 
163
 
 
164
            m.hit( o.objsize() );
 
165
        }
 
166
 
 
167
        free( buf );
 
168
 
 
169
        uassert( 10265 ,  "counts don't match" , m.done() == fileLength );
 
170
        out() << "\t "  << m.hits() << " objects" << endl;
 
171
    }
 
172
};
 
173
 
 
174
int main( int argc , char ** argv ) {
 
175
    Restore restore;
 
176
    return restore.main( argc , argv );
 
177
}