PYTHON-3280 Support for Range Indexes (#1140)

Author: juliusgeo

pymongo/encryption.py

@@ -41,7 +41,7 @@
 from pymongo import _csot
 from pymongo.cursor import Cursor
 from pymongo.daemon import _spawn_daemon
-from pymongo.encryption_options import AutoEncryptionOpts
+from pymongo.encryption_options import AutoEncryptionOpts, RangeOpts
 from pymongo.errors import (
     ConfigurationError,
     EncryptionError,
@@ -416,6 +416,14 @@ class Algorithm(str, enum.Enum):
 
     .. versionadded:: 4.2
     """
+    RANGEPREVIEW = "RangePreview"
+    """RangePreview.
+
+    .. note:: Support for Range queries is in beta.
+       Backwards-breaking changes may be made before the final release.
+
+    .. versionadded:: 4.4
+    """
 
 
 class QueryType(str, enum.Enum):
@@ -430,6 +438,9 @@ class QueryType(str, enum.Enum):
     EQUALITY = "equality"
     """Used to encrypt a value for an equality query."""
 
+    RANGEPREVIEW = "rangePreview"
+    """Used to encrypt a value for a range query."""
+
 
 class ClientEncryption(Generic[_DocumentType]):
     """Explicit client-side field level encryption."""
@@ -627,6 +638,45 @@ def create_data_key(
                 key_material=key_material,
             )
 
+    def _encrypt_helper(
+        self,
+        value,
+        algorithm,
+        key_id=None,
+        key_alt_name=None,
+        query_type=None,
+        contention_factor=None,
+        range_opts=None,
+        is_expression=False,
+    ):
+        self._check_closed()
+        if key_id is not None and not (
+            isinstance(key_id, Binary) and key_id.subtype == UUID_SUBTYPE
+        ):
+            raise TypeError("key_id must be a bson.binary.Binary with subtype 4")
+
+        doc = encode(
+            {"v": value},
+            codec_options=self._codec_options,
+        )
+        if range_opts:
+            range_opts = encode(
+                range_opts.document,
+                codec_options=self._codec_options,
+            )
+        with _wrap_encryption_errors():
+            encrypted_doc = self._encryption.encrypt(
+                value=doc,
+                algorithm=algorithm,
+                key_id=key_id,
+                key_alt_name=key_alt_name,
+                query_type=query_type,
+                contention_factor=contention_factor,
+                range_opts=range_opts,
+                is_expression=is_expression,
+            )
+            return decode(encrypted_doc)["v"]  # type: ignore[index]
+
     def encrypt(
         self,
         value: Any,
@@ -635,6 +685,7 @@ def encrypt(
         key_alt_name: Optional[str] = None,
         query_type: Optional[str] = None,
         contention_factor: Optional[int] = None,
+        range_opts: Optional[RangeOpts] = None,
     ) -> Binary:
         """Encrypt a BSON value with a given key and algorithm.
 
@@ -655,10 +706,10 @@ def encrypt(
             when the algorithm is :attr:`Algorithm.INDEXED`.  An integer value
             *must* be given when the :attr:`Algorithm.INDEXED` algorithm is
             used.
+          - `range_opts`: **(BETA)** An instance of RangeOpts.
 
-        .. note:: `query_type` and `contention_factor` are part of the
-           Queryable Encryption beta. Backwards-breaking changes may be made before the
-           final release.
+        .. note:: `query_type`, `contention_factor` and `range_opts` are part of the Queryable Encryption beta.
+           Backwards-breaking changes may be made before the final release.
 
         :Returns:
           The encrypted value, a :class:`~bson.binary.Binary` with subtype 6.
@@ -667,23 +718,66 @@ def encrypt(
            Added the `query_type` and `contention_factor` parameters.
 
         """
-        self._check_closed()
-        if key_id is not None and not (
-            isinstance(key_id, Binary) and key_id.subtype == UUID_SUBTYPE
-        ):
-            raise TypeError("key_id must be a bson.binary.Binary with subtype 4")
+        return self._encrypt_helper(
+            value=value,
+            algorithm=algorithm,
+            key_id=key_id,
+            key_alt_name=key_alt_name,
+            query_type=query_type,
+            contention_factor=contention_factor,
+            range_opts=range_opts,
+            is_expression=False,
+        )
 
-        doc = encode({"v": value}, codec_options=self._codec_options)
-        with _wrap_encryption_errors():
-            encrypted_doc = self._encryption.encrypt(
-                doc,
-                algorithm,
-                key_id=key_id,
-                key_alt_name=key_alt_name,
-                query_type=query_type,
-                contention_factor=contention_factor,
-            )
-            return decode(encrypted_doc)["v"]  # type: ignore[index]
+    def encrypt_expression(
+        self,
+        expression: Mapping[str, Any],
+        algorithm: str,
+        key_id: Optional[Binary] = None,
+        key_alt_name: Optional[str] = None,
+        query_type: Optional[str] = None,
+        contention_factor: Optional[int] = None,
+        range_opts: Optional[RangeOpts] = None,
+    ) -> RawBSONDocument:
+        """Encrypt a BSON expression with a given key and algorithm.
+
+        Note that exactly one of ``key_id`` or  ``key_alt_name`` must be
+        provided.
+
+        :Parameters:
+          - `expression`: **(BETA)** The BSON aggregate or match expression to encrypt.
+          - `algorithm` (string): The encryption algorithm to use. See
+            :class:`Algorithm` for some valid options.
+          - `key_id`: Identifies a data key by ``_id`` which must be a
+            :class:`~bson.binary.Binary` with subtype 4 (
+            :attr:`~bson.binary.UUID_SUBTYPE`).
+          - `key_alt_name`: Identifies a key vault document by 'keyAltName'.
+          - `query_type` (str): **(BETA)** The query type to execute. See
+            :class:`QueryType` for valid options.
+          - `contention_factor` (int): **(BETA)** The contention factor to use
+            when the algorithm is :attr:`Algorithm.INDEXED`.  An integer value
+            *must* be given when the :attr:`Algorithm.INDEXED` algorithm is
+            used.
+          - `range_opts`: **(BETA)** An instance of RangeOpts.
+
+        .. note:: Support for range queries is in beta.
+           Backwards-breaking changes may be made before the final release.
+
+        :Returns:
+          The encrypted expression, a :class:`~bson.RawBSONDocument`.
+
+        .. versionadded:: 4.4
+        """
+        return self._encrypt_helper(
+            value=expression,
+            algorithm=algorithm,
+            key_id=key_id,
+            key_alt_name=key_alt_name,
+            query_type=query_type,
+            contention_factor=contention_factor,
+            range_opts=range_opts,
+            is_expression=True,
+        )
 
     def decrypt(self, value: Binary) -> Any:
         """Decrypt an encrypted value.

pymongo/encryption_options.py

@@ -22,7 +22,7 @@
     _HAVE_PYMONGOCRYPT = True
 except ImportError:
     _HAVE_PYMONGOCRYPT = False
-
+from bson import int64
 from pymongo.common import validate_is_mapping
 from pymongo.errors import ConfigurationError
 from pymongo.uri_parser import _parse_kms_tls_options
@@ -219,3 +219,45 @@ def __init__(
         # Maps KMS provider name to a SSLContext.
         self._kms_ssl_contexts = _parse_kms_tls_options(kms_tls_options)
         self._bypass_query_analysis = bypass_query_analysis
+
+
+class RangeOpts:
+    """Options to configure encrypted queries using the rangePreview algorithm."""
+
+    def __init__(
+        self,
+        sparsity: int,
+        min: Optional[Any] = None,
+        max: Optional[Any] = None,
+        precision: Optional[int] = None,
+    ) -> None:
+        """Options to configure encrypted queries using the rangePreview algorithm.
+
+        .. note:: Support for Range queries is in beta.
+           Backwards-breaking changes may be made before the final release.
+
+        :Parameters:
+          - `sparsity`: An integer.
+          - `min`: A BSON scalar value corresponding to the type being queried.
+          - `max`: A BSON scalar value corresponding to the type being queried.
+          - `precision`: An integer, may only be set for double or decimal128 types.
+
+        .. versionadded:: 4.4
+        """
+        self.min = min
+        self.max = max
+        self.sparsity = sparsity
+        self.precision = precision
+
+    @property
+    def document(self) -> Mapping[str, Any]:
+        doc = {}
+        for k, v in [
+            ("sparsity", int64.Int64(self.sparsity)),
+            ("precision", self.precision),
+            ("min", self.min),
+            ("max", self.max),
+        ]:
+            if v is not None:
+                doc[k] = v
+        return doc

test/client-side-encryption/etc/data/encryptedFields-Range-Date.json

@@ -0,0 +1,36 @@
+{
+    "escCollection": "enxcol_.default.esc",
+    "eccCollection": "enxcol_.default.ecc",
+    "ecocCollection": "enxcol_.default.ecoc",
+    "fields": [
+        {
+            "keyId": {
+                "$binary": {
+                    "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+                    "subType": "04"
+                }
+            },
+            "path": "encryptedDate",
+            "bsonType": "date",
+            "queries": {
+                "queryType": "rangePreview",
+                "contention": {
+                    "$numberLong": "0"
+                },
+                "sparsity": {
+                    "$numberLong": "1"
+                },
+                "min": {
+                    "$date": {
+                        "$numberLong": "0"
+                    }
+                },
+                "max": {
+                    "$date": {
+                        "$numberLong": "200"
+                    }
+                }
+            }
+        }
+    ]
+}

test/client-side-encryption/etc/data/encryptedFields-Range-Decimal.json

@@ -0,0 +1,26 @@
+{
+  "escCollection": "enxcol_.default.esc",
+  "eccCollection": "enxcol_.default.ecc",
+  "ecocCollection": "enxcol_.default.ecoc",
+  "fields": [
+    {
+      "keyId": {
+        "$binary": {
+          "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+          "subType": "04"
+        }
+      },
+      "path": "encryptedDecimal",
+      "bsonType": "decimal",
+      "queries": {
+        "queryType": "rangePreview",
+        "contention": {
+          "$numberLong": "0"
+        },
+        "sparsity": {
+          "$numberLong": "1"
+        }
+      }
+    }
+  ]
+}

test/client-side-encryption/etc/data/encryptedFields-Range-DecimalPrecision.json

@@ -0,0 +1,35 @@
+{
+  "escCollection": "enxcol_.default.esc",
+  "eccCollection": "enxcol_.default.ecc",
+  "ecocCollection": "enxcol_.default.ecoc",
+  "fields": [
+    {
+      "keyId": {
+        "$binary": {
+          "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+          "subType": "04"
+        }
+      },
+      "path": "encryptedDecimalPrecision",
+      "bsonType": "decimal",
+      "queries": {
+        "queryType": "rangePreview",
+        "contention": {
+          "$numberLong": "0"
+        },
+        "sparsity": {
+          "$numberLong": "1"
+        },
+        "min": {
+          "$numberDecimal": "0.0"
+        },
+        "max": {
+          "$numberDecimal": "200.0"
+        },
+        "precision": {
+          "$numberInt": "2"
+        }
+      }
+    }
+  ]
+}

test/client-side-encryption/etc/data/encryptedFields-Range-Double.json

@@ -0,0 +1,26 @@
+{
+  "escCollection": "enxcol_.default.esc",
+  "eccCollection": "enxcol_.default.ecc",
+  "ecocCollection": "enxcol_.default.ecoc",
+  "fields": [
+    {
+      "keyId": {
+        "$binary": {
+          "base64": "EjRWeBI0mHYSNBI0VniQEg==",
+          "subType": "04"
+        }
+      },
+      "path": "encryptedDouble",
+      "bsonType": "double",
+      "queries": {
+        "queryType": "rangePreview",
+        "contention": {
+          "$numberLong": "0"
+        },
+        "sparsity": {
+          "$numberLong": "1"
+        }
+      }
+    }
+  ]
+}