PYTHON-1818 TypeCodec support for ChangeStreams

Author: prashantmital

bson/__init__.py

@@ -948,7 +948,13 @@ def decode_all(data, codec_options=DEFAULT_CODEC_OPTIONS):
 
 
 def _decode_selective(rawdoc, fields, codec_options):
-    doc = {}
+    if _raw_document_class(codec_options.document_class):
+        # If document_class is RawBSONDocument, use vanilla dictionary for
+        # decoding command response.
+        doc = {}
+    else:
+        # Else, use the specified document_class.
+        doc = codec_options.document_class()
     for key, value in iteritems(rawdoc):
         if key in fields:
             if fields[key] == 1:

bson/raw_bson.py

@@ -19,6 +19,7 @@
 from bson.py3compat import abc, iteritems
 from bson.codec_options import (
     DEFAULT_CODEC_OPTIONS as DEFAULT, _RAW_BSON_DOCUMENT_MARKER)
+from bson.son import SON
 
 
 class RawBSONDocument(abc.Mapping):
@@ -93,8 +94,9 @@ def __inflated(self):
         if self.__inflated_doc is None:
             # We already validated the object's size when this document was
             # created, so no need to do that again.
+            # Use SON to preserve ordering of elements.
             self.__inflated_doc = _elements_to_dict(
-                self.__raw, 4, len(self.__raw)-1, self.__codec_options, {})
+                self.__raw, 4, len(self.__raw)-1, self.__codec_options, SON())
         return self.__inflated_doc
 
     def __getitem__(self, item):

pymongo/change_stream.py

@@ -16,6 +16,8 @@
 
 import copy
 
+from bson import _bson_to_dict
+from bson.raw_bson import RawBSONDocument
 from bson.son import SON
 
 from pymongo import common
@@ -60,7 +62,16 @@ def __init__(self, target, pipeline, full_document, resume_after,
         validate_collation_or_none(collation)
         common.validate_non_negative_integer_or_none("batchSize", batch_size)
 
-        self._target = target
+        self._decode_custom = False
+        self._orig_codec_options = target.codec_options
+        if target.codec_options.type_registry._decoder_map:
+            self._decode_custom = True
+            self._target = target.with_options(
+                codec_options=target.codec_options.with_options(
+                    document_class=RawBSONDocument, type_registry=None))
+        else:
+            self._target = target
+
         self._pipeline = copy.deepcopy(pipeline)
         self._full_document = full_document
         self._resume_token = copy.deepcopy(resume_after)
@@ -145,8 +156,7 @@ def _run_aggregation_cmd(self, session, explicit_session):
                 aggregation_collection, cursor, sock_info.address,
                 batch_size=self._batch_size or 0,
                 max_await_time_ms=self._max_await_time_ms,
-                session=session, explicit_session=explicit_session
-            )
+                session=session, explicit_session=explicit_session)
 
     def _create_cursor(self):
         with self._database.client._tmp_session(self._session, close=False) as s:
@@ -263,6 +273,9 @@ def try_next(self):
                 "token is missing.")
         self._resume_token = copy.copy(resume_token)
         self._start_at_operation_time = None
+
+        if self._decode_custom:
+            return _bson_to_dict(change.raw, self._orig_codec_options)
         return change
 
     def __enter__(self):

test/test_custom_types.py

@@ -17,6 +17,7 @@
 import datetime
 import sys
 import tempfile
+from collections import OrderedDict
 from decimal import Decimal
 from random import random
 
@@ -36,16 +37,18 @@
                                 TypeEncoder, TypeRegistry)
 from bson.errors import InvalidDocument
 from bson.int64 import Int64
+from bson.raw_bson import RawBSONDocument
 from bson.py3compat import text_type
 
 from gridfs import GridIn, GridOut
 
 from pymongo.collection import ReturnDocument
 from pymongo.errors import DuplicateKeyError
+from pymongo.message import _CursorAddress
 
 from test import client_context, unittest
 from test.test_client import IntegrationTest
-from test.utils import ignore_deprecations
+from test.utils import ignore_deprecations, rs_client
 
 
 class DecimalEncoder(TypeEncoder):
@@ -115,6 +118,14 @@ def transform_bson(self, value):
     [UppercaseTextDecoder(),]))
 
 
