Skip to content

Commit

Permalink
Move protocol constants initialization to protocol.py
Browse files Browse the repository at this point in the history
  • Loading branch information
IlyaSkriblovsky committed Oct 16, 2024
1 parent 4615ed3 commit 1bb2170
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 24 deletions.
25 changes: 2 additions & 23 deletions txmongo/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,6 @@
from txmongo.protocol import MongoProtocol
from txmongo.utils import get_err, timeout

DEFAULT_MAX_BSON_SIZE = 16777216
DEFAULT_MAX_WRITE_BATCH_SIZE = 1000
DEFAULT_MAX_MESSAGE_SIZE = 48000000


_PRIMARY_READ_PREFERENCES = {
ReadPreference.PRIMARY.mode,
ReadPreference.PRIMARY_PREFERRED.mode,
Expand Down Expand Up @@ -77,7 +72,7 @@ def _initializeProto(self, proto):

@staticmethod
@timeout
def __send_ismaster(proto, _deadline):
def __send_ismaster(proto, _deadline) -> defer.Deferred[dict]:
return proto.send_op_query_command("admin", {"ismaster": 1})

@defer.inlineCallbacks
Expand Down Expand Up @@ -110,23 +105,7 @@ def configure(self, proto: MongoProtocol):
msg = "TxMongo: Mongo instance does not match requested replicaSet."
raise ConfigurationError(msg)

# FIXME: move this initialization to MongoProtocol class
# Track max bson object size limit.
proto.max_bson_size = config.get("maxBsonObjectSize", DEFAULT_MAX_BSON_SIZE)
proto.max_write_batch_size = config.get(
"maxWriteBatchSize", DEFAULT_MAX_WRITE_BATCH_SIZE
)
proto.max_message_size = config.get(
"maxMessageSizeBytes", DEFAULT_MAX_MESSAGE_SIZE
)
proto.set_wire_versions(
config.get("minWireVersion", 0), config.get("maxWireVersion", 0)
)

# MongoDB < 4.0
if proto.max_wire_version < 7:
warnings.warn("TxMongo: MongoDB version <4.0 is not supported")
raise ConfigurationError("TxMongo: MongoDB version <4.0 is not supported")
proto.init_from_hello_response(config)

# Track the other hosts in the replica set.
hosts = config.get("hosts")
Expand Down
28 changes: 27 additions & 1 deletion txmongo/protocol.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import hmac
import logging
import struct
import warnings
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass, field
from hashlib import sha1
Expand All @@ -28,6 +29,7 @@
from bson import SON, Binary, CodecOptions
from pymongo.errors import (
AutoReconnect,
ConfigurationError,
ConnectionFailure,
CursorNotFound,
NotPrimaryError,
Expand Down Expand Up @@ -67,6 +69,11 @@ def _digest(msg, mac=mac):
return to_bytes(_ui, mac.digest_size, "big")


DEFAULT_MAX_BSON_SIZE = 16777216
DEFAULT_MAX_WRITE_BATCH_SIZE = 1000
DEFAULT_MAX_MESSAGE_SIZE = 48000000


INT_MAX = 2147483647

OP_REPLY = 1
Expand Down Expand Up @@ -435,6 +442,25 @@ def __init__(self):

self.__auth_lock = defer.DeferredLock()

def init_from_hello_response(self, config: dict) -> None:
"""`config` is a dict from server hello response"""
# Track max bson object size limit.
self.max_bson_size = config.get("maxBsonObjectSize", DEFAULT_MAX_BSON_SIZE)
self.max_write_batch_size = config.get(
"maxWriteBatchSize", DEFAULT_MAX_WRITE_BATCH_SIZE
)
self.max_message_size = config.get(
"maxMessageSizeBytes", DEFAULT_MAX_MESSAGE_SIZE
)
self.set_wire_versions(
config.get("minWireVersion", 0), config.get("maxWireVersion", 0)
)

# MongoDB < 4.0
if self.max_wire_version < 7:
warnings.warn("TxMongo: MongoDB version <4.0 is not supported")
raise ConfigurationError("TxMongo: MongoDB version <4.0 is not supported")

def inflight(self):
return len(self.__deferreds)

Expand Down Expand Up @@ -558,7 +584,7 @@ def set_wire_versions(self, min_wire_version, max_wire_version):
self.min_wire_version = min_wire_version
self.max_wire_version = max_wire_version

def send_op_query_command(self, database, query):
def send_op_query_command(self, database, query) -> defer.Deferred[dict]:
cmd_collection = str(database) + ".$cmd"
return self.send_query(
Query(collection=cmd_collection, query=bson.encode(query))
Expand Down

0 comments on commit 1bb2170

Please sign in to comment.