PYTHON-861 - Implement list_indexes.

behackett · behackett · commit 14ce3af058c2 · 2015-03-16T08:00:58.000-07:00
diff --git a/doc/api/pymongo/collection.rst b/doc/api/pymongo/collection.rst
@@ -56,6 +56,7 @@
       .. automethod:: drop_index
       .. automethod:: drop_indexes
       .. automethod:: reindex
+      .. automethod:: list_indexes
       .. automethod:: index_information
       .. automethod:: drop
       .. automethod:: rename
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -268,6 +268,7 @@ The following methods have been added:
 - :meth:`~pymongo.collection.Collection.find_one_and_update`
 - :meth:`~pymongo.collection.Collection.with_options`
 - :meth:`~pymongo.collection.Collection.create_indexes`
+- :meth:`~pymongo.collection.Collection.list_indexes`
 
 The following methods have changed:
 
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -1193,6 +1193,27 @@ def reindex(self):
         cmd = SON([("reIndex", self.__name)])
         return self._command(cmd, ReadPreference.PRIMARY)[0]
 
+    def list_indexes(self):
+        """Get a cursor over the indexes for this collection."""
+        client = self.__database.client
+        if client._writable_max_wire_version() > 2:
+            cmd = SON([("listIndexes", self.__name), ("cursor", {})])
+            res, addr = self._command(cmd,
+                                      ReadPreference.PRIMARY,
+                                      CodecOptions(SON))
+            cursor = res["cursor"]
+        else:
+            namespace = "%s.%s" % (self.__database.name, "system.indexes")
+            res, addr = helpers._first_batch(
+                client, namespace, {"ns": self.__full_name},
+                0, ReadPreference.PRIMARY, CodecOptions(SON))
+            cursor = {
+                "id": res["cursor_id"],
+                "firstBatch": res["data"],
+                "ns": namespace,
+            }
+        return CommandCursor(self, cursor, addr)
+
     def index_information(self):
         """Get information on this collection's indexes.
 
@@ -1212,26 +1233,9 @@ def index_information(self):
         {u'_id_': {u'key': [(u'_id', 1)]},
          u'x_1': {u'unique': True, u'key': [(u'x', 1)]}}
         """
-        client = self.__database.client
-        if client._writable_max_wire_version() > 2:
-            cmd = SON([("listIndexes", self.__name), ("cursor", {})])
-            res, addr = self._command(cmd,
-                                      ReadPreference.PRIMARY,
-                                      CodecOptions(document_class=SON))
-            # MongoDB 2.8rc2
-            if "indexes" in res:
-                raw = res["indexes"]
-            # >= MongoDB 2.8rc3
-            else:
-                raw = CommandCursor(self, res["cursor"], addr)
-        else:
-            raw = self.__database.get_collection(
-                "system.indexes",
-                codec_options=CodecOptions(document_class=SON),
-                read_preference=ReadPreference.PRIMARY).find(
-                    {"ns": self.__full_name}, {"ns": 0})
+        cursor = self.list_indexes()
         info = {}
-        for index in raw:
+        for index in cursor:
             index["key"] = index["key"].items()
             index = dict(index)
             info[index.pop("name")] = index
diff --git a/pymongo/helpers.py b/pymongo/helpers.py
@@ -184,29 +184,38 @@ def _check_command_response(response, msg=None, allowable_errors=None):
             raise OperationFailure(msg % errmsg, code, response)
 
 
-def _command(client, namespace, command, read_preference,
-             codec_options, check=True, allowable_errors=None):
-    """Internal command helper."""
-
+def _first_batch(client, namespace, query,
+                 limit, read_preference, codec_options):
+    """Simple query helper for retrieving a first (and possibly only) batch."""
     query_opts = 0
     # XXX: Set slaveOkay flag when read preference mode is anything other
     # than primary (0). Make this more clear when we finish refactoring
     # read preferences.
     if read_preference.mode:
         query_opts = 4
 
-    query = _Query(query_opts, namespace, 0, -1, command, None, codec_options)
+    query = _Query(
+        query_opts, namespace, 0, limit, query, None, codec_options)
     response = client._send_message_with_response(query, read_preference)
-    result = _unpack_response(response.data, None, codec_options)['data'][0]
+    return _unpack_response(
+        response.data, None, codec_options), response.address
+
+
+def _command(client, namespace, command, read_preference,
+             codec_options, check=True, allowable_errors=None):
+    """Internal command helper."""
+    result, addr = _first_batch(
+        client, namespace, command, -1, read_preference, codec_options)
+    result = result["data"][0]
     if check:
         msg = "command %s on namespace %s failed: %%s" % (
             repr(command).replace("%", "%%"), namespace)
         try:
             _check_command_response(result, msg, allowable_errors)
         except NotMasterError:
-            client._reset_server_and_request_check(response.address)
+            client._reset_server_and_request_check(addr)
             raise
-    return result, response.address
+    return result, addr
 
 
 def _check_write_command_response(results):
diff --git a/test/test_collection.py b/test/test_collection.py
@@ -260,6 +260,33 @@ def check_result(result):
         else:
             check_result(reindexed)
 
+    def test_list_indexes(self):
+        db = self.db
+        db.test.drop()
+        db.test.insert_one({})  # create collection
+
+        def map_indexes(indexes):
+            return dict([(index["name"], index) for index in indexes])
+
+        indexes = list(db.test.list_indexes())
+        self.assertEqual(len(indexes), 1)
+        self.assertTrue("_id_" in map_indexes(indexes))
+
+        db.test.create_index("hello")
+        indexes = list(db.test.list_indexes())
+        self.assertEqual(len(indexes), 2)
+        self.assertEqual(map_indexes(indexes)["hello_1"]["key"],
+                         SON([("hello", ASCENDING)]))
+
+        db.test.create_index([("hello", DESCENDING), ("world", ASCENDING)],
+                             unique=True)
+        indexes = list(db.test.list_indexes())
+        self.assertEqual(len(indexes), 3)
+        index_map = map_indexes(indexes)
+        self.assertEqual(index_map["hello_-1_world_1"]["key"],
+                         SON([("hello", DESCENDING), ("world", ASCENDING)]))
+        self.assertEqual(True, index_map["hello_-1_world_1"]["unique"])
+
     def test_index_info(self):
         db = self.db
         db.test.drop()
