sharded dropDatabase SHARDING-53

6 files changed, 195 insertions, 13 deletions

diff --git a/jstests/sharding/shard3.js b/jstests/sharding/shard3.js
index b3f52f4c160..8c5b184e798 100644
--- a/jstests/sharding/shard3.js
+++ b/jstests/sharding/shard3.js
@@ -37,16 +37,18 @@ assert.eq( 3 , b.find().toArray().length , "other B" );
 
 // --- filtering ---
 
-function doCounts( name ){
-    assert.eq( 3 , a.count() , name + " count" );    
-    assert.eq( 3 , a.find().sort( { n : 1 } ).itcount() , name + " itcount - sort n" );
-    assert.eq( 3 , a.find().itcount() , name + " itcount" );
-    assert.eq( 3 , a.find().sort( { _id : 1 } ).itcount() , name + " itcount - sort _id" );
+function doCounts( name , total ){
+    total = total || ( primary.count() + secondary.count() );
+    assert.eq( total , a.count() , name + " count" );    
+    assert.eq( total , a.find().sort( { n : 1 } ).itcount() , name + " itcount - sort n" );
+    assert.eq( total , a.find().itcount() , name + " itcount" );
+    assert.eq( total , a.find().sort( { _id : 1 } ).itcount() , name + " itcount - sort _id" );
+    return total;
 }
 
-doCounts( "before wrong save" )
+var total = doCounts( "before wrong save" )
 secondary.save( { num : -3 } );
-doCounts( "after wrong save" )
+doCounts( "after wrong save" , total )
 
 // --- move all to 1 ---
 print( "MOVE ALL TO 1" );
@@ -103,5 +105,26 @@ assert.isnull( b.findOne( { num : 4 } ) , "b drop1" );
 
 s.printCollectionInfo( "test.foo" , "after b findOne tests" );
 
+print( "*** dropDatabase setup" )
+
+s.printShardingStatus()
+s.adminCommand( { shardcollection : "test.foo" , key : { num : 1 } } );
+a.save( { num : 2 } );
+a.save( { num : 3 } );
+s.adminCommand( { split : "test.foo" , middle : { num : 2 } } );
+s.adminCommand( { movechunk : "test.foo" , find : { num : 3 } , to : s.getOther( s.getServer( "test" ) ).name } );
+s.printShardingStatus();
+
+s.printCollectionInfo( "test.foo" , "after dropDatabase setup" );
+doCounts( "after dropDatabase setup2" )
+s.printCollectionInfo( "test.foo" , "after dropDatabase setup3" );
+
+print( "*** ready to call dropDatabase" )
+res = s.getDB( "test" ).dropDatabase();
+assert.eq( 1 , res.ok , "dropDatabase failed : " + tojson( res ) );
+
+s.printShardingStatus();
+s.printCollectionInfo( "test.foo" , "after dropDatabase call 1" );
+assert.eq( 0 , doCounts( "after dropDatabase called" ) )
 
 s.stop();
diff --git a/s/chunk.cpp b/s/chunk.cpp
index 0834abd220b..9f1cec1359c 100644
--- a/s/chunk.cpp
+++ b/s/chunk.cpp
@@ -472,6 +472,12 @@ namespace mongo {
         }
         return added;
     }
