Merge pull request #290 from IlyaSkriblovsky/fix-pymongo4-compat

PyMongo 4+ support

Author: psi29a

docs/source/NEWS.rst

@@ -1,6 +1,16 @@
 Changelog
 =========
 
+Release UPCOMING (yyyy-mm-dd)
+-----------------------------
+
+API Changes
+^^^^^^^^^^^
+
+- PyMongo 4+ is now supported. If you will migrate from PyMongo 3 to PyMongo 4, please be sure
+  to check their PyMongo's guide because newer version has a number of incompatible changes.
+
+
 Release 23.0.0 (2023-01-29)
 ---------------------------
 

requirements.txt

@@ -1,3 +1,3 @@
 Twisted
-pymongo>=3.0
+pymongo>=3.0,<4.9
 -e .

setup.py

@@ -10,7 +10,7 @@
     url="https://github.com/twisted/txmongo",
     keywords=["mongo", "mongodb", "pymongo", "gridfs", "txmongo"],
     packages=["txmongo", "txmongo._gridfs"],
-    install_requires=["twisted>=14.0", "pymongo>=3.0, <4.0"],
+    install_requires=["twisted>=14.0", "pymongo>=3.0, <4.9"],
     extras_require={
         'srv': ['pymongo[srv]>=3.6'],
     },

tests/basic/test_bulk.py

@@ -1,6 +1,6 @@
 from bson import BSON
 from pymongo import InsertOne
-from pymongo.errors import BulkWriteError, OperationFailure, NotMasterError
+from pymongo.errors import BulkWriteError, OperationFailure
 from pymongo.operations import UpdateOne, DeleteOne, UpdateMany, ReplaceOne
 from pymongo.results import BulkWriteResult
 from pymongo.write_concern import WriteConcern
@@ -10,6 +10,12 @@
 from tests.utils import SingleCollectionTest
 from txmongo.protocol import Reply
 
+try:
+    from pymongo.errors import NotPrimaryError
+except ImportError:
+    # For pymongo < 3.12
+    from pymongo.errors import NotMasterError as NotPrimaryError
+
 
 class TestArgsValidation(SingleCollectionTest):
 
@@ -242,4 +248,4 @@ def fake_send_query(*args):
         with patch('txmongo.protocol.MongoProtocol.send_QUERY', side_effect=fake_send_query):
             yield self.assertFailure(
                     self.coll.bulk_write([UpdateOne({}, {'$set': {'x': 42}}, upsert=True)], ordered=True),
-                    OperationFailure, NotMasterError)
+                    OperationFailure, NotPrimaryError)

tests/basic/test_collection.py

