Skip to content

Commit 6fa2e40

Browse files
ajdavisajdavis
committed
PYTHON-1332 - Send lsid with all commands
1 parent c1ec855 commit 6fa2e40

File tree

6 files changed

+157
-109
lines changed

6 files changed

+157
-109
lines changed

‎pymongo/client_session.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,13 @@ class _ServerSessionPool(collections.deque):
149149
This class is not thread-safe, access it while holding the Topology lock.
150150
"""
151151
def get_server_session(self, session_timeout_minutes):
152+
# Although the Driver Sessions Spec says we only clear stale sessions
153+
# in return_server_session, PyMongo can't take a lock when returning
154+
# sessions from a __del__ method (like in Cursor.__die), so it can't
155+
# clear stale sessions there. In case many sessions were returned via
156+
# __del__, check for stale sessions here too.
157+
self._clear_stale(session_timeout_minutes)
158+
152159
# The most recently used sessions are on the left.
153160
while self:
154161
s = self.popleft()
@@ -158,13 +165,15 @@ def get_server_session(self, session_timeout_minutes):
158165
return _ServerSession()
159166

160167
def return_server_session(self, server_session, session_timeout_minutes):
168+
self._clear_stale(session_timeout_minutes)
169+
if not server_session.timed_out(session_timeout_minutes):
170+
self.appendleft(server_session)
171+
172+
def _clear_stale(self, session_timeout_minutes):
161173
# Clear stale sessions. The least recently used are on the right.
162174
while self:
163175
if self[-1].timed_out(session_timeout_minutes):
164176
self.pop()
165177
else:
166178
# The remaining sessions also haven't timed out.
167179
break
168-
169-
if not server_session.timed_out(session_timeout_minutes):
170-
self.appendleft(server_session)

‎pymongo/collection.py‎

Lines changed: 93 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1398,6 +1398,9 @@ def find_raw_batches(self, *args, **kwargs):
13981398
>>> for batch in cursor:
13991399
... print(bson.decode_all(batch))
14001400
1401+
Unlike most PyMongo methods, this method sends no session id to the
1402+
server.
1403+
14011404
.. versionadded:: 3.6
14021405
"""
14031406
return RawBatchCursor(self, *args, **kwargs)
@@ -1443,6 +1446,9 @@ def parallel_scan(self, num_cursors, session=None, **kwargs):
14431446
- `**kwargs`: additional options for the parallelCollectionScan
14441447
command can be passed as keyword arguments.
14451448
1449+
Unlike most PyMongo methods, this method sends no session id to the
1450+
server unless an explicit ``session`` parameter is passed.
1451+
14461452
.. note:: Requires server version **>= 2.5.5**.
14471453
14481454
.. versionchanged:: 3.6
@@ -1461,14 +1467,19 @@ def parallel_scan(self, num_cursors, session=None, **kwargs):
14611467
('numCursors', num_cursors)])
14621468
cmd.update(kwargs)
14631469

1464-
s = self.__database.client._ensure_session(session)
14651470
with self._socket_for_reads() as (sock_info, slave_ok):
1466-
result = self._command(sock_info, cmd, slave_ok,
1467-
read_concern=self.read_concern,
1468-
session=s)
1471+
# Avoid auto-injecting a session.
1472+
result = sock_info.command(
1473+
self.__database.name,
1474+
cmd,
1475+
slave_ok,
1476+
self.read_preference,
1477+
self.codec_options,
1478+
read_concern=self.read_concern,
1479+
session=session)
14691480

14701481
return [CommandCursor(self, cursor['cursor'], sock_info.address,
1471-
session=s, session_owned=session is None)
1482+
session=session, explicit_session=True)
14721483
for cursor in result['cursors']]
14731484