+
+    void ChunkManager::getAllServers( set<string>& allServers ){
+        for ( vector<Chunk*>::iterator i=_chunks.begin(); i != _chunks.end(); i++  ){
+            allServers.insert( (*i)->getShard() );
+        }        
+    }
     
     void ChunkManager::ensureIndex(){
         set<string> seen;
diff --git a/s/chunk.h b/s/chunk.h
index db56f1cc548..e813de334d7 100644
--- a/s/chunk.h
+++ b/s/chunk.h
@@ -190,6 +190,8 @@ namespace mongo {
          */
         int getChunksForQuery( vector<Chunk*>& chunks , const BSONObj& query );
 
+        void getAllServers( set<string>& allServers );
+
         void save();
 
         string toString() const;
diff --git a/s/commands_public.cpp b/s/commands_public.cpp
index c25c20ae098..db28829cc00 100644
--- a/s/commands_public.cpp
+++ b/s/commands_public.cpp
@@ -102,11 +102,40 @@ namespace mongo {
                 
                 cm->drop();
 
-                result.append( "ok" , 1 );
                 return 1;
             }
         } dropCmd;
 
+        class DropDBCmd : public PublicGridCommand {
+        public:
+            DropDBCmd() : PublicGridCommand( "dropDatabase" ){}
+            bool run(const char *ns, BSONObj& cmdObj, string& errmsg, BSONObjBuilder& result, bool){
+                
+                BSONElement e = cmdObj.firstElement();
+                
+                if ( ! e.isNumber() || e.number() != 1 ){
+                    errmsg = "invalid params";
+                    return 0;
+                }
+                
+                string dbName = getDBName( ns );
+                DBConfig * conf = grid.getDBConfig( dbName , false );
+                
+                log() << "DROP DATABASE: " << dbName << endl;
+
+                if ( ! conf || ! conf->isShardingEnabled() ){
+                    log(1) << "  passing though drop database for: " << dbName << endl;
+                    return passthrough( conf , cmdObj , result );
+                }
+                
+                if ( ! conf->dropDatabase( errmsg ) )
+                    return false;
+
+                result.append( "dropped" , dbName );
+                return true;
+            }
+        } dropDBCmd;
+
         class CountCmd : public PublicGridCommand {
         public:
             CountCmd() : PublicGridCommand("count") { }
@@ -124,7 +153,6 @@ namespace mongo {
                     ScopedDbConnection conn( conf->getPrimary() );
                     result.append( "n" , (double)conn->count( fullns , filter ) );
                     conn.done();
-                    result.append( "ok" , 1 );
                     return true;
                 }
                 
@@ -141,7 +169,6 @@ namespace mongo {
                 }
                 
                 result.append( "n" , (double)total );
-                result.append( "ok" , 1 );
                 return true;
             }
         } countCmd;
@@ -224,7 +251,6 @@ namespace mongo {
                 }
                 
                 result.appendArray( "values" , b.obj() );
-                result.append( "ok" , 1 );
                 return true;
             }
         } disinctCmd;
diff --git a/s/config.cpp b/s/config.cpp
index eaf839c77e2..661945d5a71 100644
--- a/s/config.cpp
+++ b/s/config.cpp
@@ -163,9 +163,102 @@ namespace mongo {
         BSONObj q = b.done();
         return load(q);
     }
+
+    bool DBConfig::dropDatabase( string& errmsg ){
+        /**
+         * 1) make sure everything is up
+         * 2) update config server
+         * 3) drop and reset sharded collections
+         * 4) drop and reset primary
+         * 5) drop everywhere to clean up loose ends
+         */
+
+        log(1) << "DBConfig::dropDatabase: " << _name << endl;
+        
+        // 1
+        if ( ! configServer.allUp( errmsg ) ){
+            log(1) << "\t DBConfig::dropDatabase not all up" << endl;
+            return 0;
+        }
+        
+        // 2
+        grid.removeDB( _name );
+        remove( true );
+        if ( ! configServer.allUp( errmsg ) ){
+            log() << "error removing from config server even after checking!" << endl;
+            return 0;
+        }
+        log(1) << "\t removed entry from config server for: " << _name << endl;
+        
+        set<string> allServers;
+
+        // 3
+        while ( true ){
+            int num;
+            if ( ! _dropShardedCollections( num , allServers , errmsg ) )
+                return 0;
+            log() << "   DBConfig::dropDatabase: " << _name << " dropped sharded collections: " << num << endl;
+            if ( num == 0 )
+                break;
+        }
+        
+        // 4
+        {
+            ScopedDbConnection conn( _primary );
+            BSONObj res;
+            if ( ! conn->dropDatabase( _name , &res ) ){
+                errmsg = res.toString();
+                return 0;
+            }
+            conn.done();
+        }
+        
+        // 5
+        for ( set<string>::iterator i=allServers.begin(); i!=allServers.end(); i++ ){
+            string s = *i;
+            ScopedDbConnection conn( s );
+            BSONObj res;
+            if ( ! conn->dropDatabase( _name , &res ) ){
+                errmsg = res.toString();
+                return 0;
+            }
+            conn.done();            
+        }
+        
+        log(1) << "\t dropped primary db for: " << _name << endl;
+
+        return true;
+    }
+
+    bool DBConfig::_dropShardedCollections( int& num, set<string>& allServers , string& errmsg ){
+        num = 0;
+        set<string> seen;
+        while ( true ){
+            map<string,ChunkManager*>::iterator i = _shards.begin();
+
+            if ( i == _shards.end() )
+                break;
+
+            if ( seen.count( i->first ) ){
+                errmsg = "seen a collection twice!";
+                return false;
+            }
+
+            seen.insert( i->first );
+            log(1) << "\t dropping sharded collection: " << i->first << endl;
+
+            i->second->getAllServers( allServers );
+            i->second->drop();
+            
+            num++;
+            uassert( "_dropShardedCollections too many collections - bailing" , num < 100000 );
+            log(2) << "\t\t dropped " << num << " so far" << endl;
+        }
+        return true;
+    }
     
     /* --- Grid --- */
