Skip to content

Commit 9df87b6

Browse files
ajdavisajdavis
committed
prototype transaction tests
1 parent 04de169 commit 9df87b6

18 files changed

+3716
-33
lines changed

‎pymongo/bulk.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -260,15 +260,13 @@ def _execute_command(self, generator, write_concern, session,
260260
cmd['writeConcern'] = write_concern.document
261261
if self.bypass_doc_val and sock_info.max_wire_version >= 4:
262262
cmd['bypassDocumentValidation'] = True
263-
if session:
264-
cmd['lsid'] = session._use_lsid()
265263
bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id,
266264
listeners, session)
267265

268266
results = []
269267
while run.idx_offset < len(run.ops):
270-
if session and retryable:
271-
cmd['txnNumber'] = session._transaction_id()
268+
if session:
269+
session._apply_to(cmd, retryable)
272270
sock_info.send_cluster_time(cmd, session, client)
273271
check_keys = run.op_type == _INSERT
274272
ops = islice(run.ops, run.idx_offset, None)
@@ -278,6 +276,8 @@ def _execute_command(self, generator, write_concern, session,
278276
self.collection.codec_options, bwc)
279277
if not to_send:
280278
raise InvalidOperation("cannot do an empty bulk write")
279+
if session:
280+
session._advance_statement_id(len(to_send))
281281
result = bwc.write_command(request_id, msg, to_send)
282282
client._receive_cluster_time(result, session)
283283
results.append((run.idx_offset, result))

‎pymongo/client_session.py‎

Lines changed: 143 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -61,15 +61,51 @@ class SessionOptions(object):
6161
:Parameters:
6262
- `causal_consistency` (optional): If True (the default), read
6363
operations are causally ordered within the session.
64+
- `auto_start_transaction` (optional): If True, any operation using
65+
the session begins a transaction if none is in progress.
6466
"""
65-
def __init__(self, causal_consistency=True):
67+
# TODO: accept a TransactionOptions.
68+
def __init__(self,
69+
causal_consistency=True,
70+
auto_start_transaction=False):
6671
self._causal_consistency = causal_consistency
72+
self._auto_start_transaction = auto_start_transaction
6773

6874
@property
6975
def causal_consistency(self):
7076
"""Whether causal consistency is configured."""
7177
return self._causal_consistency
7278

79+
@property
80+
def auto_start_transaction(self):
81+
"""Whether the session is configured to always start a transaction."""
82+
return self._auto_start_transaction
83+
84+
85+
class TransactionOptions(object):
86+
"""Options for :meth:`ClientSession.start_transaction`.
87+
88+
:Parameters:
89+
- `read_concern`: The :class:`~read_concern.ReadConcern` to use for this
90+
transaction.
91+
- `write_concern`: The :class:`~write_concern.WriteConcern` to use for
92+
this transaction.
93+
"""
94+
def __init__(self, read_concern=None, write_concern=None):
95+
# TODO: validate arguments.
96+
self._read_concern = read_concern
97+
self._write_concern = write_concern
98+
99+
@property
100+
def read_concern(self):
101+
"""This transaction's :class:`~read_concern.ReadConcern`."""
102+
return self._read_concern
103+
104+
@property
105+
def write_concern(self):
106+
"""This transaction's :class:`~write_concern.WriteConcern`."""
107+
return self._write_concern
108+
73109

74110
class ClientSession(object):
75111
"""A session for ordering sequential operations."""
@@ -81,27 +117,41 @@ def __init__(self, client, server_session, options, authset):
81117
self._authset = authset
82118
self._cluster_time = None
83119
self._operation_time = None
120+
self._transaction_options = None # Current transaction's options.
121+
if self.options.auto_start_transaction:
122+
# TODO: Get transaction options from self.options.
123+
self._transaction_options = TransactionOptions()
124+
self._server_session.start_transaction()
84125

85126
def end_session(self):
86-
"""Finish this session.
127+
"""Finish this session. If a transaction has started, abort it.
87128
88129
It is an error to use the session or any derived
89130
:class:`~pymongo.database.Database`,
90131
:class:`~pymongo.collection.Collection`, or
91132
:class:`~pymongo.cursor.Cursor` after the session has ended.
92133
"""
93-
self._end_session(True)
134+
self._end_session(lock=True, abort_txn=True)
94135

95-
def _end_session(self, lock):
136+
def _end_session(self, lock, abort_txn):
96137
if self._server_session is not None:
97-
self._client._return_server_session(self._server_session, lock)
98-
self._server_session = None
138+
try:
139+
if self.in_transaction:
140+
if abort_txn:
141+
self.abort_txn()
142+
else:
143+
self.commit_transaction()
144+
finally:
145+
self._client._return_server_session(self._server_session, lock)
146+
self._server_session = None
99147

100148
def __enter__(self):
101149
return self
102150

103151
def __exit__(self, exc_type, exc_val, exc_tb):
104-
self.end_session()
152+
# Abort when exiting with an exception, otherwise commit.
153+
# TODO: test and document this.
154+
self._end_session(lock=True, abort_txn=exc_val is not None)
105155

106156
@property
107157
def client(self):
@@ -137,6 +187,45 @@ def operation_time(self):
137187
"""
138188
return self._operation_time
139189