14741485
def _count(self, cmd, collation=None, session=None):
@@ -1890,20 +1901,20 @@ def list_indexes(self, session=None):
18901901
with self._socket_for_primary_reads() as (sock_info, slave_ok):
18911902
cmd = SON([("listIndexes", self.__name), ("cursor", {})])
18921903
if sock_info.max_wire_version > 2:
1893-
s = self.__database.client._ensure_session(session)
1894-
try:
1895-
cursor = self._command(sock_info, cmd, slave_ok,
1896-
ReadPreference.PRIMARY,
1897-
codec_options,
1898-
session=s)["cursor"]
1899-
except OperationFailure as exc:
1900-
# Ignore NamespaceNotFound errors to match the behavior
1901-
# of reading from *.system.indexes.
1902-
if exc.code != 26:
1903-
raise
1904-
cursor = {'id': 0, 'firstBatch': []}
1905-
return CommandCursor(coll, cursor, sock_info.address,
1906-
session=s, session_owned=session is None)
1904+
with self.__database.client._tmp_session(session, False) as s:
1905+
try:
1906+
cursor = self._command(sock_info, cmd, slave_ok,
1907+
ReadPreference.PRIMARY,
1908+
codec_options,
1909+
session=s)["cursor"]
1910+
except OperationFailure as exc:
1911+
# Ignore NamespaceNotFound errors to match the behavior
1912+
# of reading from *.system.indexes.
1913+
if exc.code != 26:
1914+
raise
1915+
cursor = {'id': 0, 'firstBatch': []}
1916+
return CommandCursor(coll, cursor, sock_info.address, session=s,
1917+
explicit_session=session is not None)
19071918
else:
19081919
namespace = _UJOIN % (self.__database.name, "system.indexes")
19091920
res = helpers._first_batch(
@@ -1995,7 +2006,7 @@ def options(self, session=None):
19952006
return options
19962007

19972008
def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
1998-
**kwargs):
2009+
explicit_session, **kwargs):
19992010
common.validate_list('pipeline', pipeline)
20002011

20012012
if "explain" in kwargs:
@@ -2028,47 +2039,56 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
20282039
cmd['writeConcern'] = self.write_concern.document
20292040

20302041
cmd.update(kwargs)
2031-
session_owned = session is None
2032-
s = self.__database.client._ensure_session(session)
2033-
try:
2034-
# Apply this Collection's read concern if $out is not in the
2035-
# pipeline.
2036-
if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
2037-
if dollar_out:
2038-
result = self._command(sock_info, cmd, slave_ok,
2039-
parse_write_concern_error=True,
2040-
collation=collation,
2041-
session=s)
2042-
else:
2043-
result = self._command(sock_info, cmd, slave_ok,
2044-
read_concern=self.read_concern,
2045-
collation=collation,
2046-
session=s)
2042+
# Apply this Collection's read concern if $out is not in the
2043+
# pipeline.
2044+
if sock_info.max_wire_version >= 4 and 'readConcern' not in cmd:
2045+
if dollar_out:
2046+
# Avoid auto-injecting a session.
2047+
result = sock_info.command(
2048+
self.__database.name,
2049+
cmd,
2050+
slave_ok,
2051+
self.read_preference,
2052+
self.codec_options,
2053+
parse_write_concern_error=True,
2054+
collation=collation,
2055+
session=session)
20472056
else:
2048-
result = self._command(sock_info, cmd, slave_ok,
2049-
parse_write_concern_error=dollar_out,
2050-
collation=collation,
2051-
session=s)
2057+
result = sock_info.command(
2058+
self.__database.name,
2059+
cmd,
2060+
slave_ok,
2061+
ReadPreference.PRIMARY,
2062+
self.codec_options,
2063+
read_concern=self.read_concern,
2064+
collation=collation,
2065+
session=session)
2066+
else:
2067+
result = sock_info.command(
2068+
self.__database.name,
2069+
cmd,
2070+
slave_ok,
2071+
self.read_preference,
2072+
self.codec_options,
2073+
parse_write_concern_error=dollar_out,
2074+
collation=collation,
2075+
session=session)
20522076

