better handling of sharded upserts. SHARDING-72

more update fixes coming after 1.3.1

2 files changed, 37 insertions, 1 deletions

diff --git a/jstests/sharding/update1.js b/jstests/sharding/update1.js
new file mode 100644
index 00000000000..82c3d8a77e3
--- /dev/null
+++ b/jstests/sharding/update1.js
@@ -0,0 +1,33 @@
+s = new ShardingTest( "auto1" , 2 , 1 , 1 );
+
+s.adminCommand( { enablesharding : "test" } );
+s.adminCommand( { shardcollection : "test.update1" , key : { key : 1 } } );
+
+db = s.getDB( "test" )
+coll = db.update1;
+
+coll.insert({_id:1, key:1});
+
+// these are upserts
+coll.save({_id:2, key:2});
+coll.save({_id:3, key:3});
+
+assert.eq(coll.count(), 3, "count A")
+
+// update existing using save()
+coll.save({_id:1, key:1, other:1});
+
+// update existing using update()
+coll.update({_id:2}, {key:2, other:2});
+//coll.update({_id:3, key:3}, {other:3}); //should add key to new object (doesn't work yet)
+coll.update({_id:3}, {key:3, other:3});
+
+assert.eq(coll.count(), 3, "count B")
+coll.find().forEach(function(x){
+    assert.eq(x._id, x.key, "_id == key");
+    assert.eq(x._id, x.other, "_id == other");
+});
+
+
+s.stop()
+
diff --git a/s/strategy_shard.cpp b/s/strategy_shard.cpp
index e38eeb84f9e..34cf226f38d 100644
--- a/s/strategy_shard.cpp
+++ b/s/strategy_shard.cpp
@@ -153,8 +153,11 @@ namespace mongo {
             if ( multi )
                 uassert( 10202 ,  "can't mix multi and upsert and sharding" , ! upsert );
 
-            if ( upsert && ! manager->hasShardKey( query ) )
+            if ( upsert && !(manager->hasShardKey(toupdate) ||
+                             (toupdate.firstElement().fieldName()[0] == '$' && manager->hasShardKey(query))))
+            {
                 throw UserException( 8012 , "can't upsert something without shard key" );
+            }
 
             bool save = false;
             if ( ! manager->hasShardKey( query ) ){