PYTHON-4527 Make Range V2 opt in (#846)

bindings/python/pymongocrypt/mongocrypt.py

@@ -60,8 +60,9 @@ def __init__(self, options, callback):
         self.__crypt = lib.mongocrypt_new()
         if self.__crypt == ffi.NULL:
             raise MongoCryptError("unable to create new mongocrypt object")
-        if not lib.mongocrypt_setopt_use_range_v2(self.__crypt):
-            raise MongoCryptError("unable to enable QE Range Protocol V2")
+        if options.enable_range_v2:
+            if not lib.mongocrypt_setopt_use_range_v2(self.__crypt):
+                raise MongoCryptError("unable to enable QE Range Protocol V2")
 
         try:
             self.__init()

bindings/python/pymongocrypt/options.py

@@ -11,6 +11,7 @@ def __init__(
         crypt_shared_lib_path=None,
         crypt_shared_lib_required=False,
         bypass_encryption=False,
+        enable_range_v2=False,
     ):
         """Options for :class:`MongoCrypt`.
 
@@ -53,9 +54,13 @@ def __init__(
           - `crypt_shared_lib_required`: Whether to require a crypt_shared
             library.
           - `bypass_encryption`: Whether to bypass encryption.
+          - `enable_range_v2`: Whether to enable range V2.
+
+        .. versionadded:: 1.10
+           Added the ``enable_range_v2`` parameter.
 
         .. versionadded:: 1.3
-           ``crypt_shared_lib_path``, ``crypt_shared_lib_path``,
+           Added the ``crypt_shared_lib_path``, ``crypt_shared_lib_path``, and
            ``bypass_encryption`` parameters.
 
         .. versionadded:: 1.1
@@ -138,6 +143,7 @@ def __init__(
         self.crypt_shared_lib_path = crypt_shared_lib_path
         self.crypt_shared_lib_required = crypt_shared_lib_required
         self.bypass_encryption = bypass_encryption
+        self.enable_range_v2 = enable_range_v2
 
 
 class ExplicitEncryptOpts:

bindings/python/test/data/fle2-find-rangePreview-explicit/int32/rangeopts.json

@@ -0,0 +1,11 @@
+{
+    "min": {
+        "$numberInt": "0"
+    },
+    "max": {
+        "$numberInt": "200"
+    },
+    "sparsity": {
+        "$numberLong": "1"
+    }
+}

bindings/python/test/data/fle2-find-rangePreview-explicit/int32/value-to-encrypt.json

@@ -0,0 +1,20 @@
+{
+    "v": {
+        "$and": [
+            {
+                "age": {
+                    "$gte": {
+                        "$numberInt": "23"
+                    }
+                }
+            },
+            {
+                "age": {
+                    "$lte": {
+                        "$numberInt": "35"
+                    }
+                }
+            }
+        ]
+    }
+}

bindings/python/test/performance/perf_test.py

@@ -79,7 +79,7 @@ def tearDownModule():
 
 class TestBulkDecryption(unittest.TestCase):
     def setUp(self):
-        opts = MongoCryptOptions({"local": {"key": LOCAL_MASTER_KEY}})
+        opts = MongoCryptOptions({"local": {"key": LOCAL_MASTER_KEY}}, enable_range_v2=True)
         callback = MockCallback(key_docs=[bson_data("keyDocument.json")])
         self.mongocrypt = MongoCrypt(opts, callback)
         self.encrypter = ExplicitEncrypter(callback, opts)

bindings/python/test/test_mongocrypt.py

@@ -193,6 +193,14 @@ def test_mongocrypt_options_validation(self):
         ):
             MongoCryptOptions(valid_kms, encrypted_fields_map={})
 
+    def test_mongocrypt_options_range(self):
+        opts = MongoCryptOptions({"local": {"key": b"\x00" * 96}})
+        self.assertFalse(opts.enable_range_v2)
+        opts.enable_range_v2 = True
+        self.assertTrue(opts.enable_range_v2)
+        opts = MongoCryptOptions({"local": {"key": b"\x00" * 96}}, enable_range_v2=True)
+        self.assertTrue(opts.enable_range_v2)
+
 
 class TestMongoCrypt(unittest.TestCase):
     maxDiff = None
@@ -373,7 +381,7 @@ def test_encrypt_encrypted_fields_map(self):
         encrypted_fields_map = bson_data(
             "compact/success/encrypted-field-config-map.json"
         )
-        mc = self.create_mongocrypt(encrypted_fields_map=encrypted_fields_map)
+        mc = self.create_mongocrypt(encrypted_fields_map=encrypted_fields_map, enable_range_v2=True)
         self.addCleanup(mc.close)
         with mc.encryption_context("db", bson_data("compact/success/cmd.json")) as ctx:
             self.assertEqual(ctx.state, lib.MONGOCRYPT_CTX_NEED_MONGO_KEYS)
@@ -772,7 +780,8 @@ def mongo_crypt_opts():
                 {
                     "aws": {"accessKeyId": "example", "secretAccessKey": "example"},
                     "local": {"key": b"\x00" * 96},
-                }
+                },
+                enable_range_v2=True,
             )
 
         async def _test_encrypt_decrypt(self, key_id=None, key_alt_name=None):
@@ -993,13 +1002,7 @@ async def test_range_query_int32(self):
                 is_expression=True,
             )
             encrypted_val = bson.decode(encrypted, OPTS)
-
-            # Workaround for the internal range payload counter in libmongocrypt
-            if encrypted_val != expected:
-                expected = json_data(
-                    "fle2-find-range-explicit-v2/int32/encrypted-payload-second.json"
-                )
-            self.assertEqual(encrypted_val, expected)
+            self.assertEqual(encrypted_val, adjust_range_counter(encrypted_val, expected))
 
 
 class TestNeedKMSAzureCredentials(unittest.TestCase):
@@ -1167,6 +1170,22 @@ def insert_data_key(self, data_key):
         return bson.decode(data_key, OPTS)["_id"]
 
 
+def adjust_range_counter(encrypted_val, expected):
+    """Workaround for the internal range payload counter in libmongocrypt."""
+    if encrypted_val != expected:
+        _payload1 = expected["v"]["$and"][0]["age"]["$gte"]
+        _payload2 = expected["v"]["$and"][1]["age"]["$lte"]
+        _decoded1 = bson.decode(_payload1[1:])
+        _decoded2 = bson.decode(_payload2[1:])
+        for _ in range(10):
+            _decoded1["payloadId"] += 1
+            expected["v"]["$and"][0]["age"]["$gte"] = Binary(_payload1[0:1]+bson.encode(_decoded1), 6)
+            _decoded2["payloadId"] += 1
+            expected["v"]["$and"][1]["age"]["$lte"] = Binary(_payload2[0:1]+bson.encode(_decoded2), 6)
+            if encrypted_val == expected:
+                break
+    return expected
+
 class AsyncKeyVaultCallback(MockAsyncCallback):
     def __init__(self, kms_reply=None):
         super().__init__(kms_reply=kms_reply)
@@ -1184,12 +1203,13 @@ class TestExplicitEncryption(unittest.TestCase):
     maxDiff = None
 
     @staticmethod
-    def mongo_crypt_opts():
+    def mongo_crypt_opts(enable_range_v2=True):
         return MongoCryptOptions(
             {
                 "aws": {"accessKeyId": "example", "secretAccessKey": "example"},
                 "local": {"key": b"\x00" * 96},
-            }
+            },
+            enable_range_v2=enable_range_v2,
         )
 
     def _test_encrypt_decrypt(self, key_id=None, key_alt_name=None):
@@ -1392,13 +1412,33 @@ def test_range_query_int32(self):
             is_expression=True,
         )
         encrypted_val = bson.decode(encrypted, OPTS)
+        self.assertEqual(encrypted_val, adjust_range_counter(encrypted_val, expected))
 
-        # Workaround for the internal range payload counter in libmongocrypt
-        if encrypted_val != expected:
-            expected = json_data(
-                "fle2-find-range-explicit-v2/int32/encrypted-payload-second.json"
-            )
-        self.assertEqual(encrypted_val, expected)
+    def test_rangePreview_query_int32(self):
+        key_path = "keys/ABCDEFAB123498761234123456789012-local-document.json"
+        key_id = json_data(key_path)["_id"]
+        encrypter = ExplicitEncrypter(
+            MockCallback(
+                key_docs=[bson_data(key_path)], kms_reply=http_data("kms-reply.txt")
+            ),
+            self.mongo_crypt_opts(enable_range_v2=False),
+        )
+        self.addCleanup(encrypter.close)
+
+        range_opts = bson_data("fle2-find-rangePreview-explicit/int32/rangeopts.json")
+        value = bson_data("fle2-find-rangePreview-explicit/int32/value-to-encrypt.json")
+        expected = json_data("fle2-find-range-explicit-v2/int32/encrypted-payload.json")
+        encrypted = encrypter.encrypt(
+            value,
+            "rangePreview",
+            key_id=key_id,
+            query_type="rangePreview",
+            contention_factor=4,
+            range_opts=range_opts,
+            is_expression=True,
+        )
+        encrypted_val = bson.decode(encrypted, OPTS)
+        self.assertEqual(encrypted_val, adjust_range_counter(encrypted_val, expected))
 
 
 def read(filename, **kwargs):