190+
def start_transaction(self, **kwargs):
191+
"""Start a multi-statement transaction.
192+
193+
Takes the same arguments as :class:`TransactionOptions`.
194+
195+
Do not use this method if the session is configured to automatically
196+
start a transaction.
197+
"""
198+
self._transaction_options = TransactionOptions(**kwargs)
199+
self._server_session.start_transaction()
200+
201+
def commit_transaction(self):
202+
"""Commit a multi-statement transaction."""
203+
self._finish_transaction("commitTransaction")
204+
205+
def abort_txn(self):
206+
"""Abort a multi-statement transaction."""
207+
assert False, "Not implemented" # Await server.
208+
self._finish_transaction("abortTransaction")
209+
210+
def _finish_transaction(self, command_name):
211+
if (self.options.auto_start_transaction
212+
and self._server_session.statement_id == 0):
213+
# Not really started.
214+
return
215+
216+
try:
217+
# TODO: retryable. And it's weird to pass parse_write_concern_error
218+
# from outside database.py.
219+
self._client.admin.command(
220+
command_name,
221+
txnNumber=self._server_session.transaction_id,
222+
session=self,
223+
write_concern=self._transaction_options.write_concern,
224+
parse_write_concern_error=True)
225+
finally:
226+
self._server_session.reset_transaction()
227+
self._transaction_options = None
228+
140229
def _advance_cluster_time(self, cluster_time):
141230
"""Internal cluster time helper."""
142231
if self._cluster_time is None:
@@ -186,19 +275,28 @@ def has_ended(self):
186275
"""True if this session is finished."""
187276
return self._server_session is None
188277

189-
def _use_lsid(self):
278+
@property
279+
def in_transaction(self):
280+
"""True if this session has an active multi-statement transaction."""
281+
return (self._server_session is not None
282+
and self._server_session.in_transaction)
283+
284+
def _apply_to(self, command, is_retryable):
190285
# Internal function.
191286
if self._server_session is None:
192287
raise InvalidOperation("Cannot use ended session")
193288

194-
return self._server_session.use_lsid()
289+
if self.options.auto_start_transaction and not self.in_transaction:
290+
self.start_transaction()
291+
292+
self._server_session.apply_to(command, is_retryable)
195293

196-
def _transaction_id(self):
294+
def _advance_statement_id(self, n):
197295
# Internal function.
198296
if self._server_session is None:
199297
raise InvalidOperation("Cannot use ended session")
200298

201-
return self._server_session.transaction_id()
299+
self._server_session.advance_statement_id(n)
202300

203301
def _retry_transaction_id(self):
204302
# Internal function.
@@ -213,23 +311,53 @@ def __init__(self):
213311
# Ensure id is type 4, regardless of CodecOptions.uuid_representation.
214312
self.session_id = {'id': Binary(uuid.uuid4().bytes, 4)}
215313
self.last_use = monotonic.time()
314+
self.in_transaction = False
216315
self._transaction_id = 0
316+
self.statement_id = 0
217317

218318
def timed_out(self, session_timeout_minutes):
219319
idle_seconds = monotonic.time() - self.last_use
220320

221321
# Timed out if we have less than a minute to live.
222322
return idle_seconds > (session_timeout_minutes - 1) * 60
223323

224-
def use_lsid(self):
324+
def apply_to(self, command, is_retryable):
325+
command['lsid'] = self.session_id
326+
327+
if is_retryable:
328+
self._transaction_id += 1
329+
command['txnNumber'] = self.transaction_id
330+
elif self.in_transaction:
331+
command['txnNumber'] = self.transaction_id
332+
# TODO: Allow stmtId for find/getMore, SERVER-33213.
333+
name = next(iter(command))
334+
if name not in ('find', 'getMore'):
335+
command['stmtId'] = self.statement_id
336+
if self.statement_id == 0:
337+
command['readConcern'] = {'level': 'snapshot'}
338+
command['autocommit'] = False
339+
self.statement_id += 1
340+
225341
self.last_use = monotonic.time()
226-
return self.session_id
227342

