PYTHON-1874 Fix coll.aggregate() when result is missing the "ns" field

Author: ShaneHarvey

pymongo/aggregation.py

@@ -77,6 +77,11 @@ def _cursor_namespace(self):
         """The namespace in which the aggregate command is run."""
         raise NotImplementedError
 
+    @property
+    def _cursor_collection(self, cursor_doc):
+        """The Collection used for the aggregate command cursor."""
+        raise NotImplementedError
+
     @property
     def _database(self):
         """The database against which the aggregation command is run."""
@@ -152,18 +157,9 @@ def get_cursor(self, session, server, sock_info, slave_ok):
                 "ns": self._cursor_namespace,
             }
 
-        # Get collection to target with cursor.
-        ns = cursor["ns"]
-        _, collname = ns.split(".", 1)
-        aggregation_collection = self._database.get_collection(
-            collname, codec_options=self._target.codec_options,
-            read_preference=read_preference,
-            write_concern=self._target.write_concern,
-            read_concern=self._target.read_concern)
-
         # Create and return cursor instance.
         return self._cursor_class(
-            aggregation_collection, cursor, sock_info.address,
+            self._cursor_collection(cursor), cursor, sock_info.address,
             batch_size=self._batch_size or 0,
             max_await_time_ms=self._max_await_time_ms,
             session=session, explicit_session=self._explicit_session)
@@ -188,6 +184,10 @@ def _aggregation_target(self):
     def _cursor_namespace(self):
         return self._target.full_name
 
+    def _cursor_collection(self, cursor):
+        """The Collection used for the aggregate command cursor."""
+        return self._target
+
     @property
     def _database(self):
         return self._target.database
@@ -209,16 +209,24 @@ def _aggregation_target(self):
 
     @property
     def _cursor_namespace(self):
-        return "%s.%s.aggregate" % (self._target.name, "$cmd")
+        return "%s.$cmd.aggregate" % (self._target.name,)
 
     @property
     def _database(self):
         return self._target
 
+    def _cursor_collection(self, cursor):
+        """The Collection used for the aggregate command cursor."""
+        # Collection level aggregate may not always return the "ns" field
+        # according to our MockupDB tests. Let's handle that case for db level
+        # aggregate too by defaulting to the <db>.$cmd.aggregate namespace.
+        _, collname = cursor.get("ns", self._cursor_namespace).split(".", 1)
+        return self._database[collname]
+
     @staticmethod
     def _check_compat(sock_info):
         # Older server version don't raise a descriptive error, so we raise
         # one instead.
         if not sock_info.max_wire_version >= 6:
-            err_msg = "Database.aggregation is only supported on MongoDB 3.6+."
+            err_msg = "Database.aggregate() is only supported on MongoDB 3.6+."
             raise ConfigurationError(err_msg)

test/test_database.py

@@ -1054,7 +1054,7 @@ def test_database_aggregation_fake_cursor(self):
 
     @client_context.require_version_max(3, 6, 0, -1)
     def test_database_aggregation_unsupported(self):
-        err_msg = "Database.aggregation is only supported on MongoDB 3.6\+."
+        err_msg = "Database.aggregate\(\) is only supported on MongoDB 3.6\+."
         with self.assertRaisesRegex(ConfigurationError, err_msg):
             with self.admin.aggregate(self.pipeline) as _:
                 pass