2053-
if "cursor" in result:
2054-
cursor = result["cursor"]
2055-
else:
2056-
# Pre-MongoDB 2.6. Fake a cursor.
2057-
cursor = {
2058-
"id": 0,
2059-
"firstBatch": result["result"],
2060-
"ns": self.full_name,
2061-
}
2062-
2063-
return cursor_class(
2064-
self, cursor, sock_info.address,
2065-
batch_size=batch_size or 0,
2066-
max_await_time_ms=max_await_time_ms,
2067-
session=s, session_owned=session_owned)
2068-
except Exception:
2069-
if session_owned:
2070-
s.end_session()
2071-
raise
2077+
if "cursor" in result:
2078+
cursor = result["cursor"]
2079+
else:
2080+
# Pre-MongoDB 2.6. Fake a cursor.
2081+
cursor = {
2082+
"id": 0,
2083+
"firstBatch": result["result"],
2084+
"ns": self.full_name,
2085+
}
2086+
2087+
return cursor_class(
2088+
self, cursor, sock_info.address,
2089+
batch_size=batch_size or 0,
2090+
max_await_time_ms=max_await_time_ms,
2091+
session=session, explicit_session=explicit_session)
20722092

20732093
def aggregate(self, pipeline, session=None, **kwargs):
20742094
"""Perform an aggregation using the aggregation framework on this
@@ -2148,13 +2168,15 @@ def aggregate(self, pipeline, session=None, **kwargs):
21482168
.. _aggregate command:
21492169
https://docs.mongodb.com/manual/reference/command/aggregate
21502170
"""
2151-
return self._aggregate(pipeline,
2152-
CommandCursor,
2153-
kwargs.get('batchSize'),
2154-
session=session,
2155-
**kwargs)
2156-
2157-
def aggregate_raw_batches(self, pipeline, session=None, **kwargs):
2171+
with self.__database.client._tmp_session(session, close=False) as s:
2172+
return self._aggregate(pipeline,
2173+
CommandCursor,
2174+
kwargs.get('batchSize'),
2175+
session=s,
2176+
explicit_session=session is not None,
2177+
**kwargs)
2178+
2179+
def aggregate_raw_batches(self, pipeline, **kwargs):
21582180
"""Perform an aggregation and retrieve batches of raw BSON.
21592181
21602182
Takes the same parameters as :meth:`aggregate` but returns a
@@ -2171,13 +2193,13 @@ def aggregate_raw_batches(self, pipeline, session=None, **kwargs):
21712193
>>> for batch in cursor:
21722194
... print(bson.decode_all(batch))
21732195
2196+
Unlike most PyMongo methods, this method sends no session id to the
2197+
server.
2198+
21742199
.. versionadded:: 3.6
21752200
"""
2176-
return self._aggregate(pipeline,
2177-
RawBatchCommandCursor,
2178-
0,
2179-
session=session,
2180-
**kwargs)
2201+
return self._aggregate(pipeline, RawBatchCommandCursor, 0,
2202+
None, False, **kwargs)
21812203

21822204
def watch(self, pipeline=None, full_document='default', resume_after=None,
21832205
max_await_time_ms=None, batch_size=None, collation=None):

‎pymongo/command_cursor.py‎

Lines changed: 26 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class CommandCursor(object):
3636

