SERVER-7668 append MinKey instead of null when extending tag boundary

6 files changed, 182 insertions, 1 deletions

diff --git a/jstests/sharding/tag_auto_split.js b/jstests/sharding/tag_auto_split.js
index ac3c4c76637..6f29c589601 100644
--- a/jstests/sharding/tag_auto_split.js
+++ b/jstests/sharding/tag_auto_split.js
@@ -54,4 +54,10 @@ s.config.chunks.find().forEach(
     }
 );
 
+// check chunk mins correspond exactly to tag range boundaries, extended to match shard key
+assert.eq( 1, s.config.chunks.find( {min : {_id : 5 , a : MinKey} } ).count(),
+           "bad chunk range boundary" );
+assert.eq( 1, s.config.chunks.find( {min : {_id : 10 , a : MinKey} } ).count(),
+           "bad chunk range boundary" );
+
 s.stop();
diff --git a/src/mongo/db/keypattern.cpp b/src/mongo/db/keypattern.cpp
index 96229b70fd8..5123048d816 100644
--- a/src/mongo/db/keypattern.cpp
+++ b/src/mongo/db/keypattern.cpp
@@ -21,6 +21,8 @@
 #include "mongo/db/hasher.h"
 #include "mongo/db/queryutil.h"
 
+using namespace mongoutils;
+
 namespace mongo {
 
     BSONObj KeyPattern::extractSingleKey(const BSONObj& doc ) const {
@@ -61,6 +63,41 @@ namespace mongo {
         return true;
     }
 
+    BSONObj KeyPattern::extendRangeBound( const BSONObj& bound , bool makeUpperInclusive ) const {
+        BSONObjBuilder newBound( bound.objsize() );
+
+        BSONObjIterator src( bound );
+        BSONObjIterator pat( _pattern );
+
+        while( src.more() ){
+            massert( 16649 ,
+                     str::stream() << "keyPattern " << _pattern << " shorter than bound " << bound,
+                     pat.more() );
+            BSONElement srcElt = src.next();
+            BSONElement patElt = pat.next();
+            massert( 16634 ,
+                     str::stream() << "field names of bound " << bound
+                                   << " do not match those of keyPattern " << _pattern ,
+                                   str::equals( srcElt.fieldName() , patElt.fieldName() ) );
+            newBound.append( srcElt );
+        }
+        while( pat.more() ){
+            BSONElement patElt = pat.next();
+            // for non 1/-1 field values, like {a : "hashed"}, treat order as ascending
+            int order = patElt.isNumber() ? patElt.numberInt() : 1;
+            // flip the order semantics if this is an upper bound
+            if ( makeUpperInclusive ) order *= -1;
+
+            if( order > 0 ){
+                newBound.appendMinKey( patElt.fieldName() );
+            }
+            else {
+                newBound.appendMaxKey( patElt.fieldName() );
+            }
+        }
+        return newBound.obj();
+    }
+
     typedef vector<pair<BSONObj,BSONObj> >::const_iterator BoundListIter;
 
     BoundList KeyPattern::keyBounds( const FieldRangeSet& queryConstraints ) const {
diff --git a/src/mongo/db/keypattern.h b/src/mongo/db/keypattern.h
index 2bf02dd7bc6..fae49bb931a 100644
--- a/src/mongo/db/keypattern.h
+++ b/src/mongo/db/keypattern.h
@@ -74,6 +74,31 @@ namespace mongo {
             return _pattern.isPrefixOf( other.toBSON() );
         }
 
+        /* Takes a BSONObj whose field names are a prefix of the fields in this keyPattern, and
+         * outputs a new bound with MinKey values appended to match the fields in this keyPattern
+         * (or MaxKey values for descending -1 fields). This is useful in sharding for
+         * calculating chunk boundaries when tag ranges are specified on a prefix of the actual
+         * shard key, or for calculating index bounds when the shard key is a prefix of the actual
+         * index used.
+         *
+         * @param makeUpperInclusive If true, then MaxKeys instead of MinKeys will be appended, so
+         * that the output bound will compare *greater* than the bound being extended (note that
+         * -1's in the keyPattern will swap MinKey/MaxKey vals. See examples).
+         *
+         * Examples:
+         * If this keyPattern is {a : 1}
+         *   extendRangeBound( {a : 55}, false) --> {a : 55}
+         *
+         * If this keyPattern is {a : 1, b : 1}
+         *   extendRangeBound( {a : 55}, false) --> {a : 55, b : MinKey}
+         *   extendRangeBound( {a : 55}, true ) --> {a : 55, b : MaxKey}
+         *
+         * If this keyPattern is {a : 1, b : -1}
+         *   extendRangeBound( {a : 55}, false) --> {a : 55, b : MaxKey}
+         *   extendRangeBound( {a : 55}, true ) --> {a : 55, b : MinKey}
+         */
+        BSONObj extendRangeBound( const BSONObj& bound , bool makeUpperInclusive ) const;
+
         /**
          * Returns true if this KeyPattern contains any computed values, (e.g. {a : "hashed"}),
          * and false if this KeyPattern consists of only ascending/descending fields
diff --git a/src/mongo/dbtests/keypatterntests.cpp b/src/mongo/dbtests/keypatterntests.cpp
new file mode 100644
index 00000000000..7d8568c5706
--- /dev/null
+++ b/src/mongo/dbtests/keypatterntests.cpp
@@ -0,0 +1,109 @@
+// keypatterntests.cpp - Tests for the KeyPattern class
+//
+
+/**
+ *    Copyright (C) 2012 10gen Inc.
+ *
+ *    This program is free software: you can redistribute it and/or  modify
+ *    it under the terms of the GNU Affero General Public License, version 3,
+ *    as published by the Free Software Foundation.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    GNU Affero General Public License for more details.
+ *
+ *    You should have received a copy of the GNU Affero General Public License
+ *    along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+
+#include "mongo/db/keypattern.h"
+#include "mongo/dbtests/dbtests.h"
+
+namespace KeyPatternTests {
+
+    class ExtendRangeBoundTests {
+    public:
+        void run() {
+
+            BSONObj bound = BSON( "a" << 55 );
+            BSONObj longBound = BSON("a" << 55 << "b" << 66);
+
+            //test keyPattern shorter than bound, should fail
+            {
+                KeyPattern keyPat( BSON( "a" << 1 ) );
+                ASSERT_THROWS( keyPat.extendRangeBound( longBound, false ), MsgAssertionException );
+            }
+
+            //test keyPattern doesn't match bound, should fail
+            {
+                KeyPattern keyPat( BSON( "b" << 1 ) );
+                ASSERT_THROWS( keyPat.extendRangeBound( bound, false ), MsgAssertionException );
+            }
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "c" << 1) );
+                ASSERT_THROWS( keyPat.extendRangeBound( longBound, false ), MsgAssertionException );
+            }
+
+            //test keyPattern same as bound
+            {
+                KeyPattern keyPat( BSON( "a" << 1 ) );
+                BSONObj newB = keyPat.extendRangeBound( bound, false );
+                ASSERT_EQUALS( newB , BSON("a" << 55) );
+            }
+            {
+                KeyPattern keyPat( BSON( "a" << 1 ) );
+                BSONObj newB = keyPat.extendRangeBound( bound, false );
+                ASSERT_EQUALS( newB , BSON("a" << 55) );
+            }
+
+            //test keyPattern longer than bound, simple
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << 1) );
+                BSONObj newB = keyPat.extendRangeBound( bound, false );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MINKEY ) );
+            }
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << 1) );
+                BSONObj newB = keyPat.extendRangeBound( bound, true );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MAXKEY ) );
+            }
+
+            //test keyPattern longer than bound, more complex pattern directions
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << -1) );
+                BSONObj newB = keyPat.extendRangeBound( bound, false );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MAXKEY ) );
+            }
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << -1) );
+                BSONObj newB = keyPat.extendRangeBound( bound, true );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MINKEY ) );
+            }
+            {
+
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << -1 << "c" << 1 ) );
+                BSONObj newB = keyPat.extendRangeBound( bound, false );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MAXKEY << "c" << MINKEY ) );
+            }
+            {
+                KeyPattern keyPat( BSON( "a" << 1 << "b" << -1 << "c" << 1 ) );
+                BSONObj newB = keyPat.extendRangeBound( bound, true );
+                ASSERT_EQUALS( newB , BSON("a" << 55 << "b" << MINKEY << "c" << MAXKEY ) );
+            }
+        }
+    };
+
+    class All : public Suite {
+    public:
+        All() : Suite( "keypattern" ) {
+        }
+
+        void setupTests() {
+            add< ExtendRangeBoundTests >();
+        }
+    } myall;
+
+} // namespace JsobjTests
+
diff --git a/src/mongo/s/balance.cpp b/src/mongo/s/balance.cpp
index f8ebec32066..52712c961d2 100644
--- a/src/mongo/s/balance.cpp
+++ b/src/mongo/s/balance.cpp
@@ -278,7 +278,7 @@ namespace mongo {
             for ( unsigned i = 0; i < ranges.size(); i++ ) {
                 BSONObj min = ranges[i].min;
 
-                min = min.extractFields( cm->getShardKey().key(), true );
+                min = cm->getShardKey().extendRangeBound( min, false );
 
                 if ( allChunkMinimums.count( min ) > 0 )
                     continue;
diff --git a/src/mongo/s/shardkey.h b/src/mongo/s/shardkey.h
index 9e16268101f..01219a15752 100644
--- a/src/mongo/s/shardkey.h
+++ b/src/mongo/s/shardkey.h
@@ -80,6 +80,10 @@ namespace mongo {
             return pattern.hasField(key);
         }
 
+        BSONObj extendRangeBound( const BSONObj& bound , bool makeUpperInclusive ) const {
+            return pattern.extendRangeBound( bound , makeUpperInclusive );
+        }
+
         /**
          * @return
          * true if 'this' is a prefix (not necessarily contained) of 'otherPattern'.