Skip to content

Commit b14dede

Browse files
prashantmitalprashantmital
committed
PYTHON-1783 Decode user-facing documents but not internal driver-server
communications. (cherry picked from commit 7491162)
1 parent 3f7481f commit b14dede

File tree

13 files changed

+486
-56
lines changed

13 files changed

+486
-56
lines changed

‎bson/__init__.py‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,7 @@ def _get_array(data, position, obj_end, opts, element_name):
215215
end -= 1
216216
result = []
217217

218-
# Avoid doing global and attibute lookups in the loop.
218+
# Avoid doing global and attribute lookups in the loop.
219219
append = result.append
220220
index = data.index
221221
getter = _ELEMENT_GETTER
@@ -940,6 +940,59 @@ def decode_all(data, codec_options=DEFAULT_CODEC_OPTIONS):
940940
decode_all = _cbson.decode_all
941941

942942

943+
def _decode_selective(rawdoc, fields, codec_options):
944+
doc = codec_options.document_class()
945+
for key, value in iteritems(rawdoc):
946+
if key in fields:
947+
if fields[key] == list:
948+
doc[key] = [_bson_to_dict(r.raw, codec_options) for r in value]
949+
elif fields[key] == dict:
950+
doc[key] = _bson_to_dict(value.raw, codec_options)
951+
else:
952+
doc[key] = _decode_selective(value, fields[key], codec_options)
953+
continue
954+
doc[key] = value
955+
return doc
956+
957+
958+
def _decode_all_selective(data, codec_options, fields):
959+
"""Decode BSON data to a single document while using user-provided
960+
custom decoding logic.
961+
962+
`data` must be a string representing a valid, BSON-encoded document.
963+
964+
:Parameters:
965+
- `data`: BSON data
966+
- `codec_options`: An instance of
967+
:class:`~bson.codec_options.CodecOptions` with user-specified type
968+
decoders. If no decoders are found, this method is the same as
969+
``decode_all``.
970+
- `fields`: Map of document namespaces where data that needs
971+
to be custom decoded lives or None. For example, to custom decode a
972+
list of objects in 'field1.subfield1', the specified value should be
973+
``{'field1': {'subfield1': list}}``. Use ``dict`` instead of ``list``
974+
if the field contains a single object to custom decode. If ``fields``
975+
is an empty map or None, this method is the same as ``decode_all``.
976+
977+
:Returns:
978+
- `document_list`: Single-member list containing the decoded document.
979+
980+
.. versionadded:: 3.8
981+
"""
982+
if not codec_options.type_registry._decoder_map:
983+
return decode_all(data, codec_options)
984+
985+
if not fields:
986+
return decode_all(data, codec_options.with_options(type_registry=None))
987+
988+
# Decode documents for internal use.
989+
from bson.raw_bson import RawBSONDocument
990+
internal_codec_options = codec_options.with_options(
991+
document_class=RawBSONDocument, type_registry=None)
992+
_doc = _bson_to_dict(data, internal_codec_options)
993+
return [_decode_selective(_doc, fields, codec_options,)]
994+
995+
943996
def decode_iter(data, codec_options=DEFAULT_CODEC_OPTIONS):
944997
"""Decode BSON data to multiple documents as a generator.
945998

‎gridfs/__init__.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,8 @@
2525
from gridfs.grid_file import (GridIn,
2626
GridOut,
2727
GridOutCursor,
28-
DEFAULT_CHUNK_SIZE)
28+
DEFAULT_CHUNK_SIZE,
29+
_clear_entity_type_registry)
2930
from pymongo import (ASCENDING,
3031
DESCENDING)
3132
from pymongo.common import UNAUTHORIZED_CODES, validate_string
@@ -61,6 +62,8 @@ def __init__(self, database, collection="fs", disable_md5=False):
6162
if not isinstance(database, Database):
6263
raise TypeError("database must be an instance of Database")
6364

65+
database = _clear_entity_type_registry(database)
66+
6467
if not database.write_concern.acknowledged:
6568
raise ConfigurationError('database must use '
6669
'acknowledged write_concern')
@@ -443,6 +446,8 @@ def __init__(self, db, bucket_name="fs",
443446
if not isinstance(db, Database):
444447
raise TypeError("database must be an instance of Database")
445448

449+
db = _clear_entity_type_registry(db)
450+
446451
wtc = write_concern if write_concern is not None else db.write_concern
447452
if not wtc.acknowledged:
448453
raise ConfigurationError('write concern must be acknowledged')

‎gridfs/grid_file.py‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,12 @@ def getter(self):
9898
return property(getter, doc=docstring)
9999

100100

101+
def _clear_entity_type_registry(entity, **kwargs):
102+
"""Clear the given database/collection object's type registry."""
103+
codecopts = entity.codec_options.with_options(type_registry=None)
104+
return entity.with_options(codec_options=codecopts, **kwargs)
105+
106+
101107
class GridIn(object):
102108
"""Class to write data to GridFS.
103109
"""
@@ -168,8 +174,8 @@ def __init__(
168174
if "chunk_size" in kwargs:
169175
kwargs["chunkSize"] = kwargs.pop("chunk_size")
170176

171-
coll = root_collection.with_options(
172-
read_preference=ReadPreference.PRIMARY)
177+
coll = _clear_entity_type_registry(
178+
root_collection, read_preference=ReadPreference.PRIMARY)
173179

174180
if not disable_md5:
175181
kwargs["md5"] = hashlib.md5()
@@ -449,6 +455,8 @@ def __init__(self, root_collection, file_id=None, file_document=None,
449455
raise TypeError("root_collection must be an "
450456
"instance of Collection")
451457

458+
root_collection = _clear_entity_type_registry(root_collection)
459+
452460
self.__chunks = root_collection.chunks
453461
self.__files = root_collection.files
454462
self.__file_id = file_id
@@ -800,6 +808,8 @@ def __init__(self, collection, filter=None, skip=0, limit=0,
800808
801809
.. mongodoc:: cursors
802810
"""
811+
collection = _clear_entity_type_registry(collection)
812+
803813
# Hold on to the base "fs" collection to create GridOut objects later.
804814
self.__root_collection = collection
805815

‎pymongo/collection.py‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353

5454
_NO_OBJ_ERROR = "No matching object found"
5555
_UJOIN = u"%s.%s"
56+
_FIND_AND_MODIFY_DOC_FIELDS = {'value': dict}
5657

5758

5859
class ReturnDocument(object):
@@ -202,7 +203,8 @@ def _command(self, sock_info, command, slave_ok=False,
202203
write_concern=None,
203204
collation=None,
204205
session=None,
205-
retryable_write=False):
206+
retryable_write=False,
207+
user_fields=None):
206208
"""Internal command helper.
207209
208210
:Parameters:
@@ -241,7 +243,8 @@ def _command(self, sock_info, command, slave_ok=False,
241243
collation=collation,
242244
session=s,
243245
client=self.__database.client,
244-
retryable_write=retryable_write)
246+
retryable_write=retryable_write,
247+
user_fields=user_fields)
245248

246249
def __create(self, options, collation, session):
247250
"""Sends a create command with the given options.
@@ -314,9 +317,8 @@ def database(self):
314317
"""
315318
return self.__database
316319

317-
def with_options(
318-
self, codec_options=None, read_preference=None,
319-
write_concern=None, read_concern=None):
320+
def with_options(self, codec_options=None, read_preference=None,
321+
write_concern=None, read_concern=None):
320322
"""Get a clone of this collection changing the specified settings.
321323
322324
>>> coll1.read_preference
@@ -2307,7 +2309,8 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
23072309
write_concern=write_concern,
23082310
collation=collation,
23092311
session=session,
2310-
client=self.__database.client)
2312+
client=self.__database.client,
2313+
user_fields={'cursor': {'firstBatch': list}})
23112314

23122315
if "cursor" in result:
23132316
cursor = result["cursor"]
@@ -2565,7 +2568,8 @@ def group(self, key, condition, initial, reduce, finalize=None, **kwargs):
25652568

25662569
with self._socket_for_reads(session=None) as (sock_info, slave_ok):
25672570
return self._command(sock_info, cmd, slave_ok,
2568-
collation=collation)["retval"]
2571+
collation=collation,
2572+
user_fields={'retval': list})["retval"]
25692573

25702574
def rename(self, new_name, session=None, **kwargs):
25712575
"""Rename this collection.
@@ -2669,7 +2673,8 @@ def distinct(self, key, filter=None, session=None, **kwargs):
26692673
with self._socket_for_reads(session) as (sock_info, slave_ok):
26702674
return self._command(sock_info, cmd, slave_ok,
26712675
read_concern=self.read_concern,
2672-
collation=collation, session=session)["values"]
2676+
collation=collation,
2677+
session=session)["values"]
26732678

26742679
def map_reduce(self, map, reduce, out, full_response=False, session=None,
26752680
**kwargs):
@@ -2749,12 +2754,17 @@ def map_reduce(self, map, reduce, out, full_response=False, session=None,
27492754
write_concern = self._write_concern_for(session)
27502755
else:
27512756
write_concern = None
2757+
if inline:
2758+
user_fields = {'results': list}
2759+
else:
2760+
user_fields = None
27522761

27532762
response = self._command(
27542763
sock_info, cmd, slave_ok, read_pref,
27552764
read_concern=read_concern,
27562765
write_concern=write_concern,
2757-
collation=collation, session=session)
2766+
collation=collation, session=session,
2767+
user_fields=user_fields)
27582768

27592769
if full_response or not response.get('result'):
27602770
return response
@@ -2804,16 +2814,19 @@ def inline_map_reduce(self, map, reduce, full_response=False, session=None,
28042814
("map", map),
28052815
("reduce", reduce),
28062816
("out", {"inline": 1})])
2817+
user_fields = {'results': list}
28072818
collation = validate_collation_or_none(kwargs.pop('collation', None))
28082819
cmd.update(kwargs)
28092820
with self._socket_for_reads(session) as (sock_info, slave_ok):
28102821
if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
28112822
res = self._command(sock_info, cmd, slave_ok,
28122823
read_concern=self.read_concern,
2813-
collation=collation, session=session)
2824+
collation=collation, session=session,
2825+
user_fields=user_fields)
28142826
else:
28152827
res = self._command(sock_info, cmd, slave_ok,
2816-
collation=collation, session=session)
2828+
collation=collation, session=session,
2829+
user_fields=user_fields)
28172830

