From 1bb21705baa606c037a43c2c1b88d8f1a95c6904 Mon Sep 17 00:00:00 2001
From: Ilya Skriblovsky <ilyaskriblovsky@gmail.com>
Date: Tue, 15 Oct 2024 00:00:11 +0300
Subject: [PATCH] Move protocol constants initialization to protocol.py

---
 txmongo/connection.py | 25 ++-----------------------
 txmongo/protocol.py   | 28 +++++++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/txmongo/connection.py b/txmongo/connection.py
index 26075d5d..448acf60 100755
--- a/txmongo/connection.py
+++ b/txmongo/connection.py
@@ -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,
@@ -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
@@ -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")
diff --git a/txmongo/protocol.py b/txmongo/protocol.py
index 48381fc5..75f99343 100644
--- a/txmongo/protocol.py
+++ b/txmongo/protocol.py
@@ -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
@@ -28,6 +29,7 @@
 from bson import SON, Binary, CodecOptions
 from pymongo.errors import (
     AutoReconnect,
+    ConfigurationError,
     ConnectionFailure,
     CursorNotFound,
     NotPrimaryError,
@@ -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
@@ -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)
 
@@ -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))