@@ -16,9 +16,13 @@
 Based on pymongo driver's test_collection.py
 """
 
+import uuid
+
+from bson import UuidRepresentation, CodecOptions
 from pymongo import errors
 from twisted.internet import defer
 from twisted.trial import unittest
+
 import txmongo
 from txmongo import filter as qf
 from txmongo.collection import Collection
@@ -446,3 +450,40 @@ def test_convert(self):
             {"filter": {'x': 42}, "projection": {'a': 1}, "skip": 5, "limit": 6,
              "sort": qf.sort([('s', 1)])}
         )
+
+
+class TestUuid(unittest.TestCase):
+
+    def setUp(self):
+        self.conn = txmongo.MongoConnection(mongo_host, mongo_port, codec_options=CodecOptions(uuid_representation=UuidRepresentation.STANDARD))
+        self.db = self.conn.mydb
+        self.coll = self.db.conn
+
+    @defer.inlineCallbacks
+    def tearDown(self):
+        yield self.conn.drop_database(self.db)
+        yield self.conn.disconnect()
+
+    @defer.inlineCallbacks
+    def test_uuid_in_crud(self):
+        task_id = uuid.uuid4()
+        yield self.coll.insert({'task_id': task_id}, safe=True)
+
+        doc = yield self.coll.find_one(fields={"_id": 0})
+        self.assertEqual(doc, {'task_id': task_id})
+
+        doc = yield self.coll.find_one({'task_id': task_id}, fields={"_id": 0})
+        self.assertEqual(doc, {'task_id': task_id})
+
+        new = uuid.uuid4()
+        yield self.coll.update_one({'task_id': task_id}, {'$set': {'task_id': new}})
+        doc = yield self.coll.find_one()
+        self.assertEqual(doc['task_id'], new)
+
+        yield self.coll.delete_one({'task_id': new})
+        cnt = yield self.coll.count()
+        self.assertEqual(cnt, 0)
+
+    @defer.inlineCallbacks
+    def test_uuid_in_batch(self):
+        yield self.coll.batch.insert_many([{'task_id': uuid.uuid4()} for _ in range(10)])

tests/basic/test_write_concern.py

@@ -137,7 +137,7 @@ def test_AllOperations(self):
 
     @defer.inlineCallbacks
     def test_ConnectionUrlParams(self):
-        conn = ConnectionPool("mongodb://{0}:{1}/?w=2&j=true".format(mongo_host, mongo_port))
+        conn = ConnectionPool("mongodb://{0}:{1}/?w=2&journal=true".format(mongo_host, mongo_port))
         coll = conn.mydb.mycol
 
         try:

txmongo/collection.py

@@ -2,31 +2,33 @@
 # Use of this source code is governed by the Apache License that can be
 # found in the LICENSE file.
 
+import collections.abc
 import io
 import struct
-import collections.abc
 import warnings
+
 from bson import BSON, ObjectId
 from bson.code import Code
-from bson.son import SON
 from bson.codec_options import CodecOptions
+from bson.son import SON
 from pymongo.bulk import _Bulk, _COMMANDS
+from pymongo.collection import ReturnDocument
+from pymongo.common import validate_ok_for_update, validate_ok_for_replace, \
+    validate_is_mapping, validate_boolean
 from pymongo.errors import InvalidName, BulkWriteError, InvalidOperation, OperationFailure
 from pymongo.message import _OP_MAP, _INSERT
 from pymongo.results import InsertOneResult, InsertManyResult, UpdateResult, \
     DeleteResult, BulkWriteResult
-from pymongo.common import validate_ok_for_update, validate_ok_for_replace, \
-    validate_is_mapping, validate_boolean
-from pymongo.collection import ReturnDocument
 from pymongo.write_concern import WriteConcern
+from twisted.internet import defer
+from twisted.python.compat import comparable
+
+from txmongo import filter as qf
 from txmongo.filter import _QueryFilter
 from txmongo.protocol import DELETE_SINGLE_REMOVE, UPDATE_UPSERT, UPDATE_MULTI, \
     Query, Getmore, Insert, Update, Delete, KillCursors
 from txmongo.pymongo_internals import _check_write_command_response, _merge_command, _check_command_response
 from txmongo.utils import check_deadline, timeout
-from txmongo import filter as qf
-from twisted.internet import defer
-from twisted.python.compat import comparable
 
 
 @comparable
@@ -358,7 +360,6 @@ def query():
         return self.__real_find_with_cursor(**new_kwargs)
 
     def __real_find_with_cursor(self, filter=None, projection=None, skip=0, limit=0, sort=None, batch_size=0,**kwargs):
-        
         if filter is None:
             filter = SON()
 
@@ -376,6 +377,8 @@ def __real_find_with_cursor(self, filter=None, projection=None, skip=0, limit=0,
         projection = self._normalize_fields_projection(projection)
 
         filter = self.__apply_find_filter(filter, sort)
+        if not isinstance(filter, BSON):
+            filter = BSON.encode(filter, codec_options=self.codec_options)
 
         as_class = kwargs.get("as_class")
         proto = self._database.connection.getprotocol()
@@ -603,7 +606,7 @@ def insert(self, docs, safe=None, flags=0, **kwargs):
         else:
             raise TypeError("TxMongo: insert takes a document or a list of documents.")
 
-        docs = [BSON.encode(d) for d in docs]
+        docs = [BSON.encode(d, codec_options=self.codec_options) for d in docs]
         insert = Insert(flags=flags, collection=str(self), documents=docs)
 
         def on_proto(proto):
@@ -654,8 +657,7 @@ def on_ok(result):
             return InsertOneResult(inserted_id, self.write_concern.acknowledged)
         return self._insert_one(document, _deadline).addCallback(on_ok)
 
-    @staticmethod
-    def _generate_batch_commands(collname, command, docs_field, documents, ordered,
+    def _generate_batch_commands(self, collname, command, docs_field, documents, ordered,
                                  write_concern, max_bson, max_count):
         # Takes a list of documents and generates one or many `insert` commands
         # with documents list in each command is less or equal to max_bson bytes
@@ -669,7 +671,7 @@ def _generate_batch_commands(collname, command, docs_field, documents, ordered,
                    ("writeConcern", write_concern.document)])
 
         buf = io.BytesIO()
-        buf.write(BSON.encode(msg))
+        buf.write(BSON.encode(msg, codec_options=self.codec_options))
         buf.seek(-1, io.SEEK_END)  # -1 because we don't need final NUL from partial command
         buf.write(docs_field)  # type, name and length placeholder of 'documents' array
         docs_start = buf.tell() - 4
@@ -693,7 +695,7 @@ def prepare_command():
         idx_offset = 0
         for doc in documents:
             key = str(idx).encode('ascii')
-            value = BSON.encode(doc)
+            value = BSON.encode(doc, codec_options=self.codec_options)
 
             enough_size = buf.tell() + len(key)+2 + len(value) - docs_start > max_bson
             enough_count = idx >= max_count
@@ -796,8 +798,8 @@ def update(self, spec, document, upsert=False, multi=False, safe=None, flags=0,
         if upsert:
             flags |= UPDATE_UPSERT
 
-        spec = BSON.encode(spec)
-        document = BSON.encode(document)
+        spec = BSON.encode(spec, codec_options=self.codec_options)
+        document = BSON.encode(document, codec_options=self.codec_options)
         update = Update(flags=flags, collection=str(self),
                         selector=spec, update=document)
 
@@ -959,7 +961,7 @@ def remove(self, spec, safe=None, single=False, flags=0, **kwargs):
         if single:
             flags |= DELETE_SINGLE_REMOVE
 
-        spec = BSON.encode(spec)
+        spec = BSON.encode(spec, codec_options=self.codec_options)
         delete = Delete(flags=flags, collection=str(self), selector=spec)
 
         def on_proto(proto):

txmongo/connection.py

@@ -4,12 +4,13 @@
 
 from bson.codec_options import DEFAULT_CODEC_OPTIONS
 from pymongo.errors import AutoReconnect, ConfigurationError, OperationFailure
-from pymongo.uri_parser import parse_uri
 from pymongo.read_preferences import ReadPreference
+from pymongo.uri_parser import parse_uri
 from pymongo.write_concern import WriteConcern
 from twisted.internet import defer, reactor, task
 from twisted.internet.protocol import ReconnectingClientFactory, ClientFactory
 from twisted.python import log
+
 from txmongo.database import Database
 from txmongo.protocol import MongoProtocol, Query
 from txmongo.utils import timeout, get_err

txmongo/database.py

@@ -3,7 +3,7 @@
 # found in the LICENSE file.
 
 from bson.son import SON
-from bson.codec_options import DEFAULT_CODEC_OPTIONS
+
 from txmongo.collection import Collection
 from txmongo.pymongo_internals import _check_command_response
 from txmongo.utils import timeout
@@ -53,10 +53,12 @@ def codec_options(self):
 
     @timeout
     def command(self, command, value=1, check=True, allowable_errors=None,
-                codec_options=DEFAULT_CODEC_OPTIONS, _deadline=None, **kwargs):
-        """command(command, value=1, check=True, allowable_errors=None, codec_options=DEFAULT_CODEC_OPTIONS)"""
+                codec_options=None, _deadline=None, **kwargs):
+        """command(command, value=1, check=True, allowable_errors=None, codec_options=None)"""
         if isinstance(command, (bytes, str)):
             command = SON([(command, value)])
+        if codec_options is None:
+            codec_options = self.__codec_options
         options = kwargs.copy()
         command.update(options)
 

txmongo/protocol.py

@@ -15,22 +15,29 @@
 
 import base64
 import hashlib
-
-from bson import BSON, SON, Binary
-from collections import namedtuple
-from hashlib import sha1
 import hmac
 import logging
+import struct
+from collections import namedtuple
+from hashlib import sha1
+from random import SystemRandom
+
+from bson import BSON, SON, Binary
 from pymongo import auth
 from pymongo.errors import AutoReconnect, ConnectionFailure, DuplicateKeyError, OperationFailure, \
-    NotMasterError, CursorNotFound
-from random import SystemRandom
-import struct
+    CursorNotFound
 from twisted.internet import defer, protocol, error
 from twisted.python import failure, log
+
 from txmongo.utils import get_err
 
 
+try:
+    from pymongo.errors import NotPrimaryError
+except ImportError:
+    # For pymongo < 3.12
+    from pymongo.errors import NotMasterError as NotPrimaryError
+
 try:
     from hashlib import pbkdf2_hmac as _hi
 except ImportError:
@@ -397,7 +404,7 @@ def handle_REPLY(self, request):
                 msg = "TxMongo: " + doc.get("$err", "Unknown error")
                 fail_conn = False
                 if code == 13435:
-                    err = NotMasterError(msg)
+                    err = NotPrimaryError(msg)
                     fail_conn = True
                 else:
                     err = OperationFailure(msg, code)

txmongo/pymongo_internals.py

@@ -1,9 +1,15 @@
-from pymongo.bulk import _UOP
 from pymongo.errors import (DuplicateKeyError, WriteError, WTimeoutError, WriteConcernError, OperationFailure,
-                            NotMasterError, ExecutionTimeout, CursorNotFound)
+                            ExecutionTimeout, CursorNotFound)
 from pymongo.message import _INSERT, _DELETE, _UPDATE
 
 
+try:
+    from pymongo.errors import NotPrimaryError
+except ImportError:
+    # For pymongo < 3.12
+    from pymongo.errors import NotMasterError as NotPrimaryError
+
+
 
 # Copied from pymongo/helpers.py:32 at commit d7d94b2776098dba32686ddf3ada1f201172daaf
 
@@ -100,7 +106,7 @@ def _merge_command(run, full_result, results):
                 idx = doc["index"] + offset
                 replacement["index"] = run.index(idx)
                 # Add the failed operation to the error document.
-                replacement[_UOP] = run.ops[idx]
+                replacement["op"] = run.ops[idx]
                 full_result["writeErrors"].append(replacement)
 
         wc_error = result.get("writeConcernError")
@@ -140,10 +146,10 @@ def _check_command_response(response, msg=None, allowable_errors=None,
             code = details.get("code")
             # Server is "not master" or "recovering"
             if code in _NOT_MASTER_CODES:
-                raise NotMasterError(errmsg, response)
+                raise NotPrimaryError(errmsg, response)
             elif ("not master" in errmsg
                   or "node is recovering" in errmsg):
-                raise NotMasterError(errmsg, response)
+                raise NotPrimaryError(errmsg, response)
 
             # Server assertion failures
             if errmsg == "db assertion failure":