Skip to content

Commit 8edb2d3

Browse files
committed
Fix tests to account for later session support checking
1 parent f94ad1b commit 8edb2d3

File tree

6 files changed

+44
-62
lines changed

6 files changed

+44
-62
lines changed

pymongo/client_session.py

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -965,10 +965,12 @@ def _txn_read_preference(self) -> Optional[_ServerMode]:
965965
return self._transaction.opts.read_preference
966966
return None
967967

968-
def _materialize(self) -> None:
968+
def _materialize(self, logical_session_timeout_minutes: float) -> None:
969969
if isinstance(self._server_session, _EmptyServerSession):
970970
old = self._server_session
971-
self._server_session = self._client._topology.get_server_session()
971+
self._server_session = self._client._topology.get_server_session(
972+
logical_session_timeout_minutes, old
973+
)
972974
if old.started_retryable_write:
973975
self._server_session.inc_transaction_id()
974976

@@ -984,7 +986,7 @@ def _apply_to(
984986
raise ConfigurationError("Sessions are not supported by this MongoDB deployment")
985987
return
986988
self._check_ended()
987-
self._materialize()
989+
self._materialize(conn.logical_session_timeout_minutes)
988990
if self.options.snapshot:
989991
self._update_read_concern(command, conn)
990992

@@ -1036,11 +1038,12 @@ def __copy__(self) -> NoReturn:
10361038

10371039

10381040
class _EmptyServerSession:
1039-
__slots__ = "dirty", "started_retryable_write"
1041+
__slots__ = "dirty", "started_retryable_write", "session_id"
10401042

10411043
def __init__(self) -> None:
10421044
self.dirty = False
10431045
self.started_retryable_write = False
1046+
self.session_id = {"id": Binary(uuid.uuid4().bytes, 4)}
10441047

10451048
def mark_dirty(self) -> None:
10461049
self.dirty = True
@@ -1050,9 +1053,9 @@ def inc_transaction_id(self) -> None:
10501053

10511054

10521055
class _ServerSession:
1053-
def __init__(self, generation: int):
1056+
def __init__(self, generation: int, session_id: Optional[dict[str, Binary]] = None):
10541057
# Ensure id is type 4, regardless of CodecOptions.uuid_representation.
1055-
self.session_id = {"id": Binary(uuid.uuid4().bytes, 4)}
1058+
self.session_id = session_id or {"id": Binary(uuid.uuid4().bytes, 4)}
10561059
self.last_use = time.monotonic()
10571060
self._transaction_id = 0
10581061
self.dirty = False
@@ -1101,14 +1104,23 @@ def pop_all(self) -> list[_ServerSession]:
11011104
ids.append(self.pop().session_id)
11021105
return ids
11031106

1104-
def get_server_session(self) -> _ServerSession:
1107+
def get_server_session(
1108+
self, session_timeout_minutes: float, old: _EmptyServerSession
1109+
) -> _ServerSession:
1110+
# Although the Driver Sessions Spec says we only clear stale sessions
1111+
# in return_server_session, PyMongo can't take a lock when returning
1112+
# sessions from a __del__ method (like in Cursor.__die), so it can't
1113+
# clear stale sessions there. In case many sessions were returned via
1114+
# __del__, check for stale sessions here too.
1115+
self._clear_stale(session_timeout_minutes)
1116+
11051117
# The most recently used sessions are on the left.
11061118
while self:
1107-
# s = self.popleft()
1108-
# # if not s.timed_out(session_timeout_minutes):
1109-
return self.popleft()
1119+
s = self.popleft()
1120+
if not s.timed_out(session_timeout_minutes):
1121+
return s
11101122

1111-
return _ServerSession(self.generation)
1123+
return _ServerSession(self.generation, old.session_id)
11121124

11131125
def return_server_session(
11141126
self, server_session: _ServerSession, session_timeout_minutes: Optional[float]

pymongo/mongo_client.py

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1750,10 +1750,7 @@ def _process_periodic_tasks(self) -> None:
17501750
helpers._handle_exception()
17511751

17521752
def __start_session(self, implicit: bool, **kwargs: Any) -> ClientSession:
1753-
if implicit:
1754-
server_session: Union[_EmptyServerSession, _ServerSession] = _EmptyServerSession()
1755-
else:
1756-
server_session = self._get_server_session()
1753+
server_session = _EmptyServerSession()
17571754
opts = client_session.SessionOptions(**kwargs)
17581755
return client_session.ClientSession(self, server_session, opts, implicit)
17591756

@@ -1786,10 +1783,6 @@ def start_session(
17861783
snapshot=snapshot,
17871784
)
17881785

1789-
def _get_server_session(self) -> _ServerSession:
1790-
"""Internal: start or resume a _ServerSession."""
1791-
return self._topology.get_server_session()
1792-
17931786
def _return_server_session(
17941787
self, server_session: Union[_ServerSession, _EmptyServerSession], lock: bool
17951788
) -> None:

pymongo/pool.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -752,6 +752,7 @@ def __init__(
752752
self.active = False
753753
self.last_timeout = self.opts.socket_timeout
754754
self.connect_rtt = 0.0
755+
self.logical_session_timeout_minutes: float = 0.0
755756

756757
def set_conn_timeout(self, timeout: Optional[float]) -> None:
757758
"""Cache last timeout to avoid duplicate calls to conn.settimeout."""
@@ -870,6 +871,7 @@ def _hello(
870871
self.max_message_size = hello.max_message_size
871872
self.max_write_batch_size = hello.max_write_batch_size
872873
self.supports_sessions = hello.logical_session_timeout_minutes is not None
874+
self.logical_session_timeout_minutes = hello.logical_session_timeout_minutes or 0.0
873875
self.hello_ok = hello.hello_ok
874876
self.is_repl = hello.server_type in (
875877
SERVER_TYPE.RSPrimary,

pymongo/topology.py

Lines changed: 5 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,8 @@
2525
from typing import TYPE_CHECKING, Any, Callable, Mapping, Optional, cast
2626

2727
from pymongo import _csot, common, helpers, periodic_executor
28-
from pymongo.client_session import _ServerSession, _ServerSessionPool
28+
from pymongo.client_session import _EmptyServerSession, _ServerSession, _ServerSessionPool
2929
from pymongo.errors import (
30-
ConfigurationError,
3130
ConnectionFailure,
3231
InvalidOperation,
3332
NetworkTimeout,
@@ -47,7 +46,6 @@
4746
Selection,
4847
any_server_selector,
4948
arbiter_server_selector,
50-
readable_server_selector,
5149
secondary_server_selector,
5250
writable_server_selector,
5351
)
@@ -579,37 +577,12 @@ def pop_all_sessions(self) -> list[_ServerSession]:
579577
with self._lock:
580578
return self._session_pool.pop_all()
581579

582-
def _check_implicit_session_support(self) -> None:
583-
with self._lock:
584-
self._check_session_support()
585-
586-
def _check_session_support(self) -> float:
587-
"""Internal check for session support on clusters."""
588-
if self._settings.load_balanced:
589-
# Sessions never time out in load balanced mode.
590-
return float("inf")
591-
session_timeout = self._description.logical_session_timeout_minutes
592-
if session_timeout is None:
593-
# Maybe we need an initial scan? Can raise ServerSelectionError.
594-
if self._description.topology_type == TOPOLOGY_TYPE.Single:
595-
if not self._description.has_known_servers:
596-
self._select_servers_loop(
597-
any_server_selector, self.get_server_selection_timeout(), None
598-
)
599-
elif not self._description.readable_servers:
600-
self._select_servers_loop(
601-
readable_server_selector, self.get_server_selection_timeout(), None
602-
)
603-
604-
session_timeout = self._description.logical_session_timeout_minutes
605-
if session_timeout is None:
606-
raise ConfigurationError("Sessions are not supported by this MongoDB deployment")
607-
return session_timeout
608-
609-
def get_server_session(self) -> _ServerSession:
580+
def get_server_session(
581+
self, session_timeout_minutes: float, old: _EmptyServerSession
582+
) -> _ServerSession:
610583
"""Start or resume a server session, or raise ConfigurationError."""
611584
with self._lock:
612-
return self._session_pool.get_server_session()
585+
return self._session_pool.get_server_session(session_timeout_minutes, old)
613586

614587
def return_server_session(self, server_session: _ServerSession, lock: bool) -> None:
615588
if lock:

test/test_encryption.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,6 @@
2222
import socket
2323
import socketserver
2424
import ssl
25-
import subprocess
2625
import sys
2726
import textwrap
2827
import traceback
@@ -3009,7 +3008,6 @@ class TestNoSessionsSupport(EncryptionIntegrationTest):
30093008
MONGOCRYPTD_PORT = 27020
30103009

30113010
@classmethod
3012-
@client_context.require_sessions
30133011
@unittest.skipIf(os.environ.get("TEST_CRYPT_SHARED"), "crypt_shared lib is installed")
30143012
def setUpClass(cls):
30153013
super().setUpClass()
@@ -3021,10 +3019,10 @@ def tearDownClass(cls):
30213019

30223020
def setUp(self) -> None:
30233021
self.listener = OvertCommandListener()
3024-
self.listener.reset()
30253022
self.mongocryptd_client = MongoClient(
30263023
f"mongodb://localhost:{self.MONGOCRYPTD_PORT}", event_listeners=[self.listener]
30273024
)
3025+
self.addCleanup(self.mongocryptd_client.close)
30283026

30293027
hello = self.mongocryptd_client.db.command("hello")
30303028
self.assertNotIn("logicalSessionTimeoutMinutes", hello)
@@ -3040,16 +3038,18 @@ def test_implicit_session_ignored_when_unsupported(self):
30403038
self.mongocryptd_client.db.test.insert_one({"x": 1})
30413039

30423040
self.assertNotIn("lsid", self.listener.started_events[1].command)
3043-
self.mongocryptd_client.close()
30443041

30453042
def test_explicit_session_errors_when_unsupported(self):
30463043
self.listener.reset()
30473044
with self.mongocryptd_client.start_session() as s:
3048-
with self.assertRaises(ConfigurationError):
3045+
with self.assertRaisesRegex(
3046+
ConfigurationError, r"Sessions are not supported by this MongoDB deployment"
3047+
):
30493048
self.mongocryptd_client.db.test.find_one(session=s)
3050-
with self.assertRaises(ConfigurationError):
3049+
with self.assertRaisesRegex(
3050+
ConfigurationError, r"Sessions are not supported by this MongoDB deployment"
3051+
):
30513052
self.mongocryptd_client.db.test.insert_one({"x": 1}, session=s)
3052-
self.mongocryptd_client.close()
30533053

30543054

30553055
if __name__ == "__main__":

test/test_session.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,16 +116,12 @@ def _test_ops(self, client, *ops):
116116

117117
for f, args, kw in ops:
118118
with client.start_session() as s:
119-
last_use = s._server_session.last_use
120-
start = time.monotonic()
121-
self.assertLessEqual(last_use, start)
122119
listener.reset()
123120
# In case "f" modifies its inputs.
124121
args = copy.copy(args)
125122
kw = copy.copy(kw)
126123
kw["session"] = s
127124
f(*args, **kw)
128-
self.assertGreaterEqual(s._server_session.last_use, start)
129125
self.assertGreaterEqual(len(listener.started_events), 1)
130126
for event in listener.started_events:
131127
self.assertTrue(
@@ -239,16 +235,20 @@ def test_pool_lifo(self):
239235
# "Pool is LIFO" test from Driver Sessions Spec.
240236
a = self.client.start_session()
241237
b = self.client.start_session()
238+
self.client.admin.command("ping", session=a)
239+
self.client.admin.command("ping", session=b)
242240
a_id = a.session_id
243241
b_id = b.session_id
244242
a.end_session()
245243
b.end_session()
246244

247245
s = self.client.start_session()
246+
self.client.admin.command("ping", session=s)
248247
self.assertEqual(b_id, s.session_id)
249248
self.assertNotEqual(a_id, s.session_id)
250249

251250
s2 = self.client.start_session()
251+
self.client.admin.command("ping", session=s2)
252252
self.assertEqual(a_id, s2.session_id)
253253
self.assertNotEqual(b_id, s2.session_id)
254254

@@ -274,6 +274,8 @@ def test_end_sessions(self):
274274
client = rs_or_single_client(event_listeners=[listener])
275275
# Start many sessions.
276276
sessions = [client.start_session() for _ in range(_MAX_END_SESSIONS + 1)]
277+
for s in sessions:
278+
client.admin.command("ping", session=s)
277279
for s in sessions:
278280
s.end_session()
279281

0 commit comments

Comments
 (0)