+def type_obfuscating_decoder_factory(rt_type):
+    class ResumeTokenToNanDecoder(TypeDecoder):
+        bson_type = rt_type
+        def transform_bson(self, value):
+            return "NaN"
+    return ResumeTokenToNanDecoder
+
+
 class CustomBSONTypeTests(object):
     def roundtrip(self, doc):
         bsonbytes = BSON().encode(doc, codec_options=self.codecopts)
@@ -549,7 +560,7 @@ def test_command_errors_w_custom_type_decoder(self):
     def test_find_w_custom_type_decoder(self):
         db = self.db
         input_docs = [
-            {'x': Int64(k)} for k in [1.0, 2.0, 3.0]]
+            {'x': Int64(k)} for k in [1, 2, 3]]
         for doc in input_docs:
             db.test.insert_one(doc)
 
@@ -558,6 +569,24 @@ def test_find_w_custom_type_decoder(self):
         for doc in test.find({}, batch_size=1):
             self.assertIsInstance(doc['x'], UndecipherableInt64Type)
 
+    def test_find_w_custom_type_decoder_and_document_class(self):
+        def run_test(doc_cls):
+            db = self.db
+            input_docs = [
+                {'x': Int64(k)} for k in [1, 2, 3]]
+            for doc in input_docs:
+                db.test.insert_one(doc)
+
+            test = db.get_collection('test', codec_options=CodecOptions(
+                type_registry=TypeRegistry([UndecipherableIntDecoder()]),
+                document_class=doc_cls))
+            for doc in test.find({}, batch_size=1):
+                self.assertIsInstance(doc, doc_cls)
+                self.assertIsInstance(doc['x'], UndecipherableInt64Type)
+
+        for doc_cls in [RawBSONDocument, OrderedDict]:
+            run_test(doc_cls)
+
     @client_context.require_version_max(4, 1, 0, -1)
     def test_group_w_custom_type(self):
         db = self.db
@@ -709,5 +738,155 @@ def test_grid_out_custom_opts(self):
             self.assertRaises(AttributeError, setattr, two, attr, 5)
 
 
