PYTHON-1355 - list_collections improvements

behackett · behackett · commit f107c087543a · 2017-10-26T16:50:19.000-07:00
diff --git a/doc/changelog.rst b/doc/changelog.rst
@@ -19,6 +19,8 @@ Highlights include:
   external libraries that can parse raw batches of BSON data.
 - New methods :meth:`~pymongo.mongo_client.MongoClient.list_databases` and
   :meth:`~pymongo.mongo_client.MongoClient.list_database_names`.
+- New methods :meth:`~pymongo.database.Database.list_collections` and
+  :meth:`~pymongo.database.Database.list_collection_names`.
 - Support for mongodb+srv:// URIs. See
   :class:`~pymongo.mongo_client.MongoClient` for details.
 
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -1968,9 +1968,8 @@ def options(self, session=None):
         .. versionchanged:: 3.6
            Added ``session`` parameter.
         """
-        criteria = {"name": self.__name}
-        cursor = self.__database.list_collections(filter=criteria,
-                                                      session=session)
+        cursor = self.__database.list_collections(
+            session=session, filter={"name": self.__name})
 
         result = None
         for doc in cursor:
diff --git a/pymongo/database.py b/pymongo/database.py
@@ -17,9 +17,8 @@
 import warnings
 
 from bson.code import Code
-from bson.codec_options import CodecOptions, DEFAULT_CODEC_OPTIONS
+from bson.codec_options import DEFAULT_CODEC_OPTIONS
 from bson.dbref import DBRef
-from bson.objectid import ObjectId
 from bson.py3compat import iteritems, string_type, _unicode
 from bson.son import SON
 from pymongo import auth, common
@@ -35,6 +34,10 @@
 from pymongo.write_concern import WriteConcern
 
 
+_INDEX_REGEX = {"name": {"$regex": "^(?!.*\$)"}}
+_SYSTEM_FILTER = {"filter": {"name": {"$regex": "^(?!system\.)"}}}
+
+
 def _check_name(name):
     """Check if a database name is valid.
     """
@@ -543,56 +546,73 @@ def command(self, command, value=1, check=True,
                                  check, allowable_errors, read_preference,
                                  codec_options, session=session, **kwargs)
 
-    def _list_collections(self, sock_info, slave_okay, filter=None,
-                         session=None):
+    def _list_collections(self, sock_info, slave_okay, session=None, **kwargs):
         """Internal listCollections helper."""
-        filter = filter or {}
-        cmd = SON([("listCollections", 1), ("cursor", {})])
-
-        if filter:
-            cmd["filter"] = filter
 
+        coll = self["$cmd"]
         if sock_info.max_wire_version > 2:
-            coll = self["$cmd"]
-            with self.__client._tmp_session(session, close=False) as s:
+            cmd = SON([("listCollections", 1),
+                       ("cursor", {})])
+            cmd.update(kwargs)
+            with self.__client._tmp_session(
+                    session, close=False) as tmp_session:
                 cursor = self._command(
-                    sock_info, cmd, slave_okay, session=s)["cursor"]
-                return CommandCursor(coll, cursor, sock_info.address, session=s,
-                                     explicit_session=session is not None)
+                    sock_info, cmd, slave_okay, session=tmp_session)["cursor"]
+                return CommandCursor(
+                    coll,
+                    cursor,
+                    sock_info.address,
+                    session=tmp_session,
+                    explicit_session=session is not None)
         else:
-            coll = self["system.namespaces"]
-            if "name" in filter:
-                if not isinstance(filter["name"], string_type):
-                    raise TypeError("filter['name'] must be a string on MongoDB 2.6")
-                filter["name"] = coll.database.name + "." + filter["name"]
-            res = _first_batch(sock_info, coll.database.name, coll.name,
-                               filter, 0, slave_okay,
-                               CodecOptions(), ReadPreference.PRIMARY, cmd,
-                               self.client._event_listeners, session=None)
-            data = res["data"]
-            cursor = {
-                "id": res["cursor_id"],
-                "firstBatch": data,
-                "ns": coll.full_name,
-            }
+            match = _INDEX_REGEX
+            if "filter" in kwargs:
+                match = {"$and": [_INDEX_REGEX, kwargs["filter"]]}
+            dblen = len(self.name.encode("utf8") + b".")
+            pipeline = [
+                {"$project": {"name": {"$substr": ["$name", dblen, -1]},
+                              "options": 1}},
+                {"$match": match}
+            ]
+            cmd = SON([("aggregate", "system.namespaces"),
+                       ("pipeline", pipeline),
+                       ("cursor", kwargs.get("cursor", {}))])
+            cursor = self._command(sock_info, cmd, slave_okay)["cursor"]
             return CommandCursor(coll, cursor, sock_info.address)
 
-    def list_collections(self, filter=None, session=None):
-        """Get info about the collections in this database."""
+    def list_collections(self, session=None, **kwargs):
+        """Get a cursor over the collectons of this database.
+
+        :Parameters:
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
+          - `**kwargs` (optional): Optional parameters of the
+            `listCollections command
+            <https://docs.mongodb.com/manual/reference/command/listCollections/>`_
+            can be passed as keyword arguments to this method. The supported
+            options differ by server version.
+
+        :Returns:
+          An instance of :class:`~pymongo.command_cursor.CommandCursor`.
+
+        .. versionadded:: 3.6
+        """
         with self.__client._socket_for_reads(
                 ReadPreference.PRIMARY) as (sock_info, slave_okay):
+            return self._list_collections(
+                sock_info, slave_okay, session=session, **kwargs)
+
+    def list_collection_names(self, session=None):
+        """Get a list of all the collection names in this database.
 