28182831
if full_response:
28192832
return res
@@ -2831,6 +2844,7 @@ def __find_and_modify(self, filter, projection, sort, upsert=None,
28312844
return_document=ReturnDocument.BEFORE,
28322845
array_filters=None, session=None, **kwargs):
28332846
"""Internal findAndModify helper."""
2847+
28342848
common.validate_is_mapping("filter", filter)
28352849
if not isinstance(return_document, bool):
28362850
raise ValueError("return_document must be "
@@ -2870,8 +2884,10 @@ def _find_and_modify(session, sock_info, retryable_write):
28702884
write_concern=write_concern,
28712885
allowable_errors=[_NO_OBJ_ERROR],
28722886
collation=collation, session=session,
2873-
retryable_write=retryable_write)
2887+
retryable_write=retryable_write,
2888+
user_fields=_FIND_AND_MODIFY_DOC_FIELDS)
28742889
_check_write_command_response(out)
2890+
28752891
return out.get("value")
28762892

28772893
return self.__database.client._retryable_write(
@@ -3287,7 +3303,8 @@ def _find_and_modify(session, sock_info, retryable_write):
32873303
result = self._command(
32883304
sock_info, cmd, read_preference=ReadPreference.PRIMARY,
32893305
allowable_errors=[_NO_OBJ_ERROR], collation=collation,
3290-
session=session, retryable_write=retryable_write)
3306+
session=session, retryable_write=retryable_write,
3307+
user_fields=_FIND_AND_MODIFY_DOC_FIELDS)
32913308