+class ChangeStreamsWCustomTypesTestMixin(object):
+    def change_stream(self, *args, **kwargs):
+        return self.watched_target.watch(*args, **kwargs)
+
+    def insert_and_check(self, change_stream, insert_doc,
+                         expected_doc):
+        self.input_target.insert_one(insert_doc)
+        change = next(change_stream)
+        self.assertEqual(change['fullDocument'], expected_doc)
+
+    def kill_change_stream_cursor(self, change_stream):
+        # Cause a cursor not found error on the next getMore.
+        cursor = change_stream._cursor
+        address = _CursorAddress(cursor.address, cursor._CommandCursor__ns)
+        client = self.input_target.database.client
+        client._close_cursor_now(cursor.cursor_id, address)
+
+    def test_simple(self):
+        codecopts = CodecOptions(type_registry=TypeRegistry([
+            UndecipherableIntEncoder(), UppercaseTextDecoder()]))
+        self.create_targets(codec_options=codecopts)
+
+        input_docs = [
+            {'_id': UndecipherableInt64Type(1), 'data': 'hello'},
+            {'_id': 2, 'data': 'world'},
+            {'_id': UndecipherableInt64Type(3), 'data': '!'},]
+        expected_docs = [
+            {'_id': 1, 'data': 'HELLO'},
+            {'_id': 2, 'data': 'WORLD'},
+            {'_id': 3, 'data': '!'},]
+
+        change_stream = self.change_stream()
+
+        self.insert_and_check(change_stream, input_docs[0], expected_docs[0])
+        self.kill_change_stream_cursor(change_stream)
+        self.insert_and_check(change_stream, input_docs[1], expected_docs[1])
+        self.kill_change_stream_cursor(change_stream)
+        self.insert_and_check(change_stream, input_docs[2], expected_docs[2])
+
+    def test_break_resume_token(self):
+        # Get one document from a change stream to determine resumeToken type.
+        self.create_targets()
+        change_stream = self.change_stream()
+        self.input_target.insert_one({"data": "test"})
+        change = next(change_stream)
+        resume_token_decoder = type_obfuscating_decoder_factory(
+            type(change['_id']['_data']))
+
+        # Custom-decoding the resumeToken type breaks resume tokens.
+        codecopts = CodecOptions(type_registry=TypeRegistry([
+            resume_token_decoder(), UndecipherableIntEncoder()]))
+
+        # Re-create targets, change stream and proceed.
+        self.create_targets(codec_options=codecopts)
+
+        docs = [{'_id': 1}, {'_id': 2}, {'_id': 3}]
+
+        change_stream = self.change_stream()
+        self.insert_and_check(change_stream, docs[0], docs[0])
+        self.kill_change_stream_cursor(change_stream)
+        self.insert_and_check(change_stream, docs[1], docs[1])
+        self.kill_change_stream_cursor(change_stream)
+        self.insert_and_check(change_stream, docs[2], docs[2])
+
+    def test_document_class(self):
+        def run_test(doc_cls):
+            codecopts = CodecOptions(type_registry=TypeRegistry([
+                UppercaseTextDecoder(), UndecipherableIntEncoder()]),
+                document_class=doc_cls)
+
+            self.create_targets(codec_options=codecopts)
+            change_stream = self.change_stream()
+
+            doc = {'a': UndecipherableInt64Type(101), 'b': 'xyz'}
+            self.input_target.insert_one(doc)
+            change = next(change_stream)
+
+            self.assertIsInstance(change, doc_cls)
+            self.assertEqual(change['fullDocument']['a'], 101)
+            self.assertEqual(change['fullDocument']['b'], 'XYZ')
+
+        for doc_cls in [OrderedDict, RawBSONDocument]:
+            run_test(doc_cls)
+
+
+class TestCollectionChangeStreamsWCustomTypes(
+    IntegrationTest, ChangeStreamsWCustomTypesTestMixin):
+    @classmethod
+    @client_context.require_version_min(3, 6, 0)
+    @client_context.require_no_mmap
+    @client_context.require_no_standalone
+    def setUpClass(cls):
+        super(TestCollectionChangeStreamsWCustomTypes, cls).setUpClass()
+
+    def tearDown(self):
+        self.input_target.drop()
+
+    def create_targets(self, *args, **kwargs):
+        self.watched_target = self.db.get_collection(
+            'test', *args, **kwargs)
+        self.input_target = self.watched_target
+        # Insert a record to ensure db, coll are created.
+        self.input_target.insert_one({'data': 'dummy'})
+
+
+class TestDatabaseChangeStreamsWCustomTypes(
+    IntegrationTest, ChangeStreamsWCustomTypesTestMixin):
+    @classmethod
+    @client_context.require_version_min(4, 0, 0)
+    @client_context.require_no_mmap
+    @client_context.require_no_standalone
+    def setUpClass(cls):
+        super(TestDatabaseChangeStreamsWCustomTypes, cls).setUpClass()
+
+    def tearDown(self):
+        self.input_target.drop()
+        self.client.drop_database(self.watched_target)
+
+    def create_targets(self, *args, **kwargs):
+        self.watched_target = self.client.get_database(
+            self.db.name, *args, **kwargs)
+        self.input_target = self.watched_target.test
+        # Insert a record to ensure db, coll are created.
+        self.input_target.insert_one({'data': 'dummy'})
+
+
+class TestClusterChangeStreamsWCustomTypes(
+    IntegrationTest, ChangeStreamsWCustomTypesTestMixin):
+    @classmethod
+    @client_context.require_version_min(4, 0, 0)
+    @client_context.require_no_mmap
+    @client_context.require_no_standalone
+    def setUpClass(cls):
+        super(TestClusterChangeStreamsWCustomTypes, cls).setUpClass()
+
+    def tearDown(self):
+        self.input_target.drop()
+        self.client.drop_database(self.db)
+
+    def create_targets(self, *args, **kwargs):
+        codec_options = kwargs.pop('codec_options', None)
+        if codec_options:
+            kwargs['type_registry'] = codec_options.type_registry
+            kwargs['document_class'] = codec_options.document_class
+        self.watched_target = rs_client(*args, **kwargs)
+        self.input_target = self.watched_target[self.db.name].test
+        # Insert a record to ensure db, coll are created.
+        self.input_target.insert_one({'data': 'dummy'})
+
+
 if __name__ == "__main__":
     unittest.main()