PYTHON-952 - Publish events for uses of _first_batch

Author: behackett

pymongo/collection.py

@@ -1340,8 +1340,8 @@ def list_indexes(self):
  codec_options = CodecOptions(SON)
  coll = self.with_options(codec_options)
  with self._socket_for_primary_reads() as (sock_info, slave_ok):
+ cmd = SON([("listIndexes", self.__name), ("cursor", {})])
  if sock_info.max_wire_version > 2:
- cmd = SON([("listIndexes", self.__name), ("cursor", {})])
  cursor = self._command(sock_info, cmd, slave_ok,
  ReadPreference.PRIMARY,
  codec_options)["cursor"]
@@ -1350,7 +1350,7 @@ def list_indexes(self):
  namespace = _UJOIN % (self.__database.name, "system.indexes")
  res = helpers._first_batch(
  sock_info, namespace, {"ns": self.__full_name},
- 0, slave_ok, codec_options, ReadPreference.PRIMARY)
+ 0, slave_ok, codec_options, ReadPreference.PRIMARY, cmd)
  data = res["data"]
  cursor = {
  "id": res["cursor_id"],

pymongo/database.py

@@ -456,18 +456,19 @@ def command(self, command, value=1, check=True,
  def _list_collections(self, sock_info, slave_okay, criteria=None):
  """Internal listCollections helper."""
  criteria = criteria or {}
+ cmd = SON([("listCollections", 1), ("cursor", {})])
+ if criteria:
+ cmd["filter"] = criteria
+
  if sock_info.max_wire_version > 2:
- cmd = SON([("listCollections", 1), ("cursor", {})])
- if criteria:
- cmd["filter"] = criteria
  coll = self["$cmd"]
  cursor = self._command(sock_info, cmd, slave_okay)["cursor"]
  return CommandCursor(coll, cursor, sock_info.address)
  else:
  coll = self["system.namespaces"]
  res = _first_batch(sock_info, coll.full_name,
  criteria, 0, slave_okay,
- CodecOptions(), ReadPreference.PRIMARY)
+ CodecOptions(), ReadPreference.PRIMARY, cmd)
  data = res["data"]
  cursor = {
  "id": res["cursor_id"],

pymongo/helpers.py

@@ -15,14 +15,14 @@
 """Bits and pieces used by the driver that don't really fit elsewhere."""
 
 import collections
+import datetime
 import struct
-from pymongo.message import _Query
 
 import bson
-import pymongo
 from bson.codec_options import CodecOptions
 from bson.py3compat import itervalues, string_type, iteritems, u
 from bson.son import SON
+from pymongo import ASCENDING, monitoring
 from pymongo.errors import (CursorNotFound,
  DuplicateKeyError,
  ExecutionTimeout,
@@ -31,6 +31,7 @@
  WriteError,
  WriteConcernError,
  WTimeoutError)
+from pymongo.message import _Query
 
 
 _UUNDER = u("_")
@@ -50,7 +51,7 @@ def _index_list(key_or_list, direction=None):
  return [(key_or_list, direction)]
  else:
  if isinstance(key_or_list, string_type):
- return [(key_or_list, pymongo.ASCENDING)]
+ return [(key_or_list, ASCENDING)]
  elif not isinstance(key_or_list, (list, tuple)):
  raise TypeError("if no direction is specified, "
  "key_or_list must be an instance of list")
@@ -233,16 +234,43 @@ def _check_gle_response(response):
 
 
 def _first_batch(sock_info, namespace, query,
- ntoreturn, slave_ok, codec_options, read_preference):
+ ntoreturn, slave_ok, codec_options, read_preference, cmd):
  """Simple query helper for retrieving a first (and possibly only) batch."""
  query = _Query(
  0, namespace, 0, ntoreturn, query, None,
  codec_options, read_preference, 0, ntoreturn)
+
+ name = next(iter(cmd))
+ duration = None
+ publish = monitoring.enabled()
+ if publish:
+ start = datetime.datetime.now()
+
  request_id, msg, max_doc_size = query.get_message(slave_ok,
  sock_info.is_mongos)
+
+ if publish:
+ encoding_duration = datetime.datetime.now() - start
+ monitoring.publish_command_start(
+ cmd, namespace.split('.', 1)[0], request_id, sock_info.address)
+ start = datetime.datetime.now()
+
  sock_info.send_message(msg, max_doc_size)
  response = sock_info.receive_message(1, request_id)
- return _unpack_response(response, None, codec_options)
+ try:
+ result = _unpack_response(response, None, codec_options)
+ except (NotMasterError, OperationFailure) as exc:
+ if publish:
+ duration = (datetime.datetime.now() - start) + encoding_duration
+ monitoring.publish_command_failure(
+ duration, exc.details, name, request_id, sock_info.address)
+ raise
+ if publish:
+ duration = (datetime.datetime.now() - start) + encoding_duration
+ monitoring.publish_command_success(
+ duration, result, name, request_id, sock_info.address)
+
+ return result
 
 
 def _check_write_command_response(results):

test/test_monitoring.py

@@ -1187,6 +1187,46 @@ def test_write_errors(self):
  for error in errors:
  self.assertEqual(fields, set(error))
 
+ def test_first_batch_helper(self):
+ # Regardless of server version and use of helpers._first_batch
+ # this test should still pass.
+ self.listener.results.clear()
+ self.client.pymongo_test.collection_names()
+ results = self.listener.results
+ started = results['started'][0]
+ succeeded = results['succeeded'][0]
+ self.assertEqual(0, len(results['failed']))
+ self.assertIsInstance(started, monitoring.CommandStartedEvent)
+ expected = SON([('listCollections', 1), ('cursor', {})])
+ self.assertEqual(expected, started.command)
+ self.assertEqual('pymongo_test', started.database_name)
+ self.assertEqual('listCollections', started.command_name)
+ self.assertIsInstance(started.request_id, int)
+ self.assertEqual(self.client.address, started.connection_id)
+ self.assertIsInstance(succeeded, monitoring.CommandSucceededEvent)
+ self.assertIsInstance(succeeded.duration_micros, int)
+ self.assertEqual(started.command_name, succeeded.command_name)
+ self.assertEqual(started.request_id, succeeded.request_id)
+ self.assertEqual(started.connection_id, succeeded.connection_id)
+
+ self.listener.results.clear()
+ tuple(self.client.pymongo_test.test.list_indexes())
+ started = results['started'][0]
+ succeeded = results['succeeded'][0]
+ self.assertEqual(0, len(results['failed']))
+ self.assertIsInstance(started, monitoring.CommandStartedEvent)
+ expected = SON([('listIndexes', 'test'), ('cursor', {})])
+ self.assertEqual(expected, started.command)
+ self.assertEqual('pymongo_test', started.database_name)
+ self.assertEqual('listIndexes', started.command_name)
+ self.assertIsInstance(started.request_id, int)
+ self.assertEqual(self.client.address, started.connection_id)
+ self.assertIsInstance(succeeded, monitoring.CommandSucceededEvent)
+ self.assertIsInstance(succeeded.duration_micros, int)
+ self.assertEqual(started.command_name, succeeded.command_name)
+ self.assertEqual(started.request_id, succeeded.request_id)
+ self.assertEqual(started.connection_id, succeeded.connection_id)
+
 
 if __name__ == "__main__":
  unittest.main()