-             wire_version = sock_info.max_wire_version
-             results = self._list_collections(sock_info, slave_okay, filter=filter,
-                                              session=session)
-        for result in results:
-            if wire_version <= 2:
-                name = result["name"]
-                if "$" in name:
-                    continue
-                result["name"] = name.split(".", 1)[1]
-            yield result
+        :Parameters:
+          - `session` (optional): a
+            :class:`~pymongo.client_session.ClientSession`.
 
+        .. versionadded:: 3.6
+        """
+        return [result["name"]
+                for result in self.list_collections(session=session)]
 
     def collection_names(self, include_system_collections=True,
                          session=None):
@@ -607,17 +627,9 @@ def collection_names(self, include_system_collections=True,
         .. versionchanged:: 3.6
            Added ``session`` parameter.
         """
-
-        results = self.list_collections(session=session)
-
-        # Iterating the cursor to completion may require a socket for getmore.
-        # Ensure we do that outside the "with" block so we don't require more
-        # than one socket at a time.
-        names = [result["name"] for result in results]
-
-        if not include_system_collections:
-            names = [name for name in names if not name.startswith("system.")]
-        return names
+        kws = {} if include_system_collections else _SYSTEM_FILTER
+        return [result["name"]
+                for result in self.list_collections(session=session, **kws)]
 
     def drop_collection(self, name_or_collection, session=None):
         """Drop a collection.
diff --git a/test/test_database.py b/test/test_database.py
@@ -155,31 +155,41 @@ def test_create_collection(self):
         self.assertTrue(u"test.foo" in db.collection_names())
         self.assertRaises(CollectionInvalid, db.create_collection, "test.foo")
 
-    def test_collection_names(self):
+    def _test_collection_names(self, meth, test_no_system):
         db = Database(self.client, "pymongo_test")
         db.test.insert_one({"dummy": u"object"})
         db.test.mike.insert_one({"dummy": u"object"})
 
-        colls = db.collection_names()
+        colls = getattr(db, meth)()
         self.assertTrue("test" in colls)
         self.assertTrue("test.mike" in colls)
         for coll in colls:
             self.assertTrue("$" not in coll)
 
-        colls_without_systems = db.collection_names(False)
-        for coll in colls_without_systems:
-            self.assertTrue(not coll.startswith("system."))
+        if test_no_system:
+            db.systemcoll.test.insert_one({})
+            no_system_collections = getattr(
+                db, meth)(include_system_collections=False)
+            for coll in no_system_collections:
+                self.assertTrue(not coll.startswith("system."))
+            self.assertIn("systemcoll.test", no_system_collections)
 
         # Force more than one batch.
         db = self.client.many_collections
         for i in range(101):
             db["coll" + str(i)].insert_one({})
         # No Error
         try:
-            db.collection_names()
+            getattr(db, meth)()
         finally:
             self.client.drop_database("many_collections")
 
+    def test_collection_names(self):
+        self._test_collection_names('collection_names', True)
+
+    def test_list_collection_names(self):
+        self._test_collection_names('list_collection_names', False)
+
     def test_list_collections(self):
         self.client.drop_database("pymongo_test")
         db = Database(self.client, "pymongo_test")
@@ -215,6 +225,11 @@ def test_list_collections(self):
         else:
             self.assertTrue(False)
 
+        colls = db.list_collections(filter={"name": {"$regex": "^test$"}})
+        self.assertEqual(1, len(list(colls)))
+
+        colls = db.list_collections(filter={"name": {"$regex": "^test.mike$"}})
+        self.assertEqual(1, len(list(colls)))
 
         db.drop_collection("test")
 
diff --git a/test/test_monitoring.py b/test/test_monitoring.py
@@ -1217,27 +1217,9 @@ def test_write_errors(self):
     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.assertEqualCommand(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())
+        results = self.listener.results
         started = results['started'][0]
         succeeded = results['succeeded'][0]
         self.assertEqual(0, len(results['failed']))