-
+    
     string Grid::pickShardForNewDB(){
         ScopedDbConnection conn( configServer.getPrimary() );
         
@@ -203,6 +296,8 @@ namespace mongo {
         if ( database == "config" )
             return &configServer;
 
+        boostlock l( _lock );
+
         DBConfig*& cc = _databases[database];
         if ( cc == 0 ){
             cc = new DBConfig( database );
@@ -235,6 +330,13 @@ namespace mongo {
         return cc;
     }
 
+    void Grid::removeDB( string database ){
+        uassert( "removeDB expects db name" , database.find( '.' ) == string::npos );
+        boostlock l( _lock );
+        _databases.erase( database );
+        
+    }
+
     unsigned long long Grid::getNextOpTime() const {
         ScopedDbConnection conn( configServer.getPrimary() );
         
@@ -295,6 +397,11 @@ namespace mongo {
     }
 
     bool ConfigServer::allUp(){
+        string errmsg;
+        return allUp( errmsg );
+    }
+    
+    bool ConfigServer::allUp( string& errmsg ){
         try {
             ScopedDbConnection conn( _primary );
             conn->getLastError();
@@ -303,6 +410,7 @@ namespace mongo {
         }
         catch ( DBException& ){
             log() << "ConfigServer::allUp : " << _primary << " seems down!" << endl;
+            errmsg = _primary + " seems down";
             return false;
         }
         
diff --git a/s/config.h b/s/config.h
index 4bcaed47858..eade0265220 100644
--- a/s/config.h
+++ b/s/config.h
@@ -48,7 +48,8 @@ namespace mongo {
     };
     
     /**
-       top level grid configuration for an entire database
+     * top level grid configuration for an entire database
+     * TODO: use shared_ptr for ChunkManager
     */
     class DBConfig : public Model {
     public:
@@ -90,6 +91,8 @@ namespace mongo {
         }
 
         bool reload();
+
+        bool dropDatabase( string& errmsg );
         
         virtual void save( bool check=true);
 
@@ -102,6 +105,8 @@ namespace mongo {
         virtual void unserialize(const BSONObj& from);
         
     protected:
+
+        bool _dropShardedCollections( int& num, set<string>& allServers , string& errmsg );
         
         bool doload();
         
@@ -121,6 +126,10 @@ namespace mongo {
         friend class ChunkManager;
     };
 
+    /**
+     * stores meta-information about the grid
+     * TODO: used shard_ptr for DBConfig pointers
+     */
     class Grid {
     public:
         /**
@@ -129,6 +138,12 @@ namespace mongo {
          */
         DBConfig * getDBConfig( string ns , bool create=true);
         
+        /**
+         * removes db entry.
+         * on next getDBConfig call will fetch from db
+         */
+        void removeDB( string db );
+
         string pickShardForNewDB();
         
         bool knowAboutShard( string name ) const;
@@ -136,6 +151,7 @@ namespace mongo {
         unsigned long long getNextOpTime() const;
     private:
         map<string,DBConfig*> _databases;
+        boost::mutex _lock; // TODO: change to r/w lock
     };
 
     class ConfigServer : public DBConfig {
@@ -160,6 +176,7 @@ namespace mongo {
         bool init( vector<string> configHosts );
 
         bool allUp();
+        bool allUp( string& errmsg );
         
         int dbConfigVersion();
         int dbConfigVersion( DBClientBase& conn );