PYTHON-3241 Add Queryable Encryption API to AutoEncryptionOpts (#957)

Author: juliusgeo

pymongo/collection.py

@@ -35,7 +35,7 @@
 from bson.raw_bson import RawBSONDocument
 from bson.son import SON
 from bson.timestamp import Timestamp
-from pymongo import common, helpers, message
+from pymongo import ASCENDING, common, helpers, message
 from pymongo.aggregation import (
     _CollectionAggregationCommand,
     _CollectionRawAggregationCommand,
@@ -44,6 +44,7 @@
 from pymongo.change_stream import CollectionChangeStream
 from pymongo.collation import validate_collation_or_none
 from pymongo.command_cursor import CommandCursor, RawBatchCommandCursor
+from pymongo.common import _ecc_coll_name, _ecoc_coll_name, _esc_coll_name
 from pymongo.cursor import Cursor, RawBatchCursor
 from pymongo.errors import (
     ConfigurationError,
@@ -115,6 +116,7 @@ def __init__(
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional["ReadConcern"] = None,
         session: Optional["ClientSession"] = None,
+        encrypted_fields: Optional[Mapping[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         """Get / create a Mongo collection.
@@ -197,7 +199,6 @@ def __init__(
             write_concern or database.write_concern,
             read_concern or database.read_concern,
         )
-
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
 
@@ -215,7 +216,16 @@ def __init__(
         self.__name = name
         self.__full_name = "%s.%s" % (self.__database.name, self.__name)
         if create or kwargs or collation:
-            self.__create(kwargs, collation, session)
+            if encrypted_fields:
+                common.validate_is_mapping("encrypted_fields", encrypted_fields)
+                opts = {"clusteredIndex": {"key": {"_id": 1}, "unique": True}}
+                self.__create(_esc_coll_name(encrypted_fields, name), opts, None, session)
+                self.__create(_ecc_coll_name(encrypted_fields, name), opts, None, session)
+                self.__create(_ecoc_coll_name(encrypted_fields, name), opts, None, session)
+                self.__create(name, kwargs, collation, session, encrypted_fields=encrypted_fields)
+                self.create_index([("__safeContent__", ASCENDING)], session)
+            else:
+                self.__create(name, kwargs, collation, session)
 
         self.__write_response_codec_options = self.codec_options._replace(
             unicode_decode_error_handler="replace", document_class=dict
@@ -286,9 +296,12 @@ def _command(
                 user_fields=user_fields,
             )
 
-    def __create(self, options, collation, session):
+    def __create(self, name, options, collation, session, encrypted_fields=None):
         """Sends a create command with the given options."""
-        cmd = SON([("create", self.__name)])
+        cmd = SON([("create", name)])
+        if encrypted_fields:
+            cmd["encryptedFields"] = encrypted_fields
+
         if options:
             if "size" in options:
                 options["size"] = float(options["size"])

pymongo/common.py

@@ -792,6 +792,18 @@ def get_validated_options(
     return validated_options
 
 
+def _esc_coll_name(encrypted_fields, name):
+    return encrypted_fields.get("escCollection", f"enxcol_.{name}.esc")
+
+
+def _ecc_coll_name(encrypted_fields, name):
+    return encrypted_fields.get("eccCollection", f"enxcol_.{name}.ecc")
+
+
+def _ecoc_coll_name(encrypted_fields, name):
+    return encrypted_fields.get("ecocCollection", f"enxcol_.{name}.ecoc")
+
+
 # List of write-concern-related options.
 WRITE_CONCERN_OPTIONS = frozenset(["w", "wtimeout", "wtimeoutms", "fsync", "j", "journal"])
 

pymongo/database.py

@@ -38,6 +38,7 @@
 from pymongo.change_stream import DatabaseChangeStream
 from pymongo.collection import Collection
 from pymongo.command_cursor import CommandCursor
+from pymongo.common import _ecc_coll_name, _ecoc_coll_name, _esc_coll_name
 from pymongo.errors import CollectionInvalid, InvalidName
 from pymongo.read_preferences import ReadPreference, _ServerMode
 from pymongo.typings import _CollationIn, _DocumentType, _Pipeline
@@ -290,6 +291,7 @@ def create_collection(
         write_concern: Optional["WriteConcern"] = None,
         read_concern: Optional["ReadConcern"] = None,
         session: Optional["ClientSession"] = None,
+        encrypted_fields: Optional[Mapping[str, Any]] = None,
         **kwargs: Any,
     ) -> Collection[_DocumentType]:
         """Create a new :class:`~pymongo.collection.Collection` in this
@@ -321,6 +323,29 @@ def create_collection(
             :class:`~pymongo.collation.Collation`.
           - `session` (optional): a
             :class:`~pymongo.client_session.ClientSession`.
+          - `encrypted_fields`: Document that describes the encrypted fields for Queryable
+            Encryption.
+            For example::
+
+                {
+                  "escCollection": "enxcol_.encryptedCollection.esc",
+                  "eccCollection": "enxcol_.encryptedCollection.ecc",
+                  "ecocCollection": "enxcol_.encryptedCollection.ecoc",
+                  "fields": [
+                      {
+                          "path": "firstName",
+                          "keyId": Binary.from_uuid(UUID('00000000-0000-0000-0000-000000000000')),
+                          "bsonType": "string",
+                          "queries": {"queryType": "equality"}
+                      },
+                      {
+                          "path": "ssn",
+                          "keyId": Binary.from_uuid(UUID('04104104-1041-0410-4104-104104104104')),
+                          "bsonType": "string"
+                      }
+                  ]
+
+                }                }
           - `**kwargs` (optional): additional keyword arguments will
             be passed as options for the `create collection command`_
 
@@ -369,14 +394,24 @@ def create_collection(
         .. _create collection command:
             https://mongodb.com/docs/manual/reference/command/create
         """
+        if (
+            not encrypted_fields
+            and self.client.options.auto_encryption_opts
+            and self.client.options.auto_encryption_opts._encrypted_fields_map
+        ):
+            encrypted_fields = self.client.options.auto_encryption_opts._encrypted_fields_map.get(
+                "%s.%s" % (self.name, name)
+            )
+        if encrypted_fields:
+            common.validate_is_mapping("encrypted_fields", encrypted_fields)
+
         with self.__client._tmp_session(session) as s:
             # Skip this check in a transaction where listCollections is not
             # supported.
             if (not s or not s.in_transaction) and name in self.list_collection_names(
                 filter={"name": name}, session=s
             ):
                 raise CollectionInvalid("collection %s already exists" % name)
-
             return Collection(
                 self,
                 name,
@@ -386,6 +421,7 @@ def create_collection(
                 write_concern,
                 read_concern,
                 session=s,
+                encrypted_fields=encrypted_fields,
                 **kwargs,
             )
 
@@ -874,11 +910,27 @@ def list_collection_names(
 
         return [result["name"] for result in self.list_collections(session=session, **kwargs)]
 
+    def _drop_helper(self, name, session=None, comment=None):
+        command = SON([("drop", name)])
+        if comment is not None:
+            command["comment"] = comment
+
+        with self.__client._socket_for_writes(session) as sock_info:
+            return self._command(
+                sock_info,
+                command,
+                allowable_errors=["ns not found", 26],
+                write_concern=self._write_concern_for(session),
+                parse_write_concern_error=True,
+                session=session,
+            )
+
     def drop_collection(
         self,
         name_or_collection: Union[str, Collection],
         session: Optional["ClientSession"] = None,
         comment: Optional[Any] = None,
+        encrypted_fields: Optional[Mapping[str, Any]] = None,
     ) -> Dict[str, Any]:
         """Drop a collection.
 
@@ -889,6 +941,29 @@ def drop_collection(
             :class:`~pymongo.client_session.ClientSession`.
           - `comment` (optional): A user-provided comment to attach to this
             command.
+          - `encrypted_fields`: Document that describes the encrypted fields for Queryable
+            Encryption.
+            For example::
+
+                {
+                  "escCollection": "enxcol_.encryptedCollection.esc",
+                  "eccCollection": "enxcol_.encryptedCollection.ecc",
+                  "ecocCollection": "enxcol_.encryptedCollection.ecoc",
+                  "fields": [
+                      {
+                          "path": "firstName",
+                          "keyId": Binary.from_uuid(UUID('00000000-0000-0000-0000-000000000000')),
+                          "bsonType": "string",
+                          "queries": {"queryType": "equality"}
+                      },
+                      {
+                          "path": "ssn",
+                          "keyId": Binary.from_uuid(UUID('04104104-1041-0410-4104-104104104104')),
+                          "bsonType": "string"
+                      }
+                  ]
+
+                }
 
 
         .. note:: The :attr:`~pymongo.database.Database.write_concern` of
@@ -911,20 +986,34 @@ def drop_collection(
 
         if not isinstance(name, str):
             raise TypeError("name_or_collection must be an instance of str")
-
-        command = SON([("drop", name)])
-        if comment is not None:
-            command["comment"] = comment
-
-        with self.__client._socket_for_writes(session) as sock_info:
-            return self._command(
-                sock_info,
-                command,
-                allowable_errors=["ns not found", 26],
-                write_concern=self._write_concern_for(session),
-                parse_write_concern_error=True,
-                session=session,
+        full_name = "%s.%s" % (self.name, name)
+        if (
+            not encrypted_fields
+            and self.client.options.auto_encryption_opts
+            and self.client.options.auto_encryption_opts._encrypted_fields_map
+        ):
+            encrypted_fields = self.client.options.auto_encryption_opts._encrypted_fields_map.get(
+                full_name
+            )
+        if not encrypted_fields and self.client.options.auto_encryption_opts:
+            colls = list(
+                self.list_collections(filter={"name": name}, session=session, comment=comment)
             )
+            if colls and colls[0]["options"].get("encryptedFields"):
+                encrypted_fields = colls[0]["options"]["encryptedFields"]
+        if encrypted_fields:
+            common.validate_is_mapping("encrypted_fields", encrypted_fields)
+            self._drop_helper(
+                _esc_coll_name(encrypted_fields, name), session=session, comment=comment
+            )
+            self._drop_helper(
+                _ecc_coll_name(encrypted_fields, name), session=session, comment=comment
+            )
+            self._drop_helper(
+                _ecoc_coll_name(encrypted_fields, name), session=session, comment=comment
+            )
+
+        return self._drop_helper(name, session, comment)
 
     def validate_collection(
         self,

pymongo/encryption.py

@@ -264,6 +264,11 @@ def __init__(self, client, opts):
             schema_map = None
         else:
             schema_map = _dict_to_bson(opts._schema_map, False, _DATA_KEY_OPTS)
+
+        if opts._encrypted_fields_map is None:
+            encrypted_fields_map = None
+        else:
+            encrypted_fields_map = _dict_to_bson(opts._encrypted_fields_map, False, _DATA_KEY_OPTS)
         self._bypass_auto_encryption = opts._bypass_auto_encryption
         self._internal_client = None
 
@@ -304,6 +309,7 @@ def _get_internal_client(encrypter, mongo_client):
                 crypt_shared_lib_path=opts._crypt_shared_lib_path,
                 crypt_shared_lib_required=opts._crypt_shared_lib_required,
                 bypass_encryption=opts._bypass_auto_encryption,
+                encrypted_fields_map=encrypted_fields_map,
                 bypass_query_analysis=opts._bypass_query_analysis,
             ),
         )

pymongo/encryption_options.py

@@ -23,6 +23,7 @@
 except ImportError:
     _HAVE_PYMONGOCRYPT = False
 
+from pymongo.common import validate_is_mapping
 from pymongo.errors import ConfigurationError
 from pymongo.uri_parser import _parse_kms_tls_options
 
@@ -48,6 +49,7 @@ def __init__(
         crypt_shared_lib_path: Optional[str] = None,
         crypt_shared_lib_required: bool = False,
         bypass_query_analysis: bool = False,
+        encrypted_fields_map: Optional[Mapping] = None,
     ) -> None:
         """Options to configure automatic client-side field level encryption.
 
@@ -150,10 +152,33 @@ def __init__(
             outgoing commands. Set `bypass_query_analysis` to use explicit
             encryption on indexed fields without the MongoDB Enterprise Advanced
             licensed crypt_shared library.
+          - `encrypted_fields_map`: Map of collection namespace ("db.coll") to documents that
+            described the encrypted fields for Queryable Encryption. For example::
+
+                {
+                  "db.encryptedCollection": {
+                      "escCollection": "enxcol_.encryptedCollection.esc",
+                      "eccCollection": "enxcol_.encryptedCollection.ecc",
+                      "ecocCollection": "enxcol_.encryptedCollection.ecoc",
+                      "fields": [
+                          {
+                              "path": "firstName",
+                              "keyId": Binary.from_uuid(UUID('00000000-0000-0000-0000-000000000000')),
+                              "bsonType": "string",
+                              "queries": {"queryType": "equality"}
+                          },
+                          {
+                              "path": "ssn",
+                              "keyId": Binary.from_uuid(UUID('04104104-1041-0410-4104-104104104104')),
+                              "bsonType": "string"
+                          }
+                      ]
+                  }
+                }
 
         .. versionchanged:: 4.2
-           Added `crypt_shared_lib_path`, `crypt_shared_lib_required`, and `bypass_query_analysis`
-           parameters.
+           Added `encrypted_fields_map` `crypt_shared_lib_path`, `crypt_shared_lib_required`,
+           and `bypass_query_analysis` parameters.
 
         .. versionchanged:: 4.0
            Added the `kms_tls_options` parameter and the "kmip" KMS provider.
@@ -166,6 +191,10 @@ def __init__(
                 "install a compatible version with: "
                 "python -m pip install 'pymongo[encryption]'"
             )
+        if encrypted_fields_map:
+            validate_is_mapping("encrypted_fields_map", encrypted_fields_map)
+        self._encrypted_fields_map = encrypted_fields_map
+        self._bypass_query_analysis = bypass_query_analysis
         self._crypt_shared_lib_path = crypt_shared_lib_path
         self._crypt_shared_lib_required = crypt_shared_lib_required
         self._kms_providers = kms_providers