distinct() for sharded collection SHARDING-35

2 files changed, 87 insertions, 0 deletions

diff --git a/jstests/sharding/features2.js b/jstests/sharding/features2.js
new file mode 100644
index 00000000000..10548b209dc
--- /dev/null
+++ b/jstests/sharding/features2.js
@@ -0,0 +1,27 @@
+// features2.js
+
+s = new ShardingTest( "features2" , 2 , 1 , 1 );
+s.adminCommand( { enablesharding : "test" } );
+
+a = s._connections[0].getDB( "test" );
+b = s._connections[1].getDB( "test" );
+
+// ---- distinct ----
+
+db = s.getDB( "test" );
+db.foo.save( { x : 1 } );
+db.foo.save( { x : 2 } );
+db.foo.save( { x : 3 } );
+
+assert.eq( "1,2,3" , db.foo.distinct( "x" ) , "distinct 1" );
+assert( a.foo.distinct("x").length == 3 || b.foo.distinct("x").length == 3 , "distinct 2" );
+assert( a.foo.distinct("x").length == 0 || b.foo.distinct("x").length == 0 , "distinct 3" );
+
+s.adminCommand( { shardcollection : "test.foo" , key : { x : 1 } } );
+s.adminCommand( { split : "test.foo" , middle : { x : 2 } } );
+s.adminCommand( { movechunk : "test.foo" , find : { x : 3 } , to : s.getOther( s.getServer( "test" ) ).name } );
+
+assert.eq( "1,2,3" , db.foo.distinct( "x" ) , "distinct 4" );
+
+// -----
+s.stop();
diff --git a/s/commands_public.cpp b/s/commands_public.cpp
index f94b62ff30c..34214954c02 100644
--- a/s/commands_public.cpp
+++ b/s/commands_public.cpp
@@ -139,6 +139,66 @@ namespace mongo {
             
         } groupCmd;
 
+        class DistinctCmd : public PublicGridCommand {
+        public:
+            DistinctCmd() : PublicGridCommand("distinct"){}
+            virtual void help( stringstream &help ) const {
+                help << "{ distinct : 'collection name' , key : 'a.b' }";
+            }
+            bool run(const char *ns, BSONObj& cmdObj, string& errmsg, BSONObjBuilder& result, bool){
+
+                string dbName = getDBName( ns );
+                string collection = cmdObj.firstElement().valuestrsafe();
+                string fullns = dbName + "." + collection;
 
+                DBConfig * conf = grid.getDBConfig( dbName , false );
+                
+                if ( ! conf || ! conf->isShardingEnabled() || ! conf->isSharded( fullns ) ){
+                    return passthrough( conf , cmdObj , result );
+                }
+                
+                ChunkManager * cm = conf->getChunkManager( fullns );
+                massert( "how could chunk manager be null!" , cm );
+                
+                vector<Chunk*> chunks;
+                cm->getChunksForQuery( chunks , BSONObj() );
+                
+                set<BSONObj,BSONObjCmp> all;
+                int size = 32;
+                
+                for ( vector<Chunk*>::iterator i = chunks.begin() ; i != chunks.end() ; i++ ){
+                    Chunk * c = *i;
+
+                    ScopedDbConnection conn( c->getShard() );
+                    BSONObj res;
+                    bool ok = conn->runCommand( conf->getName() , cmdObj , res );
+                    conn.done();
+                    
+                    if ( ! ok ){
+                        result.appendElements( res );
+                        return false;
+                    }
+                    
+                    BSONObjIterator it( res["values"].embeddedObjectUserCheck() );
+                    while ( it.more() ){
+                        BSONElement nxt = it.next();
+                        BSONObjBuilder temp(32);
+                        temp.appendAs( nxt , "x" );
+                        all.insert( temp.obj() );
+                    }
+
+                }
+                
+                BSONObjBuilder b( size );
+                int n=0;
+                for ( set<BSONObj,BSONObjCmp>::iterator i = all.begin() ; i != all.end(); i++ ){
+                    b.appendAs( i->firstElement() , b.numStr( n++ ).c_str() );
+                }
+                
+                result.appendArray( "values" , b.obj() );
+                result.append( "ok" , 1 );
+                return true;
+            }
+        } disinctCmd;
     }
 }