32923309
_check_write_command_response(result)
32933310
return result

‎pymongo/command_cursor.py‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -149,12 +149,17 @@ def duration(): return datetime.datetime.now() - start
149149
reply = response.data
150150

151151
try:
152-
docs = self._unpack_response(reply,
153-
self.__id,
154-
self.__collection.codec_options)
152+
user_fields = None
153+
legacy_response = True
154+
if from_command:
155+
user_fields = {'cursor': {'nextBatch': list}}
156+
legacy_response = False
157+
docs = self._unpack_response(
158+
reply, self.__id, self.__collection.codec_options,
159+
legacy_response=legacy_response, user_fields=user_fields)
155160
if from_command:
156161
first = docs[0]
157-
client._receive_cluster_time(first, self.__session)
162+
client._process_response(first, self.__session)
158163
helpers._check_command_response(first)
159164

160165
except OperationFailure as exc:
@@ -174,7 +179,6 @@ def duration(): return datetime.datetime.now() - start
174179
listeners.publish_command_failure(
175180
duration(), exc.details, "getMore", rqst_id, self.__address)
176181

177-
client._reset_server_and_request_check(self.address)
178182
raise
179183
except Exception as exc:
180184
if publish:
@@ -208,8 +212,10 @@ def duration(): return datetime.datetime.now() - start
208212
kill()
209213
self.__data = deque(documents)
210214

211-
def _unpack_response(self, response, cursor_id, codec_options):
212-
return response.unpack_response(cursor_id, codec_options)
215+
def _unpack_response(self, response, cursor_id, codec_options,
216+
user_fields=None, legacy_response=False):
217+
return response.unpack_response(cursor_id, codec_options, user_fields,
218+
legacy_response)
213219

214220
def _refresh(self):
215221
"""Refreshes the cursor with more data from the server.
@@ -330,7 +336,8 @@ def __init__(self, collection, cursor_info, address, retrieved=0,
330336
collection, cursor_info, address, retrieved, batch_size,
331337
max_await_time_ms, session, explicit_session)
332338

333-
def _unpack_response(self, response, cursor_id, codec_options):
339+
def _unpack_response(self, response, cursor_id, codec_options,
340+
user_fields=None, legacy_response=False):
334341
return response.raw_response(cursor_id)
335342

336343
def __getitem__(self, index):

0 commit comments

Comments
 (0)