From d71cb1d7d2c33430bf57d3ca93e684cb39914435 Mon Sep 17 00:00:00 2001 From: Eliot Horowitz Date: Tue, 27 Jan 2009 15:16:09 -0500 Subject: [PATCH] moving to unified Tool system --- SConstruct | 5 +- tools/Tool.cpp | 63 ++++++++++++++++ tools/Tool.h | 47 ++++++++++++ tools/dump.cpp | 200 +++++++++++++++++++------------------------------ 4 files changed, 191 insertions(+), 124 deletions(-) create mode 100644 tools/Tool.cpp create mode 100644 tools/Tool.h diff --git a/SConstruct b/SConstruct index e83fc9ca99d..11a1fb76381 100644 --- a/SConstruct +++ b/SConstruct @@ -306,8 +306,9 @@ testEnv.Append( LIBPATH=["."] ) Default( env.Program( "mongod" , commonFiles + coreDbFiles + serverOnlyFiles + [ "db/db.cpp" ] ) ) # tools -env.Program( "mongodump" , allClientFiles + [ "tools/dump.cpp" ] ) -env.Program( "mongoimport" , allClientFiles + [ "tools/import.cpp" ] ) +allToolFiles = allClientFiles + [ "tools/Tool.cpp" ] +env.Program( "mongodump" , allToolFiles + [ "tools/dump.cpp" ] ) +env.Program( "mongoimport" , allToolFiles + [ "tools/import.cpp" ] ) # dbgrid env.Program( "mongogrid" , commonFiles + coreDbFiles + Glob( "dbgrid/*.cpp" ) ) diff --git a/tools/Tool.cpp b/tools/Tool.cpp new file mode 100644 index 00000000000..abcaf099142 --- /dev/null +++ b/tools/Tool.cpp @@ -0,0 +1,63 @@ +// Tool.cpp + +#include "Tool.h" + +#include + +#include + +using namespace std; +using namespace mongo; + +namespace po = boost::program_options; + +mongo::Tool::Tool( string name , string defaultDB , string defaultCollection ) : + _name( name ) , _db( defaultDB ) , _coll( defaultCollection ){ + + _options = new po::options_description( name + " options" ); + _options->add_options() + ("help","produce help message") + ("host,h",po::value(), "mongo host to connect to" ) + ("db,d",po::value(), "database to use" ) + ("collection,c",po::value(), "collection to use (some commands)" ) + ; + +} + +mongo::Tool::~Tool(){ + delete( _options ); +} + +int mongo::Tool::main( int argc , char ** argv ){ + boost::filesystem::path::default_name_check( boost::filesystem::no_check ); + + po::store( po::parse_command_line( argc, argv, *_options ), _params ); + po::notify( _params ); + + if ( _params.count( "help" ) ){ + _options->print( cerr ); + return 0; + } + + const char * host = "127.0.0.1"; + if ( _params.count( "host" ) ) + host = _params["host"].as().c_str(); + + string errmsg; + if ( ! _conn.connect( host , errmsg ) ){ + cerr << "couldn't connect to [" << host << "] " << errmsg << endl; + return -1; + } + + cout << "connected to: " << host << endl; + + if ( _params.count( "db" ) ) + _db = _params["db"].as(); + + if ( _params.count( "collection" ) ) + _coll = _params["collection"].as(); + + run(); + + return 0; +} diff --git a/tools/Tool.h b/tools/Tool.h new file mode 100644 index 00000000000..ae8a6917807 --- /dev/null +++ b/tools/Tool.h @@ -0,0 +1,47 @@ +// Tool.h + +#pragma once + +#include + +#include + +#include "client/dbclient.h" + +using std::string; + +namespace mongo { + + class Tool { + public: + Tool( string name , string defaultDB="test" , string defaultCollection=""); + virtual ~Tool(); + + int main( int argc , char ** argv ); + + boost::program_options::options_description_easy_init add_options(){ + return _options->add_options(); + } + + string getParam( string name , string def="" ){ + if ( _params.count( name ) ) + return _params[name.c_str()].as(); + return def; + } + + virtual void run() = 0; + + protected: + string _name; + mongo::DBClientConnection _conn; + + string _db; + string _coll; + + private: + boost::program_options::options_description * _options; + boost::program_options::variables_map _params; + + }; + +} diff --git a/tools/dump.cpp b/tools/dump.cpp index a8b4e1fb139..7ec0bcae70c 100644 --- a/tools/dump.cpp +++ b/tools/dump.cpp @@ -18,140 +18,96 @@ #include "../stdafx.h" #include "../client/dbclient.h" +#include "Tool.h" #include -namespace mongo { - - namespace po = boost::program_options; - - namespace dump { - - void doCollection( DBClientConnection & conn , const char * coll , path outputFile ) { - mongo::out() << "\t" << coll << " to " << outputFile.string() << endl; - - int out = open( outputFile.string().c_str() , O_WRONLY | O_CREAT | O_TRUNC , 0666 ); - assert( out ); - - BSONObjBuilder query; - auto_ptr cursor = conn.query( coll , query.doneAndDecouple() ); - - int num = 0; - while ( cursor->more() ) { - BSONObj obj = cursor->next(); - write( out , obj.objdata() , obj.objsize() ); - num++; - } - - mongo::out() << "\t\t " << num << " objects" << endl; - - close( out ); - } - - void go( DBClientConnection & conn , const char * db , const path outdir ) { - mongo::out() << "DATABASE: " << db << endl; - - create_directories( outdir ); - - string sns = db; - sns += ".system.namespaces"; - - BSONObjBuilder query; - auto_ptr cursor = conn.query( sns.c_str() , query.doneAndDecouple() ); - while ( cursor->more() ) { - BSONObj obj = cursor->next(); - if ( obj.toString().find( ".$" ) != string::npos ) - continue; - - const string name = obj.getField( "name" ).valuestr(); - const string filename = name.substr( strlen( db ) + 1 ); - - doCollection( conn , name.c_str() , outdir / ( filename + ".bson" ) ); - - } - - } - - void go( const char * host , const char * db , const char * outdir ) { - DBClientConnection conn; - string errmsg; - if ( ! conn.connect( host , errmsg ) ) { - mongo::out() << "couldn't connect : " << errmsg << endl; - throw -11; - } - - path root(outdir); - - if ( strlen( db ) == 1 && db[0] == '*' ) { - mongo::out() << "all dbs" << endl; - - BSONObjBuilder query; - query.appendBool( "listDatabases" , 1 ); - - BSONObj res = conn.findOne( "admin.$cmd" , query.doneAndDecouple() ); - BSONObj dbs = res.getField( "databases" ).embeddedObjectUserCheck(); - set keys; - dbs.getFieldNames( keys ); - for ( set::iterator i = keys.begin() ; i != keys.end() ; i++ ) { - string key = *i; - - BSONObj db = dbs.getField( key ).embeddedObjectUserCheck(); - - const char * dbName = db.getField( "name" ).valuestr(); - if ( (string)dbName == "local" ) - continue; - go ( conn , dbName , root / dbName ); - } - } - else { - go( conn , db , root / db ); - } - } - - } // namespace dump - -} // namespace mongo - using namespace mongo; -int main( int argc , char ** argv ) { +namespace po = boost::program_options; - boost::filesystem::path::default_name_check( boost::filesystem::no_check ); +class Dump : public Tool { +public: + Dump() : Tool( "dump" , "*" ){ + add_options() + ("out,o" , po::value() , "output directory" ) + ; + } + + void doCollection( const string coll , path outputFile ) { + cout << "\t" << coll << " to " << outputFile.string() << endl; + + int out = open( outputFile.string().c_str() , O_WRONLY | O_CREAT | O_TRUNC , 0666 ); + assert( out ); + + auto_ptr cursor = _conn.query( coll.c_str() , emptyObj ); - po::options_description options("dump parameters"); - options.add_options() - ("help", "produce help message") - ("host,h", po::value() , "mongo host to connect to") - ("db,d" , po::value() , "database to dump" ) - ("out" , po::value() , "output directory" ) - ; + int num = 0; + while ( cursor->more() ) { + BSONObj obj = cursor->next(); + write( out , obj.objdata() , obj.objsize() ); + num++; + } + + cout << "\t\t " << num << " objects" << endl; + + close( out ); + } + + void go( const string db , const path outdir ) { + cout << "DATABASE: " << db << "\t to \t" << outdir.string() << endl; + + create_directories( outdir ); + + string sns = db + ".system.namespaces"; - po::variables_map vm; - po::store(po::parse_command_line(argc, argv, options), vm); - po::notify(vm); - - if ( vm.count("help") ) { - options.print( cerr ); - return 1; + auto_ptr cursor = _conn.query( sns.c_str() , emptyObj ); + while ( cursor->more() ) { + BSONObj obj = cursor->next(); + if ( obj.toString().find( ".$" ) != string::npos ) + continue; + + const string name = obj.getField( "name" ).valuestr(); + const string filename = name.substr( db.size() + 1 ); + + doCollection( name.c_str() , outdir / ( filename + ".bson" ) ); + + } + } - const char * host = "127.0.0.1"; - const char * db = "*"; - const char * outdir = "dump"; + void run(){ - if ( vm.count( "host" ) ) - host = vm["host"].as().c_str(); + path root( getParam( "out" , "dump" ) ); + string db = _db; + + if ( db == "*" ){ + cout << "all dbs" << endl; - if ( vm.count( "db" ) ) - db = vm["db"].as().c_str(); + BSONObj res = _conn.findOne( "admin.$cmd" , BUILDOBJ( "listDatabases" << 1 ) ); + BSONObj dbs = res.getField( "databases" ).embeddedObjectUserCheck(); + set keys; + dbs.getFieldNames( keys ); + for ( set::iterator i = keys.begin() ; i != keys.end() ; i++ ) { + string key = *i; + + BSONObj dbobj = dbs.getField( key ).embeddedObjectUserCheck(); + + const char * dbName = dbobj.getField( "name" ).valuestr(); + if ( (string)dbName == "local" ) + continue; - if ( vm.count( "out" ) ) - outdir = vm["out"].as().c_str(); + go ( dbName , root / dbName ); + } + } + else { + go( db , root / db ); + } + + } +}; - mongo::out() << "mongo dump" << endl; - mongo::out() << "\t host \t" << host << endl; - mongo::out() << "\t db \t" << db << endl; - mongo::out() << "\t output dir \t" << outdir << endl; - - dump::go( host , db , outdir ); +int main( int argc , char ** argv ) { + Dump d; + return d.main( argc , argv ); }