PYTHON-1884 Support auto encryption in bulk write

Close KMS sockets.
Call pymongocrypt's init method.

Author: ShaneHarvey

pymongo/bulk.py

@@ -36,9 +36,9 @@
                             OperationFailure)
 from pymongo.message import (_INSERT, _UPDATE, _DELETE,
                              _do_batched_insert,
-                             _do_bulk_write_command,
                              _randint,
-                             _BulkWriteContext)
+                             _BulkWriteContext,
+                             _EncryptedBulkWriteContext)
 from pymongo.read_preferences import ReadPreference
 from pymongo.write_concern import WriteConcern
 
@@ -152,8 +152,6 @@ def __init__(self, collection, ordered, bypass_document_validation):
                 document_class=dict))
         self.ordered = ordered
         self.ops = []
-        self.name = "%s.%s" % (collection.database.name, collection.name)
-        self.namespace = collection.database.name + '.$cmd'
         self.executed = False
         self.bypass_doc_val = bypass_document_validation
         self.uses_collation = False
@@ -164,6 +162,14 @@ def __init__(self, collection, ordered, bypass_document_validation):
         # Extra state so that we know where to pick up on a retry attempt.
         self.current_run = None
 
+    @property
+    def bulk_ctx_class(self):
+        encrypter = self.collection.database.client._encrypter
+        if encrypter and not encrypter._bypass_auto_encryption:
+            return _EncryptedBulkWriteContext
+        else:
+            return _BulkWriteContext
+
     def add_insert(self, document):
         """Add an insert document to the list of ops.
         """
@@ -271,8 +277,9 @@ def _execute_command(self, generator, write_concern, session,
                 cmd['writeConcern'] = write_concern.document
             if self.bypass_doc_val and sock_info.max_wire_version >= 4:
                 cmd['bypassDocumentValidation'] = True
-            bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id,
-                                    listeners, session)
+            bwc = self.bulk_ctx_class(
+                db_name, cmd, sock_info, op_id, listeners, session,
+                run.op_type, self.collection.codec_options)
 
             while run.idx_offset < len(run.ops):
                 if session:
@@ -283,16 +290,9 @@ def _execute_command(self, generator, write_concern, session,
                         self.started_retryable_write = True
                     session._apply_to(cmd, retryable, ReadPreference.PRIMARY)
                 sock_info.send_cluster_time(cmd, session, client)
-                check_keys = run.op_type == _INSERT
                 ops = islice(run.ops, run.idx_offset, None)
-                # Run as many ops as possible.
-                request_id, msg, to_send = _do_bulk_write_command(
-                    self.namespace, run.op_type, cmd, ops, check_keys,
-                    self.collection.codec_options, bwc)
-                if not to_send:
-                    raise InvalidOperation("cannot do an empty bulk write")
-                result = bwc.write_command(request_id, msg, to_send)
-                client._process_response(result, session)
+                # Run as many ops as possible in one command.
+                result, to_send = bwc.execute(ops, client)
 
                 # Retryable writeConcernErrors halt the execution of this run.
                 wce = result.get('writeConcernError', {})
@@ -361,7 +361,7 @@ def execute_insert_no_results(self, sock_info, run, op_id, acknowledged):
         db = self.collection.database
         bwc = _BulkWriteContext(
             db.name, command, sock_info, op_id, db.client._event_listeners,
-            session=None)
+            None, _INSERT, self.collection.codec_options)
         # Legacy batched OP_INSERT.
         _do_batched_insert(
             self.collection.full_name, run.ops, True, acknowledged, concern,
@@ -383,25 +383,15 @@ def execute_op_msg_no_results(self, sock_info, generator):
             cmd = SON([(_COMMANDS[run.op_type], self.collection.name),
                        ('ordered', False),
                        ('writeConcern', {'w': 0})])
-            bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id,
-                                    listeners, None)
+            bwc = self.bulk_ctx_class(
+                db_name, cmd, sock_info, op_id, listeners, None,
+                run.op_type, self.collection.codec_options)
 
             while run.idx_offset < len(run.ops):
-                check_keys = run.op_type == _INSERT
                 ops = islice(run.ops, run.idx_offset, None)
                 # Run as many ops as possible.
-                request_id, msg, to_send = _do_bulk_write_command(
-                    self.namespace, run.op_type, cmd, ops, check_keys,
-                    self.collection.codec_options, bwc)
-                if not to_send:
-                    raise InvalidOperation("cannot do an empty bulk write")
+                to_send = bwc.execute_unack(ops, client)
                 run.idx_offset += len(to_send)