343+
def advance_statement_id(self, n):
344+
# Every command advances the statement id by 1 already.
345+
self.statement_id += (n - 1)
346+
347+
@property
228348
def transaction_id(self):
229-
"""Monotonically increasing positive 64-bit integer."""
230-
self._transaction_id += 1
349+
"""Positive 64-bit integer."""
231350
return Int64(self._transaction_id)
232351

352+
def start_transaction(self):
353+
self._transaction_id += 1
354+
self.statement_id = 0
355+
self.in_transaction = True
356+
357+
def reset_transaction(self):
358+
self.in_transaction = False
359+
self.statement_id = 0
360+
233361
def retry_transaction_id(self):
234362
self._transaction_id -= 1
235363

‎pymongo/command_cursor.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ def __die(self, synchronous=False):
8787

8888
def __end_session(self, synchronous):
8989
if self.__session and not self.__explicit_session:
90-
self.__session._end_session(lock=synchronous)
90+
self.__session._end_session(lock=synchronous, abort_txn=False)
9191
self.__session = None
9292

9393
def close(self):

‎pymongo/cursor.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -316,7 +316,7 @@ def __die(self, synchronous=False):
316316
if self.__exhaust and self.__exhaust_mgr:
317317
self.__exhaust_mgr.close()
318318
if self.__session and not self.__explicit_session:
319-
self.__session._end_session(lock=synchronous)
319+
self.__session._end_session(lock=synchronous, abort_txn=False)
320320
self.__session = None
321321

322322
def close(self):

‎pymongo/message.py‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -279,10 +279,11 @@ def as_command(self, sock_info):
279279
cmd = SON([('explain', cmd)])
280280
session = self.session
281281
if session:
282-
cmd['lsid'] = session._use_lsid()
282+
session._apply_to(cmd, False)
283283
# Explain does not support readConcern.
284284
if (not explain and session.options.causal_consistency
285-
and session.operation_time is not None):
285+
and session.operation_time is not None
286+
and not session.in_transaction):
286287
cmd.setdefault(
287288
'readConcern', {})[
288289
'afterClusterTime'] = session.operation_time
@@ -353,7 +354,7 @@ def as_command(self, sock_info):
353354
self.max_await_time_ms)
354355

355356
if self.session:
356-
cmd['lsid'] = self.session._use_lsid()
357+
self.session._apply_to(cmd, False)
357358
sock_info.send_cluster_time(cmd, self.session, self.client)
358359
return cmd, self.db
359360

@@ -653,9 +654,6 @@ def legacy_write(self, request_id, msg, max_doc_size, acknowledged, docs):
653654
def write_command(self, request_id, msg, docs):
654655
"""A proxy for SocketInfo.write_command that handles event publishing.
655656
"""
656-
if self.session:
657-
# Update last_use time.
658-
self.session._use_lsid()
659657
if self.publish:
660658
duration = datetime.datetime.now() - self.start_time
661659
self._start(request_id, docs)

‎pymongo/mongo_client.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1344,7 +1344,9 @@ def _process_periodic_tasks(self):
13441344
except Exception:
13451345
helpers._handle_exception()
13461346

1347-
def start_session(self, causal_consistency=True):
1347+
def start_session(self,
1348+
causal_consistency=True,
1349+
auto_start_transaction=False):
13481350
"""Start a logical session.
13491351
13501352
This method takes the same parameters as
@@ -1374,7 +1376,8 @@ def start_session(self, causal_consistency=True):
13741376
# Raises ConfigurationError if sessions are not supported.
13751377
server_session = self._get_server_session()
13761378
opts = client_session.SessionOptions(
1377-
causal_consistency=causal_consistency)
1379+
causal_consistency=causal_consistency,
1380+
auto_start_transaction=auto_start_transaction)
13781381
return client_session.ClientSession(
13791382
self, server_session, opts, authset)
13801383

‎pymongo/network.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ def command(sock, dbname, spec, slave_ok, is_mongos,
8888
if read_concern.level:
8989
spec['readConcern'] = read_concern.document
9090
if (session and session.options.causal_consistency
91-
and session.operation_time is not None):
91+
and session.operation_time is not None
92+
and not session.in_transaction):
9293
spec.setdefault(
9394
'readConcern', {})['afterClusterTime'] = session.operation_time
9495
if collation is not None:

‎pymongo/pool.py‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -502,9 +502,7 @@ def command(self, dbname, spec, slave_ok=False,
502502
# Ensure command name remains in first place.
503503
spec = SON(spec)
504504
if session:
505-
spec['lsid'] = session._use_lsid()
506-
if retryable_write:
507-
spec['txnNumber'] = session._transaction_id()
505+
session._apply_to(spec, retryable_write)
508506
self.send_cluster_time(spec, session, client)
509507
listeners = self.listeners if publish_events else None
510508
try:

0 commit comments

Comments
 (0)