3737
def __init__(self, collection, cursor_info, address, retrieved=0,
3838
batch_size=0, max_await_time_ms=None, session=None,
39-
session_owned=False):
39+
explicit_session=False):
4040
"""Create a new command cursor.
4141
4242
The parameter 'retrieved' is unused.
@@ -47,21 +47,24 @@ def __init__(self, collection, cursor_info, address, retrieved=0,
4747
self.__address = address
4848
self.__data = deque(cursor_info['firstBatch'])
4949
self.__batch_size = batch_size
50-
if (not isinstance(max_await_time_ms, integer_types)
51-
and max_await_time_ms is not None):
52-
raise TypeError("max_await_time_ms must be an integer or None")
5350
self.__max_await_time_ms = max_await_time_ms
51+
self.__session = session
52+
self.__explicit_session = explicit_session
5453
self.__killed = (self.__id == 0)
54+
if self.__killed:
55+
self.__end_session(True)
5556

5657
if "ns" in cursor_info:
5758
self.__ns = cursor_info["ns"]
5859
else:
5960
self.__ns = collection.full_name
6061

61-
self.__session = session
62-
self.__session_owned = session_owned
6362
self.batch_size(batch_size)
6463

64+
if (not isinstance(max_await_time_ms, integer_types)
65+
and max_await_time_ms is not None):
66+
raise TypeError("max_await_time_ms must be an integer or None")
67+
6568
def __del__(self):
6669
if self.__id and not self.__killed:
6770
self.__die()
@@ -80,7 +83,10 @@ def __die(self, synchronous=False):
8083
self.__collection.database.client.close_cursor(
8184
self.__id, address)
8285
self.__killed = True
83-
if self.__session and self.__session_owned:
86+
self.__end_session(synchronous)
87+
88+
def __end_session(self, synchronous):
89+
if self.__session and not self.__explicit_session:
8490
self.__session._end_session(lock=synchronous)
8591
self.__session = None
8692

@@ -116,6 +122,10 @@ def batch_size(self, batch_size):
116122
def __send_message(self, operation):
117123
"""Send a getmore message and handle the response.
118124
"""
125+
def kill():
126+
self.__killed = True
127+
self.__end_session(True)
128+
119129
client = self.__collection.database.client
120130
listeners = client._event_listeners
121131
publish = listeners.enabled_for_commands
@@ -127,7 +137,7 @@ def __send_message(self, operation):
127137
# or to another server. It can cause a _pinValue
128138
# assertion on some server releases if we get here
129139
# due to a socket timeout.
130-
self.__killed = True
140+
kill()
131141
raise
132142

133143
cmd_duration = response.duration
@@ -144,7 +154,7 @@ def __send_message(self, operation):
144154
helpers._check_command_response(doc['data'][0])
145155

146156
except OperationFailure as exc:
147-
self.__killed = True
157+
kill()
148158

149159
if publish:
150160
duration = (datetime.datetime.now() - start) + cmd_duration
@@ -155,7 +165,7 @@ def __send_message(self, operation):
155165
except NotMasterError as exc:
156166
# Don't send kill cursors to another server after a "not master"
157167
# error. It's completely pointless.
158-
self.__killed = True
168+
kill()
159169

160170
if publish:
161171
duration = (datetime.datetime.now() - start) + cmd_duration
@@ -191,7 +201,7 @@ def __send_message(self, operation):
191201
duration, res, "getMore", rqst_id, self.__address)
192202

193203
if self.__id == 0:
194-
self.__killed = True
204+
kill()
195205
self.__data = deque(documents)
196206

197207
def _unpack_response(self, response, cursor_id, codec_options):
@@ -219,6 +229,7 @@ def _refresh(self):
219229
self.__max_await_time_ms))
220230
else: # Cursor id is zero nothing else to return
221231
self.__killed = True
232+
self.__end_session(True)
222233

223234
return len(self.__data)
224235

@@ -258,7 +269,7 @@ def session(self):
258269
259270
.. versionadded:: 3.6
260271
"""
261-
if not self.__session_owned:
272+
if self.__explicit_session:
262273
return self.__session
263274

264275
def __iter__(self):
@@ -288,7 +299,8 @@ class RawBatchCommandCursor(CommandCursor):
288299
_getmore_class = _RawBatchGetMore
289300

290301
def __init__(self, collection, cursor_info, address, retrieved=0,
291-
batch_size=0, max_await_time_ms=None, session=None):
302+
batch_size=0, max_await_time_ms=None, session=None,
303+
explicit_session=False):
292304
"""Create a new cursor / iterator over raw batches of BSON data.
293305
294306
Should not be called directly by application developers -
@@ -300,7 +312,7 @@ def __init__(self, collection, cursor_info, address, retrieved=0,
300312
assert not cursor_info.get('firstBatch')
301313
super(RawBatchCommandCursor, self).__init__(
302314
collection, cursor_info, address, retrieved, batch_size,
303-
max_await_time_ms, session)
315+
max_await_time_ms, session, explicit_session)
304316

305317
def _unpack_response(self, response, cursor_id, codec_options):
306318
return helpers._raw_response(response, cursor_id)

0 commit comments

Comments
 (0)