PYTHON-952 - Publish explain commands when $explain is used

Author: behackett

pymongo/cursor.py

@@ -804,7 +804,6 @@ def __send_message(self, operation):
  publish = monitoring.enabled()
 
  if operation:
- cmd_name = operation.name
  kwargs = {
  "read_preference": self.__read_preference,
  "exhaust": self.__exhaust,
@@ -821,6 +820,7 @@ def __send_message(self, operation):
  self.__exhaust_mgr = _SocketManager(response.socket_info,
  response.pool)
 
+ cmd_name = operation.name
  data = response.data
  cmd_duration = response.duration
  rqst_id = response.request_id
@@ -896,14 +896,17 @@ def __send_message(self, operation):
 
  if publish:
  duration = (datetime.datetime.now() - start) + cmd_duration
- # Must publish in find / getMore command response format.
- res = {"cursor": {"id": doc["cursor_id"],
- "ns": self.__collection.full_name},
- "ok": 1}
- if cmd_name == "find":
- res["cursor"]["firstBatch"] = doc["data"]
+ # Must publish in find / getMore / explain command response format.
+ if cmd_name == "explain":
+ res = doc["data"][0] if doc["number_returned"] else {}
  else:
- res["cursor"]["nextBatch"] = doc["data"]
+ res = {"cursor": {"id": doc["cursor_id"],
+ "ns": self.__collection.full_name},
+ "ok": 1}
+ if cmd_name == "find":
+ res["cursor"]["firstBatch"] = doc["data"]
+ else:
+ res["cursor"]["nextBatch"] = doc["data"]
  monitoring.publish_command_success(
  duration, res, cmd_name, rqst_id, self.__address)
 

pymongo/message.py

@@ -75,7 +75,6 @@ def _maybe_add_read_preference(spec, read_preference):
  return spec
 
 
-# XXX: What to do about exhaust?
 _OPTIONS = SON([
  ('tailable', 2),
  ('oplogReplay', 8),
@@ -84,7 +83,6 @@ def _maybe_add_read_preference(spec, read_preference):
  ('allowPartialResults', 128)])
 
 
-# XXX: What about $explain? Time to switch to explain command?
 _MODIFIERS = SON([
  ('$query', 'filter'),
  ('$orderby', 'sort'),
@@ -101,11 +99,20 @@ def _maybe_add_read_preference(spec, read_preference):
  ('$snapshot', 'snapshot')])
 
 
+def _gen_explain_command(
+ coll, spec, projection, skip, limit, batch_size, options):
+ """Generate an explain command document."""
+ cmd = _gen_find_command(
+ coll, spec, projection, skip, limit, batch_size, options)
+ return SON([('explain', cmd)])
+
+
 def _gen_find_command(coll, spec, projection, skip, limit, batch_size, options):
  """Generate a find command document."""
  cmd = SON([('find', coll)])
  if '$query' in spec:
- cmd.update([(_MODIFIERS[key], val) for key, val in spec.items()])
+ cmd.update([(_MODIFIERS[key], val)
+ for key, val in spec.items() if key in _MODIFIERS])
  else:
  cmd['filter'] = spec
 
@@ -143,9 +150,8 @@ class _Query(object):
  """A query operation."""
 
  __slots__ = ('flags', 'ns', 'ntoskip', 'ntoreturn', 'spec', 'fields',
- 'codec_options', 'read_preference', 'limit', 'batch_size')
-
- name = 'find'
+ 'codec_options', 'read_preference', 'limit', 'batch_size',
+ 'name')
 
  def __init__(self, flags, ns, ntoskip, ntoreturn, spec, fields,
  codec_options, read_preference, limit, batch_size):
@@ -159,13 +165,19 @@ def __init__(self, flags, ns, ntoskip, ntoreturn, spec, fields,
  self.read_preference = read_preference
  self.limit = limit
  self.batch_size = batch_size
+ self.name = 'find'
 
  def as_command(self):
  """Return a find command document for this query.
 
  Should be called *after* get_message.
  """
  dbn, coll = self.ns.split('.', 1)
+ if '$explain' in self.spec:
+ self.name = 'explain'
+ return _gen_explain_command(
+ coll, self.spec, self.fields, self.ntoskip,
+ self.limit, self.batch_size, self.flags), dbn
  return _gen_find_command(coll, self.spec, self.fields, self.ntoskip,
  self.limit, self.batch_size, self.flags), dbn
 

test/test_monitoring.py

@@ -198,6 +198,33 @@ def test_find_and_get_more(self):
  'ok': 1}
  self.assertEqual(expected_result, succeeded.reply)
 
+ def test_find_with_explain(self):
+ self.client.pymongo_test.test.drop()
+ self.client.pymongo_test.test.insert_one({})
+ self.listener.results = {}
+ res = self.client.pymongo_test.test.find().explain()
+ results = self.listener.results
+ started = results.get('started')
+ succeeded = results.get('succeeded')
+ self.assertIsNone(results.get('failed'))
+ self.assertTrue(
+ isinstance(started, monitoring.CommandStartedEvent))
+ self.assertEqual(
+ SON([('explain', SON([('find', 'test'),
+ ('filter', {})]))]),
+ started.command)
+ self.assertEqual('explain', started.command_name)
+ self.assertEqual(self.client.address, started.connection_id)
+ self.assertEqual('pymongo_test', started.database_name)
+ self.assertTrue(isinstance(started.request_id, int))
+ self.assertTrue(
+ isinstance(succeeded, monitoring.CommandSucceededEvent))
+ self.assertTrue(isinstance(succeeded.duration_micros, int))
+ self.assertEqual('explain', succeeded.command_name)
+ self.assertTrue(isinstance(succeeded.request_id, int))
+ self.assertEqual(self.client.address, succeeded.connection_id)
+ self.assertEqual(res, succeeded.reply)
+
  @client_context.require_version_min(2, 6, 0)
  def test_command_and_get_more(self):
  self.client.pymongo_test.test.drop()