-                # Though this isn't strictly a "legacy" write, the helper
-                # handles publishing commands and sending our message
-                # without receiving a result. Send 0 for max_doc_size
-                # to disable size checking. Size checking is handled while
-                # the documents are encoded to BSON.
-                bwc.legacy_write(request_id, msg, 0, False, to_send)
             self.current_run = run = next(generator, None)
 
     def execute_command_no_results(self, sock_info, generator):

pymongo/encryption.py

@@ -74,15 +74,14 @@ def kms_request(self, kms_context):
                            socket_timeout=_KMS_CONNECT_TIMEOUT,
                            ssl_context=ctx)
         try:
-            conn = _configured_socket((endpoint, _HTTPS_PORT), opts)
-            conn.sendall(message)
+            with _configured_socket((endpoint, _HTTPS_PORT), opts) as conn:
+                conn.sendall(message)
+                while kms_context.bytes_needed > 0:
+                    data = conn.recv(kms_context.bytes_needed)
+                    kms_context.feed(data)
         except Exception as exc:
             raise MongoCryptError(str(exc))
 
-        while kms_context.bytes_needed > 0:
-            data = conn.recv(kms_context.bytes_needed)
-            kms_context.feed(data)
-
     def collection_info(self, database, filter):
         """Get the collection info for a namespace.
 

pymongo/message.py

@@ -29,6 +29,7 @@
                   _dict_to_bson,
                   _make_c_string)
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
+from bson.raw_bson import _inflate_bson, DEFAULT_RAW_BSON_OPTIONS
 from bson.py3compat import b, StringIO
 from bson.son import SON
 
@@ -47,6 +48,7 @@
                             ProtocolError)
 from pymongo.read_concern import DEFAULT_READ_CONCERN
 from pymongo.read_preferences import ReadPreference
+from pymongo.write_concern import WriteConcern
 
 
 MAX_INT32 = 2147483647
@@ -862,10 +864,10 @@ class _BulkWriteContext(object):
 
     __slots__ = ('db_name', 'command', 'sock_info', 'op_id',
                  'name', 'field', 'publish', 'start_time', 'listeners',
-                 'session', 'compress')
+                 'session', 'compress', 'op_type', 'codec')
 
     def __init__(self, database_name, command, sock_info, operation_id,
-                 listeners, session):
+                 listeners, session, op_type, codec):
         self.db_name = database_name
         self.command = command
         self.sock_info = sock_info
@@ -877,6 +879,38 @@ def __init__(self, database_name, command, sock_info, operation_id,
         self.start_time = datetime.datetime.now() if self.publish else None
         self.session = session
         self.compress = True if sock_info.compression_context else False
+        self.op_type = op_type
+        self.codec = codec
+
+    def _batch_command(self, docs):
+        namespace = self.db_name + '.$cmd'
+        request_id, msg, to_send = _do_bulk_write_command(
+            namespace, self.op_type, self.command, docs, self.check_keys,
+            self.codec, self)
+        if not to_send:
+            raise InvalidOperation("cannot do an empty bulk write")
+        return request_id, msg, to_send
+
+    def execute(self, docs, client):
+        request_id, msg, to_send = self._batch_command(docs)
+        result = self.write_command(request_id, msg, to_send)
+        client._process_response(result, self.session)
+        return result, to_send
+
+    def execute_unack(self, docs, client):
+        request_id, msg, to_send = self._batch_command(docs)
+        # Though this isn't strictly a "legacy" write, the helper
+        # handles publishing commands and sending our message
+        # without receiving a result. Send 0 for max_doc_size
+        # to disable size checking. Size checking is handled while
+        # the documents are encoded to BSON.
+        self.legacy_write(request_id, msg, 0, False, to_send)
+        return to_send
+
+    @property
+    def check_keys(self):
+        """Should we check keys for this operation type?"""
+        return self.op_type == _INSERT
 
     @property
     def max_bson_size(self):
@@ -975,6 +1009,54 @@ def _fail(self, request_id, failure, duration):
             request_id, self.sock_info.address, self.op_id)
 
 
+# 2MiB
+_MAX_ENC_BSON_SIZE = 2 * (1024 * 1024)
+# 6MB
+_MAX_ENC_MESSAGE_SIZE = 6 * (1000 * 1000)
+
+
+class _EncryptedBulkWriteContext(_BulkWriteContext):
+    __slots__ = ()
+
+    def _batch_command(self, docs):
+        namespace = self.db_name + '.$cmd'
+        msg, to_send = _encode_batched_write_command(
+            namespace, self.op_type, self.command, docs, self.check_keys,
+            self.codec, self)
+        if not to_send:
+            raise InvalidOperation("cannot do an empty bulk write")
+
+        # Chop off the OP_QUERY header to get a properly batched write command.
+        cmd_start = msg.index(b"\x00", 4) + 9
+        cmd = _inflate_bson(memoryview(msg)[cmd_start:],
+                            DEFAULT_RAW_BSON_OPTIONS)
+        return cmd, to_send
+
+    def execute(self, docs, client):
+        cmd, to_send = self._batch_command(docs)
+        result = self.sock_info.command(
+            self.db_name, cmd, codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
+            session=self.session, client=client)
+        return result, to_send
+
+    def execute_unack(self, docs, client):
+        cmd, to_send = self._batch_command(docs)
+        self.sock_info.command(
+            self.db_name, cmd, write_concern=WriteConcern(w=0),
+            session=self.session, client=client)
+        return to_send
+
+    @property
+    def max_bson_size(self):
+        """A proxy for SockInfo.max_bson_size."""
+        return min(self.sock_info.max_bson_size, _MAX_ENC_BSON_SIZE)
+
+    @property
+    def max_message_size(self):
+        """A proxy for SockInfo.max_message_size."""
+        return min(self.sock_info.max_message_size, _MAX_ENC_MESSAGE_SIZE)
+
+
 def _raise_document_too_large(operation, doc_size, max_size):
     """Internal helper for raising DocumentTooLarge."""
     if operation == "insert":

test/test_encryption.py

@@ -30,8 +30,16 @@
 from pymongo.errors import ConfigurationError
 from pymongo.mongo_client import MongoClient
 from pymongo.encryption_options import AutoEncryptionOpts, _HAVE_PYMONGOCRYPT
+from pymongo.write_concern import WriteConcern
 
 from test import unittest, IntegrationTest, PyMongoTestCase, client_context
+from test.utils import wait_until
+
+
+if _HAVE_PYMONGOCRYPT:
+    # Load the mongocrypt library.
+    from pymongocrypt.binding import init
+    init(os.environ.get('MONGOCRYPT_LIB', 'mongocrypt'))
 
 
 def get_client_opts(client):
@@ -156,19 +164,27 @@ def _test_auto_encrypt(self, opts):
         key_vault.insert_one(data_key)
         self.addCleanup(key_vault.drop)
 
-        # Collection.insert_one auto encrypts.
-        docs = [{'_id': 1, 'ssn': '123'},
-                {'_id': 2, 'ssn': '456'},
-                {'_id': 3, 'ssn': '789'}]
+        # Collection.insert_one/insert_many auto encrypts.
+        docs = [{'_id': 0, 'ssn': '000'},
+                {'_id': 1, 'ssn': '111'},
+                {'_id': 2, 'ssn': '222'},
+                {'_id': 3, 'ssn': '333'},
+                {'_id': 4, 'ssn': '444'},
+                {'_id': 5, 'ssn': '555'}]
         encrypted_coll = client.pymongo_test.test
-        for doc in docs:
-            encrypted_coll.insert_one(doc)
+        encrypted_coll.insert_one(docs[0])
+        encrypted_coll.insert_many(docs[1:3])
+        unack = encrypted_coll.with_options(write_concern=WriteConcern(w=0))
+        unack.insert_one(docs[3])
+        unack.insert_many(docs[4:], ordered=False)
+        wait_until(lambda: self.db.test.count_documents({}) == len(docs),
+                   'insert documents with w=0')
 
         # Database.command auto decrypts.
         res = client.pymongo_test.command(
-            'find', 'test', filter={'ssn': '123'})
+            'find', 'test', filter={'ssn': '000'})
         decrypted_docs = res['cursor']['firstBatch']
-        self.assertEqual(decrypted_docs, [{'_id': 1, 'ssn': '123'}])
+        self.assertEqual(decrypted_docs, [{'_id': 0, 'ssn': '000'}])
 
         # Collection.find auto decrypts.
         decrypted_docs = list(encrypted_coll.find())
@@ -188,13 +204,13 @@ def _test_auto_encrypt(self, opts):
 
         # Collection.distinct auto decrypts.
         decrypted_ssns = encrypted_coll.distinct('ssn')
-        self.assertEqual(decrypted_ssns, ['123', '456', '789'])
+        self.assertEqual(decrypted_ssns, [d['ssn'] for d in docs])
 
         # Make sure the field is actually encrypted.
-        encrypted_doc = self.db.test.find_one()
-        self.assertEqual(encrypted_doc['_id'], 1)
-        self.assertIsInstance(encrypted_doc['ssn'], Binary)
-        self.assertEqual(encrypted_doc['ssn'].subtype, 6)
+        for encrypted_doc in self.db.test.find():
+            self.assertIsInstance(encrypted_doc['_id'], int)
+            self.assertIsInstance(encrypted_doc['ssn'], Binary)
+            self.assertEqual(encrypted_doc['ssn'].subtype, 6)
 
     def test_auto_encrypt(self):
         # Configure the encrypted field via jsonSchema.