sharded update

1 files changed, 38 insertions, 12 deletions

diff --git a/s/strategy_shard.cpp b/s/strategy_shard.cpp
index 1b9cd5de046..d07fd1e4e4e 100644
--- a/s/strategy_shard.cpp
+++ b/s/strategy_shard.cpp
@@ -102,6 +102,42 @@ namespace mongo {
             cursorCache.remove( id );
         }
         
+        void receivedInsert( Request& r , DbMessage& d, ShardManager* manager ){
+            while ( d.moreJSObjs() ){
+                BSONObj o = d.nextJsObj();
+                if ( ! manager->hasShardKey( o ) ){
+                    log() << "tried to insert object without shard key: " << r.getns() << "  " << o << endl;
+                    throw UserException( "tried to insert object without shard key" );
+                }
+                
+                Shard& s = manager->findShard( o );
+                log(4) << "  server:" << s.getServer() << " " << o << endl;
+                insert( s.getServer() , r.getns() , o );
+            }            
+        }
+
+        void receivedUpdate( Request& r , DbMessage& d, ShardManager* manager ){
+            int flags = d.pullInt();
+            
+            BSONObj query = d.nextJsObj();
+            uassert( "invalid update" , d.moreJSObjs() );
+            BSONObj toupdate = d.nextJsObj();
+
+            
+            bool upsert = flags & 1;
+            if ( upsert && ! manager->hasShardKey( toupdate ) )
+                throw UserException( "can't upsert something without shard key" );
+
+            if ( ! manager->hasShardKey( query ) )
+                throw UserException( "can't do update with query that doesn't have the shard key" );
+            
+            if ( manager->hasShardKey( toupdate ) && manager->getShardKey().compare( query , toupdate ) )
+                throw UserException( "change would move shards!" );
+
+            Shard& s = manager->findShard( toupdate );
+            doWrite( dbUpdate , r , s.getServer() );
+        }
+
         virtual void writeOp( int op , Request& r ){
             
             const char *ns = r.getns();
@@ -112,20 +148,10 @@ namespace mongo {
             assert( info );
             
             if ( op == dbInsert ){
-                while ( d.moreJSObjs() ){
-                    BSONObj o = d.nextJsObj();
-                    if ( ! info->hasShardKey( o ) ){
-                        log() << "tried to insert object without shard key: " << ns << "  " << o << endl;
-                        throw UserException( "tried to insert object without shard key" );
-                    }
-                    
-                    Shard& s = info->findShard( o );
-                    log(4) << "  server:" << s.getServer() << " " << o << endl;
-                    insert( s.getServer() , ns , o );
-                }
+                receivedInsert( r , d , info );
             }
             else if ( op == dbUpdate ){
-                throw UserException( "can't do update yet on sharded collection" );
+                receivedUpdate( r , d , info );    
             }
             else {
                 log() << "sharding can't do write op: " << op << endl;