cache Query/GetMore.as_command, send stmtId with commit/abortTransaction

fix comparison of insertedIds in test_transactions.py

Author: ajdavis

pymongo/client_session.py

@@ -231,6 +231,7 @@ def _finish_transaction(self, command_name):
  self._client.admin.command(
  command_name,
  txnNumber=self._server_session.transaction_id,
+ stmtId=self._server_session.statement_id,
  session=self,
  write_concern=self._current_transaction_opts.write_concern,
  parse_write_concern_error=True)
@@ -314,11 +315,13 @@ def _apply_to(self, command, is_retryable):
  command['autocommit'] = False
 
  command['txnNumber'] = self._server_session.transaction_id
+
  # TODO: Allow stmtId for find/getMore, SERVER-33213.
  name = next(iter(command))
  if name not in ('find', 'getMore'):
  command['stmtId'] = self._server_session.statement_id
- self._server_session.statement_id += 1
+
+ self._server_session.statement_id += 1
 
  def _advance_statement_id(self, n):
  self._check_ended()

pymongo/message.py

@@ -222,7 +222,7 @@ class _Query(object):
  __slots__ = ('flags', 'db', 'coll', 'ntoskip', 'spec',
  'fields', 'codec_options', 'read_preference', 'limit',
  'batch_size', 'name', 'read_concern', 'collation',
- 'session', 'client')
+ 'session', 'client', '__as_command')
 
  def __init__(self, flags, db, coll, ntoskip, spec, fields,
  codec_options, read_preference, limit,
@@ -242,6 +242,7 @@ def __init__(self, flags, db, coll, ntoskip, spec, fields,
  self.session = session
  self.client = client
  self.name = 'find'
+ self.__as_command = None
 
  def use_command(self, sock_info, exhaust):
  use_find_cmd = False
@@ -265,10 +266,13 @@ def use_command(self, sock_info, exhaust):
  return use_find_cmd
 
  def as_command(self, sock_info):
- """Return a find command document for this query.
+ """Return a find command document for this query."""
+ # We use the command twice: on the wire and for command monitoring.
+ # Generate it once, for speed and to avoid repeating side-effects
+ # like incrementing the session's statement id.
+ if self.__as_command is not None:
+ return self.__as_command
 
- Should be called *after* get_message.
- """
  explain = '$explain' in self.spec
  cmd = _gen_find_command(
  self.coll, self.spec, self.fields, self.ntoskip,
@@ -288,7 +292,8 @@ def as_command(self, sock_info):
  'readConcern', {})[
  'afterClusterTime'] = session.operation_time
  sock_info.send_cluster_time(cmd, session, self.client)
- return cmd, self.db
+ self.__as_command = cmd, self.db
+ return self.__as_command
 
  def get_message(self, set_slave_ok, sock_info, use_cmd=False):
  """Get a query message, possibly setting the slaveOk bit."""
@@ -328,7 +333,7 @@ class _GetMore(object):
  """A getmore operation."""
 
  __slots__ = ('db', 'coll', 'ntoreturn', 'cursor_id', 'max_await_time_ms',
- 'codec_options', 'session', 'client')
+ 'codec_options', 'session', 'client', '__as_command')
 
  name = 'getMore'
 
@@ -342,21 +347,27 @@ def __init__(self, db, coll, ntoreturn, cursor_id, codec_options, session,
  self.session = session
  self.client = client
  self.max_await_time_ms = max_await_time_ms
+ self.__as_command = None
 
  def use_command(self, sock_info, exhaust):
  sock_info.validate_session(self.client, self.session)
  return sock_info.max_wire_version >= 4 and not exhaust
 
  def as_command(self, sock_info):
  """Return a getMore command document for this query."""
+ # See _Query.as_command for an explanation of this caching.
+ if self.__as_command is not None:
+ return self.__as_command
+
  cmd = _gen_get_more_command(self.cursor_id, self.coll,
  self.ntoreturn,
  self.max_await_time_ms)
 
  if self.session:
  self.session._apply_to(cmd, False)
  sock_info.send_cluster_time(cmd, self.session, self.client)
- return cmd, self.db
+ self.__as_command = cmd, self.db
+ return self.__as_command
 
  def get_message(self, dummy0, sock_info, use_cmd=False):
  """Get a getmore message."""

test/test_transactions.py

@@ -75,8 +75,8 @@ def check_result(self, expected_result, result):
  for res in expected_result:
  prop = camel_to_snake(res)
  # SPEC-869: Only BulkWriteResult has upserted_count.
- if (prop == "upserted_count" and
- not isinstance(result, BulkWriteResult)):
+ if (prop == "upserted_count"
+  and not isinstance(result, BulkWriteResult)):
  if result.upserted_id is not None:
  upserted_count = 1
  else:
@@ -85,22 +85,27 @@ def check_result(self, expected_result, result):
  elif prop == "inserted_ids":
  # BulkWriteResult does not have inserted_ids.
  if isinstance(result, BulkWriteResult):
- self.assertEqual(len(expected_result[res]), result.inserted_count)
+ self.assertEqual(len(expected_result[res]),
+ result.inserted_count)
  else:
  # InsertManyResult may be compared to [id1] from the
  # crud spec or {"0": id1} from the retryable write spec.
  ids = expected_result[res]
  if isinstance(ids, dict):
  ids = [ids[str(i)] for i in range(len(ids))]
- self.assertEqual(ids, result.inserted_ids)
+ self.assertEqual(ids,
+ result.inserted_ids)
  elif prop == "upserted_ids":
  # Convert indexes from strings to integers.
  ids = expected_result[res]
  expected_ids = {}
  for str_index in ids:
  expected_ids[int(str_index)] = ids[str_index]
- self.assertEqual(expected_ids != result.upserted_ids)
- self.assertEqual(getattr(result, prop), expected_result[res])
+ self.assertEqual(expected_ids,
+ result.upserted_ids)
+ else:
+ self.assertEqual(getattr(result, prop),
+ expected_result[res])
  return True
  else:
  self.assertEqual(result, expected_result)

test/transactions/abort.json

@@ -81,6 +81,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -113,6 +114,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -176,6 +178,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -254,6 +257,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -265,6 +269,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -362,6 +367,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",

test/transactions/auto-start.json

@@ -85,6 +85,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -137,6 +138,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 2,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -223,6 +225,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -313,6 +316,7 @@
  "abortTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "abortTransaction",
@@ -345,6 +349,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",

test/transactions/commit.json

@@ -80,6 +80,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -112,6 +113,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -196,6 +198,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",

test/transactions/reads.json

@@ -25,18 +25,6 @@
  "arguments": {
  "session": "session0"
  }
- },
- {
- "command_started_event": {
- "command": {
- "commitTransaction": 1,
- "lsid": "session0",
- "txnNumber": 1,
- "writeConcern": null
- },
- "command_name": "commitTransaction",
- "database_name": "admin"
- }
  }
  ],
  "tests": [
@@ -122,6 +110,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 2,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -270,6 +259,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 4,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -358,6 +348,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 2,
  "writeConcern": null
  },
  "command_name": "commitTransaction",

test/transactions/snapshot-reads.json

@@ -219,6 +219,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 3,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -313,6 +314,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 3,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -799,6 +801,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 3,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -893,6 +896,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 3,
  "writeConcern": null
  },
  "command_name": "commitTransaction",

test/transactions/statement-ids.json

@@ -36,10 +36,10 @@
  "session": "session0"
  },
  "result": {
- "insertedIds": [
- 2,
- 3
- ]
+ "insertedIds": {
+ "0": 2,
+ "1": 3
+ }
  }
  },
  {
@@ -153,6 +153,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 4,
  "writeConcern": null
  },
  "command_name": "commitTransaction",
@@ -185,6 +186,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 2,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",

test/transactions/write-concern.json

@@ -58,6 +58,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": {
  "w": "majority"
  }
@@ -131,6 +132,7 @@
  "commitTransaction": 1,
  "lsid": "session0",
  "txnNumber": 1,
+ "stmtId": 1,
  "writeConcern": null
  },
  "command_name": "commitTransaction",