From d7889fb311bdf7d174bbf2b9dc4bd5d1829c5c42 Mon Sep 17 00:00:00 2001
From: Eliot Horowitz <eliot@10gen.com>
Date: Thu, 4 Jun 2009 08:49:06 -0400
Subject: [PATCH] Sharding: save() support SHARDING-17

---
 jstests/sharding/shard6.js | 16 ++++++++++++++++
 s/strategy_shard.cpp       | 11 ++++++++---
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/jstests/sharding/shard6.js b/jstests/sharding/shard6.js
index 88c809955b7..fab251bc2cf 100644
--- a/jstests/sharding/shard6.js
+++ b/jstests/sharding/shard6.js
@@ -20,4 +20,20 @@ for ( ; num<100; num++ ){
 
 assert.eq( 100 , db.data.find().toArray().length );
 
+// limit
+
+assert.eq( 77 , db.data.find().limit(77).itcount() , "limit test 1" );
+assert.eq( 1 , db.data.find().limit(1).itcount() , "limit test 2" );
+for ( var i=1; i<10; i++ ){
+    assert.eq( i , db.data.find().limit(i).itcount() , "limit test 3 : " + i );
+}
+
+
+// --- test save support ---
+
+o = db.data.findOne();
+o.x = 16;
+db.data.save( o );
+assert.eq( 16 , db.data.findOne( { _id : o._id } ).x , "x1 - did save fail?" );
+
 s.stop();
diff --git a/s/strategy_shard.cpp b/s/strategy_shard.cpp
index 761245cc50a..ad6ef99aafe 100644
--- a/s/strategy_shard.cpp
+++ b/s/strategy_shard.cpp
@@ -126,11 +126,16 @@ namespace mongo {
             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" );
+            bool save = false;
+            if ( ! manager->hasShardKey( query ) ){
+                if ( query.nFields() != 1 || strcmp( query.firstElement().fieldName() , "_id" ) )
+                    throw UserException( "can't do update with query that doesn't have the shard key" );
+                save = true;
+            }
             
-            if ( manager->hasShardKey( toupdate ) && manager->getShardKey().compare( query , toupdate ) )
+            if ( ! save && manager->hasShardKey( toupdate ) && manager->getShardKey().compare( query , toupdate ) ){
                 throw UserException( "change would move shards!" );
+            }
 
             Shard& s = manager->findShard( toupdate );
             doWrite( dbUpdate , r , s.getServer() );