SERVER-20374 added bsonExtractIntegerFieldWithDefaultIf to support user-defined constraint on extracted value

3 files changed, 63 insertions, 0 deletions

diff --git a/src/mongo/bson/util/bson_extract.cpp b/src/mongo/bson/util/bson_extract.cpp
index 54d1e1f47ae..4caaf700406 100644
--- a/src/mongo/bson/util/bson_extract.cpp
+++ b/src/mongo/bson/util/bson_extract.cpp
@@ -196,4 +196,31 @@ Status bsonExtractIntegerFieldWithDefault(const BSONObj& object,
     return status;
 }
 
+Status bsonExtractIntegerFieldWithDefaultIf(const BSONObj& object,
+                                            StringData fieldName,
+                                            long long defaultValue,
+                                            stdx::function<bool(long long)> pred,
+                                            const std::string& predDescription,
+                                            long long* out) {
+    auto status = bsonExtractIntegerFieldWithDefault(object, fieldName, defaultValue, out);
+    if (!status.isOK()) {
+        return status;
+    }
+    if (!pred(*out)) {
+        return Status(ErrorCodes::BadValue,
+                      mongoutils::str::stream() << "Invalid value in field \"" << fieldName
+                                                << "\": " << *out << ": " << predDescription);
+    }
+    return Status::OK();
+}
+
+Status bsonExtractIntegerFieldWithDefaultIf(const BSONObj& object,
+                                            StringData fieldName,
+                                            long long defaultValue,
+                                            stdx::function<bool(long long)> pred,
+                                            long long* out) {
+    return bsonExtractIntegerFieldWithDefaultIf(
+        object, fieldName, defaultValue, pred, "constraint failed", out);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/bson/util/bson_extract.h b/src/mongo/bson/util/bson_extract.h
index c5e34d3e3ba..ddc571d99aa 100644
--- a/src/mongo/bson/util/bson_extract.h
+++ b/src/mongo/bson/util/bson_extract.h
@@ -34,6 +34,7 @@
 #include "mongo/base/string_data.h"
 #include "mongo/bson/bsontypes.h"
 #include "mongo/db/repl/optime.h"
+#include "mongo/stdx/functional.h"
 
 namespace mongo {
 
@@ -187,4 +188,28 @@ Status bsonExtractOIDFieldWithDefault(const BSONObj& object,
                                       const OID& defaultValue,
                                       OID* out);
 
+/**
+ * Finds an element named "fieldName" in "object" that represents an integral value for which
+ * 'pred' is true.
+ *
+ * If a field named "fieldName" is present and is a value of numeric type with an exact 64-bit
+ * integer representation, returns that representation in *out and returns Status::OK().
+ * If there is no field named "fieldName", stores defaultValue into *out and returns Status::OK().
+ * If the field is found, but has non-numeric type, returns ErrorCodes::TypeMismatch.
+ * If the value has numeric type, but cannot be represented as a 64-bit integer, returns BadValue.
+ * If the parsed value (or default) fails the predicate, returns ErrorCodes::BadValue.
+ */
+Status bsonExtractIntegerFieldWithDefaultIf(const BSONObj& object,
+                                            StringData fieldName,
+                                            long long defaultValue,
+                                            stdx::function<bool(long long)> pred,
+                                            const std::string& predDescription,
+                                            long long* out);
+
+Status bsonExtractIntegerFieldWithDefaultIf(const BSONObj& object,
+                                            StringData fieldName,
+                                            long long defaultValue,
+                                            stdx::function<bool(long long)> pred,
+                                            long long* out);
+
 }  // namespace mongo
diff --git a/src/mongo/bson/util/bson_extract_test.cpp b/src/mongo/bson/util/bson_extract_test.cpp
index 95951efb87e..0c7cb031edb 100644
--- a/src/mongo/bson/util/bson_extract_test.cpp
+++ b/src/mongo/bson/util/bson_extract_test.cpp
@@ -25,12 +25,14 @@
  *    then also delete it in the license file.
  */
 
+#include <functional>
 #include <limits>
 #include <string>
 
 #include "mongo/bson/util/bson_extract.h"
 #include "mongo/db/jsobj.h"
 #include "mongo/db/repl/optime.h"
+#include "mongo/stdx/functional.h"
 #include "mongo/unittest/unittest.h"
 
 using namespace mongo;
@@ -163,4 +165,13 @@ TEST(ExtractBSON, ExtractIntegerField) {
     ASSERT_EQUALS(-(1LL << 55), v);
     ASSERT_OK(bsonExtractIntegerField(BSON("a" << 5178), "a", &v));
     ASSERT_EQUALS(5178, v);
+    auto pred = stdx::bind(std::greater<long long>(), stdx::placeholders::_1, 0);
+    ASSERT_OK(bsonExtractIntegerFieldWithDefaultIf(BSON("a" << 1), "a", -1LL, pred, &v));
+    ASSERT_OK(bsonExtractIntegerFieldWithDefaultIf(BSON("a" << 1), "b", 1LL, pred, &v));
+    auto msg = "'a' has to be greater than zero";
+    auto status = bsonExtractIntegerFieldWithDefaultIf(BSON("a" << -1), "a", 1LL, pred, msg, &v);
+    ASSERT_EQUALS(ErrorCodes::BadValue, status);
+    ASSERT_STRING_CONTAINS(status.reason(), msg);
+    ASSERT_EQUALS(ErrorCodes::BadValue,
+                  bsonExtractIntegerFieldWithDefaultIf(BSON("a" << 1), "b", -1LL, pred, &v));
 }