MOTOR-1136 Ability to create and manage Atlas search indexes in the database (#220)

Author: blink1073

motor/core.py

@@ -722,6 +722,92 @@ def watch(
             show_expanded_events,
         )
 
+    async def _cursor_command(
+        self,
+        command,
+        value=1,
+        read_preference=None,
+        codec_options=None,
+        session=None,
+        comment=None,
+        max_await_time_ms=None,
+        **kwargs,
+    ):
+        """Issue a MongoDB command and parse the response as a cursor.
+
+        If the response from the server does not include a cursor field, an error will be thrown.
+
+        Otherwise, behaves identically to issuing a normal MongoDB command.
+
+        :Parameters:
+          - `command`: document representing the command to be issued,
+            or the name of the command (for simple commands only).
+
+            .. note:: the order of keys in the `command` document is
+               significant (the "verb" must come first), so commands
+               which require multiple keys (e.g. `findandmodify`)
+               should use an instance of :class:`~bson.son.SON` or
+               a string and kwargs instead of a Python `dict`.
+
+          - `value` (optional): value to use for the command verb when
+            `command` is passed as a string
+          - `read_preference` (optional): The read preference for this
+            operation. See :mod:`~pymongo.read_preferences` for options.
+            If the provided `session` is in a transaction, defaults to the
+            read preference configured for the transaction.
+            Otherwise, defaults to
+            :attr:`~pymongo.read_preferences.ReadPreference.PRIMARY`.
+          - `codec_options`: A :class:`~bson.codec_options.CodecOptions`
+            instance.
+          - `session` (optional): A
+            :class:`MotorClientSession`.
+          - `comment` (optional): A user-provided comment to attach to future getMores for this
+            command.
+          - `max_await_time_ms` (optional): The number of ms to wait for more data on future getMores for this command.
+          - `**kwargs` (optional): additional keyword arguments will
+            be added to the command document before it is sent
+
+        .. note:: :meth:`command` does **not** obey this Database's
+           :attr:`read_preference` or :attr:`codec_options`. You must use the
+           ``read_preference`` and ``codec_options`` parameters instead.
+
+        .. note:: :meth:`command` does **not** apply any custom TypeDecoders
+           when decoding the command response.
+
+        .. note:: If this client has been configured to use MongoDB Stable
+           API (see :ref:`versioned-api-ref`), then :meth:`command` will
+           automatically add API versioning options to the given command.
+           Explicitly adding API versioning options in the command and
+           declaring an API version on the client is not supported.
+
+        .. seealso:: The MongoDB documentation on `commands <https://dochub.mongodb.org/core/commands>`_.
+        """
+        args = (command,)
+        kwargs["value"] = value
+        kwargs["read_preference"] = read_preference
+        kwargs["codec_options"] = codec_options
+        kwargs["session"] = session
+        kwargs["comment"] = comment
+        kwargs["max_await_time_ms"] = max_await_time_ms
+
+        def inner():
+            return self.delegate.cursor_command(
+                *unwrap_args_session(args), **unwrap_kwargs_session(kwargs)
+            )
+
+        loop = self.get_io_loop()
+        cursor = await self._framework.run_on_executor(loop, inner)
+
+        cursor_class = create_class_with_framework(
+            AgnosticCommandCursor, self._framework, self.__module__
+        )
+
+        return cursor_class(cursor, self)
+
+    # TODO: MOTOR-1169
+    if hasattr(Database, "cursor_command"):
+        cursor_command = _cursor_command
+
     @property
     def client(self):
         """This MotorDatabase's :class:`MotorClient`."""
@@ -805,8 +891,17 @@ class AgnosticCollection(AgnosticBaseProperties):
     replace_one = AsyncCommand(doc=docstrings.replace_one_doc)
     update_many = AsyncCommand(doc=docstrings.update_many_doc)
     update_one = AsyncCommand(doc=docstrings.update_one_doc)
+
     with_options = DelegateMethod().wrap(Collection)
 
+    # TODO: MOTOR-1169
+    if hasattr(Collection, "create_search_index"):
+        create_search_index = AsyncCommand()
+        create_search_indexes = AsyncCommand()
+        drop_search_index = AsyncCommand()
+        update_search_index = AsyncCommand()
+        _async_list_search_indexes = AsyncRead(attr_name="list_search_indexes")
+
     _async_aggregate = AsyncRead(attr_name="aggregate")
     _async_aggregate_raw_batches = AsyncRead(attr_name="aggregate_raw_batches")
     _async_list_indexes = AsyncRead(attr_name="list_indexes")
@@ -1224,6 +1319,19 @@ async def print_indexes():
         # Latent cursor that will send initial command on first "async for".
         return cursor_class(self, self._async_list_indexes, session=session, **kwargs)
 
+    def _list_search_indexes(self, session=None, **kwargs):
+        """Return a cursor over search indexes for the current collection."""
+        cursor_class = create_class_with_framework(
+            AgnosticLatentCommandCursor, self._framework, self.__module__
+        )
+
+        # Latent cursor that will send initial command on first "async for".
+        return cursor_class(self, self._async_list_search_indexes, session=session, **kwargs)
+
+    # TODO: MOTOR-1169
+    if hasattr(Collection, "list_search_indexes"):
+        list_search_indexes = _list_search_indexes
+
     def wrap(self, obj):
         if obj.__class__ is Collection:
             # Replace pymongo.collection.Collection with MotorCollection.
@@ -1670,6 +1778,32 @@ class AgnosticCommandCursor(AgnosticBaseCursor):
 
     _CommandCursor__die = AsyncRead()
 
+    async def _try_next(self):
+        """Advance the cursor without blocking indefinitely.
+
+        This method returns the next document without waiting
+        indefinitely for data.
+
+        If no document is cached locally then this method runs a single
+        getMore command. If the getMore yields any documents, the next
+        document is returned, otherwise, if the getMore returns no documents
+        (because there is no additional data) then ``None`` is returned.
+
+        :Returns:
+          The next document or ``None`` when no document is available
+          after running a single getMore or when the cursor is closed.
+        """
+
+        def inner():
+            return self.delegate.try_next()
+
+        loop = self.get_io_loop()
+        return await self._framework.run_on_executor(loop, inner)
+
+    # TODO: MOTOR-1169
+    if hasattr(CommandCursor, "try_next"):
+        try_next = _try_next
+
     def _query_flags(self):
         return 0
 

synchro/__init__.py

@@ -96,7 +96,7 @@
 from pymongo.periodic_executor import *
 from pymongo.periodic_executor import _EXECUTORS
 from pymongo.pool import *
-from pymongo.pool import _METADATA, Pool, SocketInfo, _PoolClosedError
+from pymongo.pool import _METADATA, Connection, Pool, _PoolClosedError
 from pymongo.read_concern import *
 from pymongo.read_preferences import *
 from pymongo.read_preferences import _ServerMode
@@ -468,6 +468,13 @@ def __init__(self, client, name, **kwargs):
     def client(self):
         return self._client
 
+    def cursor_command(self, *args, **kwargs):
+        if "session" in kwargs:
+            # Workaround for validation added in PYTHON-3228.
+            kwargs["session"] = kwargs["session"].delegate
+        cursor = self.synchronize(self.delegate.cursor_command)(*args, **kwargs)
+        return CommandCursor(cursor)
+
     def __getattr__(self, name):
         return Collection(self, name, delegate=getattr(self.delegate, name))
 
@@ -483,6 +490,11 @@ class Collection(Synchro):
     aggregate = WrapOutgoing()
     aggregate_raw_batches = WrapOutgoing()
     list_indexes = WrapOutgoing()
+    create_search_index = WrapOutgoing()
+    create_search_indexes = WrapOutgoing()
+    drop_search_index = WrapOutgoing()
+    list_search_indexes = WrapOutgoing()
+    update_search_index = WrapOutgoing()
     watch = WrapOutgoing()
     __bool__ = WrapOutgoing()
 
@@ -618,6 +630,8 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 class CommandCursor(Cursor):
     __delegate_class__ = motor.motor_tornado.MotorCommandCursor
 
+    try_next = Sync("try_next")
+
 
 class GridOutCursor(Cursor):
     __delegate_class__ = motor.motor_tornado.MotorGridOutCursor

test/asyncio_tests/test_asyncio_client.py

@@ -107,9 +107,16 @@ async def test_reconnect_in_case_connection_closed_by_mongo(self):
         # lost, as result we should have AutoReconnect instead of
         # IncompleteReadError
         pool = get_primary_pool(cx)
-        socket = pool.sockets.pop()
-        socket.sock.close()
-        pool.sockets.appendleft(socket)
+
+        # TODO: MOTOR-1169
+        if hasattr(pool, "conns"):
+            conn = pool.conns.pop()
+            conn.conn.close()
+            pool.conns.appendleft(conn)
+        else:
+            socket = pool.sockets.pop()
+            socket.sock.close()
+            pool.sockets.appendleft(socket)
 
         with self.assertRaises(pymongo.errors.AutoReconnect):
             await cx.motor_test.test_collection.find_one()

test/asyncio_tests/test_asyncio_cursor.py

@@ -456,53 +456,58 @@ async def test_exhaust(self):
         # Ensure a pool.
         await client.db.collection.find_one()
 
-        socks = get_primary_pool(client).sockets
+        # TODO: MOTOR-1169
+        pool = get_primary_pool(client)
+        if hasattr(pool, "conns"):
+            conns = pool.conns
+        else:
+            conns = pool.sockets
 
         # Make sure the socket is returned after exhaustion.
         cur = client[self.db.name].test.find(cursor_type=CursorType.EXHAUST)
         has_next = await cur.fetch_next
         self.assertTrue(has_next)
-        self.assertEqual(0, len(socks))
+        self.assertEqual(0, len(conns))
 
         while await cur.fetch_next:
             cur.next_object()
 
-        self.assertEqual(1, len(socks))
+        self.assertEqual(1, len(conns))
 
         # Same as previous but with to_list instead of next_object.
         docs = await client[self.db.name].test.find(cursor_type=CursorType.EXHAUST).to_list(None)
-        self.assertEqual(1, len(socks))
+        self.assertEqual(1, len(conns))
         self.assertEqual((await self.db.test.count_documents({})), len(docs))
 
         # If the Cursor instance is discarded before being
         # completely iterated we have to close and
         # discard the socket.
-        sock = one(socks)
+        conn = one(conns)
         cur = client[self.db.name].test.find(cursor_type=CursorType.EXHAUST).batch_size(1)
         await cur.fetch_next
         self.assertTrue(cur.next_object())
         # Run at least one getMore to initiate the OP_MSG exhaust protocol.
         if env.version.at_least(4, 2):
             await cur.fetch_next
             self.assertTrue(cur.next_object())
-        self.assertEqual(0, len(socks))
+        self.assertEqual(0, len(conns))
         if "PyPy" in sys.version:
             # Don't wait for GC or use gc.collect(), it's unreliable.
             await cur.close()
 
         del cur
 
-        async def sock_closed():
-            return sock not in socks and sock.closed
+        async def conn_closed():
+            return conn not in conns and conn.closed
 
         await wait_until(
-            sock_closed, "close exhaust cursor socket", timeout=get_async_test_timeout()
+            conn_closed, "close exhaust cursor socket", timeout=get_async_test_timeout()
         )
 
         # The exhaust cursor's socket was discarded, although another may
         # already have been opened to send OP_KILLCURSORS.
-        self.assertNotIn(sock, socks)
-        self.assertTrue(sock.closed)
+        self.assertNotIn(conn, conns)
+        self.assertTrue(conn.closed)
 
     @asyncio_test
     async def test_close_with_docs_in_batch(self):

test/asyncio_tests/test_asyncio_database.py

@@ -15,7 +15,7 @@
 """Test AsyncIOMotorDatabase."""
 
 import unittest
-from test import env
+from test import SkipTest, env
 from test.asyncio_tests import AsyncIOTestCase, asyncio_test
 
 import pymongo.database
@@ -156,6 +156,21 @@ def test_with_options(self):
         self.assertEqual(db.codec_options, db2.codec_options)
         self.assertEqual(db.write_concern, db2.write_concern)
 
+    @asyncio_test
+    async def test_cursor_command(self):
+        db = self.db
+        if not hasattr(pymongo.database.Database, "cursor_command"):
+            raise SkipTest("MOTOR-1169")
+        await db.test.drop()
+
+        docs = [{"_id": i, "doc": i} for i in range(3)]
+        await db.test.insert_many(docs)
+
+        cursor = await db.cursor_command("find", "test")
+        for i in range(3):
+            item = await cursor.try_next()
+            self.assertEqual(item, docs[i])
+
 
 if __name__ == "__main__":
     unittest.main()

test/tornado_tests/test_motor_client.py

@@ -200,6 +200,13 @@ async def test_timeout(self):
         client.close()
 
 
+def get_conns(pool):
+    # TODO: MOTOR-1169
+    if hasattr(pool, "conns"):
+        return pool.conns
+    return pool.sockets
+
+
 class MotorClientExhaustCursorTest(MotorMockServerTest):
     def primary_server(self):
         primary = self.server()
@@ -221,7 +228,8 @@ async def _test_exhaust_query_server_error(self, rs):
         client = motor.MotorClient(server.uri, maxPoolSize=1)
         await client.admin.command("ismaster")
         pool = get_primary_pool(client)
-        sock_info = one(pool.sockets)
+        conns = get_conns(pool)
+        conn = one(conns)
         cursor = client.db.collection.find(cursor_type=CursorType.EXHAUST)
 
         # With Tornado, simply accessing fetch_next starts the fetch.
@@ -232,8 +240,9 @@ async def _test_exhaust_query_server_error(self, rs):
         with self.assertRaises(pymongo.errors.OperationFailure):
             await fetch_next
 
-        self.assertFalse(sock_info.closed)
-        self.assertEqual(sock_info, one(pool.sockets))
+        self.assertFalse(conn.closed)
+        conns = get_conns(pool)
+        self.assertEqual(conn, one(conns))
 
     @gen_test
     async def test_exhaust_query_server_error_standalone(self):
@@ -252,7 +261,8 @@ async def _test_exhaust_query_network_error(self, rs):
         await client.admin.command("ismaster")
         pool = get_primary_pool(client)
         pool._check_interval_seconds = None  # Never check.
-        sock_info = one(pool.sockets)
+        conns = get_conns(pool)
+        conn = one(conns)
 
         cursor = client.db.collection.find(cursor_type=CursorType.EXHAUST)
 
@@ -264,9 +274,10 @@ async def _test_exhaust_query_network_error(self, rs):
         with self.assertRaises(pymongo.errors.ConnectionFailure):
             await fetch_next
 
-        self.assertTrue(sock_info.closed)
+        self.assertTrue(conn.closed)
         del cursor
-        self.assertNotIn(sock_info, pool.sockets)
+        conns = get_conns(pool)
+        self.assertNotIn(conn, conns)
 
     @gen_test
     async def test_exhaust_query_network_error_standalone(self):