Skip to content

Commit be6822b

Browse files
prashantmitalprashantmital
committed
PYTHON-2472 add a metadataClient for CSFLE (#539)
(cherry picked from commit 3e97712)
1 parent 2814d59 commit be6822b

31 files changed

+307
-475
lines changed

‎pymongo/encryption.py‎

Lines changed: 43 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -247,23 +247,57 @@ def close(self):
247247

248248

249249
class _Encrypter(object):
250-
def __init__(self, io_callbacks, opts):
251-
"""Encrypts and decrypts MongoDB commands.
250+
"""Encrypts and decrypts MongoDB commands.
252251
253-
This class is used to support automatic encryption and decryption of
254-
MongoDB commands.
252+
This class is used to support automatic encryption and decryption of
253+
MongoDB commands."""
254+
def __init__(self, client, opts):
255+
"""Create a _Encrypter for a client.
255256
256257
:Parameters:
257-
- `io_callbacks`: A :class:`MongoCryptCallback`.
258+
- `client`: The encrypted MongoClient.
258259
- `opts`: The encrypted client's :class:`AutoEncryptionOpts`.
259260
"""
260261
if opts._schema_map is None:
261262
schema_map = None
262263
else:
263264
schema_map = _dict_to_bson(opts._schema_map, False, _DATA_KEY_OPTS)
265+
self._bypass_auto_encryption = opts._bypass_auto_encryption
266+
self._internal_client = None
267+
268+
def _get_internal_client(encrypter, mongo_client):
269+
if mongo_client.max_pool_size is None:
270+
# Unlimited pool size, use the same client.
271+
return mongo_client
272+
# Else - limited pool size, use an internal client.
273+
if encrypter._internal_client is not None:
274+
return encrypter._internal_client
275+
internal_client = mongo_client._duplicate(
276+
minPoolSize=0, auto_encryption_opts=None)
277+
encrypter._internal_client = internal_client
278+
return internal_client
279+
280+
if opts._key_vault_client is not None:
281+
key_vault_client = opts._key_vault_client
282+
else:
283+
key_vault_client = _get_internal_client(self, client)
284+
285+
if opts._bypass_auto_encryption:
286+
metadata_client = None
287+
else:
288+
metadata_client = _get_internal_client(self, client)
289+
290+
db, coll = opts._key_vault_namespace.split('.', 1)
291+
key_vault_coll = key_vault_client[db][coll]
292+
293+
mongocryptd_client = MongoClient(
294+
opts._mongocryptd_uri, connect=False,
295+
serverSelectionTimeoutMS=_MONGOCRYPTD_TIMEOUT_MS)
296+
297+
io_callbacks = _EncryptionIO(
298+
metadata_client, key_vault_coll, mongocryptd_client, opts)
264299
self._auto_encrypter = AutoEncrypter(io_callbacks, MongoCryptOptions(
265300
opts._kms_providers, schema_map))
266-
self._bypass_auto_encryption = opts._bypass_auto_encryption
267301
self._closed = False
268302

269303
def encrypt(self, database, cmd, check_keys, codec_options):
@@ -313,29 +347,9 @@ def close(self):
313347
"""Cleanup resources."""
314348
self._closed = True
315349
self._auto_encrypter.close()
316-
317-
@staticmethod
318-
def create(client, opts):
319-
"""Create a _CommandEncyptor for a client.
320-
321-
:Parameters:
322-
- `client`: The encrypted MongoClient.
323-
- `opts`: The encrypted client's :class:`AutoEncryptionOpts`.
324-
325-
:Returns:
326-
A :class:`_CommandEncrypter` for this client.
327-
"""
328-
key_vault_client = opts._key_vault_client or client
329-
db, coll = opts._key_vault_namespace.split('.', 1)
330-
key_vault_coll = key_vault_client[db][coll]
331-
332-
mongocryptd_client = MongoClient(
333-
opts._mongocryptd_uri, connect=False,
334-
serverSelectionTimeoutMS=_MONGOCRYPTD_TIMEOUT_MS)
335-
336-
io_callbacks = _EncryptionIO(
337-
client, key_vault_coll, mongocryptd_client, opts)
338-
return _Encrypter(io_callbacks, opts)
350+
if self._internal_client:
351+
self._internal_client.close()
352+
self._internal_client = None
339353

340354

341355
class Algorithm(object):

‎pymongo/encryption_options.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ class AutoEncryptionOpts(object):
2929
"""Options to configure automatic client-side field level encryption."""
3030

3131
def __init__(self, kms_providers, key_vault_namespace,
32-
key_vault_client=None, schema_map=None,
32+
key_vault_client=None,
33+
schema_map=None,
3334
bypass_auto_encryption=False,
3435
mongocryptd_uri='mongodb://localhost:27020',
3536
mongocryptd_bypass_spawn=False,

‎pymongo/mongo_client.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -497,6 +497,15 @@ def __init__(
497497
configures this client to automatically encrypt collection commands
498498
and automatically decrypt results. See
499499
:ref:`automatic-client-side-encryption` for an example.
500+
If a :class:`MongoClient` is configured with
501+
``auto_encryption_opts`` and a non-None ``maxPoolSize``, a
502+
separate internal ``MongoClient`` is created if any of the
503+
following are true:
504+
505+
- A ``key_vault_client`` is not passed to
506+
:class:`~pymongo.encryption_options.AutoEncryptionOpts`
507+
- ``bypass_auto_encrpytion=False`` is passed to
508+
:class:`~pymongo.encryption_options.AutoEncryptionOpts`
500509
501510
| **Versioned API options:**
502511
| (If not set explicitly, Versioned API will not be enabled.)
@@ -615,6 +624,14 @@ def __init__(
615624
616625
client.__my_database__
617626
"""
627+
self.__init_kwargs = {'host': host,
628+
'port': port,
629+
'document_class': document_class,
630+
'tz_aware': tz_aware,
631+
'connect': connect,
632+
'type_registry': type_registry,
633+
**kwargs}
634+
618635
if host is None:
619636
host = self.HOST
620637
if isinstance(host, string_type):
@@ -763,9 +780,14 @@ def target():
763780
self._encrypter = None
764781
if self.__options.auto_encryption_opts:
765782
from pymongo.encryption import _Encrypter
766-
self._encrypter = _Encrypter.create(
783+
self._encrypter = _Encrypter(
767784
self, self.__options.auto_encryption_opts)
768785

786+
def _duplicate(self, **kwargs):
787+
args = self.__init_kwargs.copy()
788+
args.update(kwargs)
789+
return MongoClient(**args)
790+
769791
def _cache_credentials(self, source, credentials, connect=False):
770792
"""Save a set of authentication credentials.
771793

‎test/client-side-encryption/spec/aggregate.json‎

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -150,18 +150,6 @@
150150
"command_name": "listCollections"
151151
}
152152
},
153-
{
154-
"command_started_event": {
155-
"command": {
156-
"listCollections": 1,
157-
"filter": {
158-
"name": "datakeys"
159-
},
160-
"$db": "keyvault"
161-
},
162-
"command_name": "listCollections"
163-
}
164-
},
165153
{
166154
"command_started_event": {
167155
"command": {
@@ -273,18 +261,6 @@
273261
"command_name": "aggregate"
274262
}
275263
},
276-
{
277-
"command_started_event": {
278-
"command": {
279-
"listCollections": 1,
280-
"filter": {
281-
"name": "datakeys"
282-
},
283-
"$db": "keyvault"
284-
},
285-
"command_name": "listCollections"
286-
}
287-
},
288264
{
289265
"command_started_event": {
290266
"command": {

‎test/client-side-encryption/spec/awsTemporary.json‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -130,18 +130,6 @@
130130
"command_name": "listCollections"
131131
}
132132
},
133-
{
134-
"command_started_event": {
135-
"command": {
136-
"listCollections": 1,
137-
"filter": {
138-
"name": "datakeys"
139-
},
140-
"$db": "keyvault"
141-
},
142-
"command_name": "listCollections"
143-
}
144-
},
145133
{
146134
"command_started_event": {
147135
"command": {

‎test/client-side-encryption/spec/azureKMS.json‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -139,18 +139,6 @@
139139
"command_name": "listCollections"
140140
}
141141
},
142-
{
143-
"command_started_event": {
144-
"command": {
145-
"listCollections": 1,
146-
"filter": {
147-
"name": "datakeys"
148-
},
149-
"$db": "keyvault"
150-
},
151-
"command_name": "listCollections"
152-
}
153-
},
154142
{
155143
"command_started_event": {
156144
"command": {

‎test/client-side-encryption/spec/basic.json‎

Lines changed: 0 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -144,18 +144,6 @@
144144
"command_name": "listCollections"
145145
}
146146
},
147-
{
148-
"command_started_event": {
149-
"command": {
150-
"listCollections": 1,
151-
"filter": {
152-
"name": "datakeys"
153-
},
154-
"$db": "keyvault"
155-
},
156-
"command_name": "listCollections"
157-
}
158-
},
159147
{
160148
"command_started_event": {
161149
"command": {
@@ -283,18 +271,6 @@
283271
"command_name": "listCollections"
284272
}
285273
},
286-
{
287-
"command_started_event": {
288-
"command": {
289-
"listCollections": 1,
290-
"filter": {
291-
"name": "datakeys"
292-
},
293-
"$db": "keyvault"
294-
},
295-
"command_name": "listCollections"
296-
}
297-
},
298274
{
299275
"command_started_event": {
300276
"command": {

‎test/client-side-encryption/spec/bulk.json‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -178,18 +178,6 @@
178178
"command_name": "listCollections"
179179
}
180180
},
181-
{
182-
"command_started_event": {
183-
"command": {
184-
"listCollections": 1,
185-
"filter": {
186-
"name": "datakeys"
187-
},
188-
"$db": "keyvault"
189-
},
190-
"command_name": "listCollections"
191-
}
192-
},
193181
{
194182
"command_started_event": {
195183
"command": {

‎test/client-side-encryption/spec/count.json‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -149,18 +149,6 @@
149149
"command_name": "listCollections"
150150
}
151151
},
152-
{
153-
"command_started_event": {
154-
"command": {
155-
"listCollections": 1,
156-
"filter": {
157-
"name": "datakeys"
158-
},
159-
"$db": "keyvault"
160-
},
161-
"command_name": "listCollections"
162-
}
163-
},
164152
{
165153
"command_started_event": {
166154
"command": {

‎test/client-side-encryption/spec/countDocuments.json‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -150,18 +150,6 @@
150150
"command_name": "listCollections"
151151
}
152152
},
153-
{
154-
"command_started_event": {
155-
"command": {
156-
"listCollections": 1,
157-
"filter": {
158-
"name": "datakeys"
159-
},
160-
"$db": "keyvault"
161-
},
162-
"command_name": "listCollections"
163-
}
164-
},
165153
{
166154
"command_started_event": {
167155
"command": {

0 commit comments

Comments
 (0)