PYTHON-1580 - Implement new count API

This commit is contained in:
Bernie Hackett 2018-06-08 14:27:04 -07:00
parent bb8130abd8
commit 91c3793703
8 changed files with 230 additions and 13 deletions

View File

@ -53,7 +53,8 @@
.. automethod:: find_one_and_delete
.. automethod:: find_one_and_replace(filter, replacement, projection=None, sort=None, return_document=ReturnDocument.BEFORE, session=None, **kwargs)
.. automethod:: find_one_and_update(filter, update, projection=None, sort=None, return_document=ReturnDocument.BEFORE, array_filters=None, session=None, **kwargs)
.. automethod:: count
.. automethod:: count_documents
.. automethod:: estimated_document_count
.. automethod:: distinct
.. automethod:: create_index
.. automethod:: create_indexes
@ -71,6 +72,7 @@
.. automethod:: initialize_unordered_bulk_op
.. automethod:: initialize_ordered_bulk_op
.. automethod:: group
.. automethod:: count
.. automethod:: insert(doc_or_docs, manipulate=True, check_keys=True, continue_on_error=False, **kwargs)
.. automethod:: save(to_save, manipulate=True, check_keys=True, **kwargs)
.. automethod:: update(spec, document, upsert=False, manipulate=False, multi=False, check_keys=True, **kwargs)

View File

@ -1533,7 +1533,9 @@ class Collection(common.BaseObject):
"""Internal count helper."""
with self._socket_for_reads(session) as (sock_info, slave_ok):
res = self._command(
sock_info, cmd, slave_ok,
sock_info,
cmd,
slave_ok,
allowable_errors=["ns missing"],
codec_options=self.__write_response_codec_options,
read_concern=self.read_concern,
@ -1543,23 +1545,117 @@ class Collection(common.BaseObject):
return 0
return int(res["n"])
def _aggregate_one_result(
self, sock_info, slave_ok, cmd, collation=None, session=None):
"""Internal helper to run an aggregate that returns a single result."""
result = self._command(
sock_info,
cmd,
slave_ok,
codec_options=self.__write_response_codec_options,
read_concern=self.read_concern,
collation=collation,
session=session)
batch = result['cursor']['firstBatch']
return batch[0] if batch else None
def estimated_document_count(self, **kwargs):
"""Get an estimate of the number of documents in this collection using
collection metadata.
The :meth:`estimated_document_count` method is **not** supported in a
transaction.
All optional parameters should be passed as keyword arguments
to this method. Valid options include:
- `maxTimeMS` (int): The maximum amount of time to allow this
operation to run, in milliseconds.
:Parameters:
- `**kwargs` (optional): See list of options above.
.. versionadded:: 3.7
"""
cmd = SON([('count', self.__name)])
cmd.update(kwargs)
return self._count(cmd)
def count_documents(self, filter, session=None, **kwargs):
"""Count the number of documents in this collection.
The :meth:`count_documents` method is supported in a transaction.
All optional parameters should be passed as keyword arguments
to this method. Valid options include:
- `skip` (int): The number of matching documents to skip before
returning results.
- `limit` (int): The maximum number of documents to count.
- `maxTimeMS` (int): The maximum amount of time to allow this
operation to run, in milliseconds.
- `collation` (optional): An instance of
:class:`~pymongo.collation.Collation`. This option is only supported
on MongoDB 3.4 and above.
- `hint` (string or list of tuples): The index to use. Specify either
the index name as a string or the index specification as a list of
tuples (e.g. [('a', pymongo.ASCENDING), ('b', pymongo.ASCENDING)]).
This option is only supported on MongoDB 3.6 and above.
The :meth:`count_documents` method obeys the :attr:`read_preference` of
this :class:`Collection`.
:Parameters:
- `filter` (required): A query document that selects which documents
to count in the collection. Can be an empty document to count all
documents.
- `session` (optional): a
:class:`~pymongo.client_session.ClientSession`.
- `**kwargs` (optional): See list of options above.
.. versionadded:: 3.7
"""
pipeline = [{'$match': filter}]
if 'skip' in kwargs:
pipeline.append({'$skip': kwargs.pop('skip')})
if 'limit' in kwargs:
pipeline.append({'$limit': kwargs.pop('limit')})
pipeline.append({'$group': {'_id': None, 'n': {'$sum': 1}}})
cmd = SON([('aggregate', self.__name),
('pipeline', pipeline),
('cursor', {})])
if "hint" in kwargs and not isinstance(kwargs["hint"], string_type):
kwargs["hint"] = helpers._index_document(kwargs["hint"])
collation = validate_collation_or_none(kwargs.pop('collation', None))
cmd.update(kwargs)
with self._socket_for_reads(session) as (sock_info, slave_ok):
result = self._aggregate_one_result(
sock_info, slave_ok, cmd, collation, session)
if not result:
return 0
return result['n']
def count(self, filter=None, session=None, **kwargs):
"""Get the number of documents in this collection.
"""**DEPRECATED** - Get the number of documents in this collection.
The :meth:`count` method is deprecated and **not** supported in a
transaction. Please use :meth:`count_documents` or
:meth:`estimated_document_count` instead.
All optional count parameters should be passed as keyword arguments
to this method. Valid options include:
- `hint` (string or list of tuples): The index to use. Specify either
the index name as a string or the index specification as a list of
tuples (e.g. [('a', pymongo.ASCENDING), ('b', pymongo.ASCENDING)]).
- `limit` (int): The maximum number of documents to count.
- `skip` (int): The number of matching documents to skip before
returning results.
- `limit` (int): The maximum number of documents to count.
- `maxTimeMS` (int): The maximum amount of time to allow the count
command to run, in milliseconds.
- `collation` (optional): An instance of
:class:`~pymongo.collation.Collation`. This option is only supported
on MongoDB 3.4 and above.
- `hint` (string or list of tuples): The index to use. Specify either
the index name as a string or the index specification as a list of
tuples (e.g. [('a', pymongo.ASCENDING), ('b', pymongo.ASCENDING)]).
The :meth:`count` method obeys the :attr:`read_preference` of
this :class:`Collection`.
@ -1571,6 +1667,9 @@ class Collection(common.BaseObject):
:class:`~pymongo.client_session.ClientSession`.
- `**kwargs` (optional): See list of options above.
.. versionchanged:: 3.7
Deprecated.
.. versionchanged:: 3.6
Added ``session`` parameter.

View File

@ -710,7 +710,11 @@ class Cursor(object):
return self
def count(self, with_limit_and_skip=False):
"""Get the size of the results set for this query.
"""**DEPRECATED** - Get the size of the results set for this query.
The :meth:`count` method is deprecated and **not** supported in a
transaction. Please use
:meth:`~pymongo.collection.Collection.count_documents` instead.
Returns the number of documents in the results set for this query. Does
not take :meth:`limit` and :meth:`skip` into account by default - set
@ -736,6 +740,9 @@ class Cursor(object):
.. note:: The `with_limit_and_skip` parameter requires server
version **>= 1.1.4-**
.. versionchanged:: 3.7
Deprecated.
.. versionchanged:: 2.8
The :meth:`~count` method now supports :meth:`~hint`.
"""

View File

@ -8,7 +8,25 @@
"minServerVersion": "3.4",
"tests": [
{
"description": "Count with collation",
"description": "Count documents with collation",
"operation": {
"name": "countDocuments",
"arguments": {
"filter": {
"x": "ping"
},
"collation": {
"locale": "en_US",
"strength": 2
}
}
},
"outcome": {
"result": 1
}
},
{
"description": "Deprecated count with collation",
"operation": {
"name": "count",
"arguments": {

View File

@ -15,7 +15,59 @@
],
"tests": [
{
"description": "Count without a filter",
"description": "Estimated document count",
"operation": {
"name": "estimatedDocumentCount",
"arguments": {}
},
"outcome": {
"result": 3
}
},
{
"description": "Count documents without a filter",
"operation": {
"name": "countDocuments",
"arguments": {
"filter": {}
}
},
"outcome": {
"result": 3
}
},
{
"description": "Count documents with a filter",
"operation": {
"name": "countDocuments",
"arguments": {
"filter": {
"_id": {
"$gt": 1
}
}
}
},
"outcome": {
"result": 2
}
},
{
"description": "Count documents with skip and limit",
"operation": {
"name": "countDocuments",
"arguments": {
"filter": {},
"skip": 1,
"limit": 3
}
},
"outcome": {
"result": 2
}
},
{
"description": "Deprecated count without a filter",
"operation": {
"name": "count",
"arguments": {
@ -27,7 +79,7 @@
}
},
{
"description": "Count with a filter",
"description": "Deprecated count with a filter",
"operation": {
"name": "count",
"arguments": {
@ -43,9 +95,9 @@
}
},
{
"description": "Count with skip and limit",
"description": "Deprecated count with skip and limit",
"operation": {
"name": "count",
"name": "countDocuments",
"arguments": {
"filter": {},
"skip": 1,

View File

@ -169,6 +169,11 @@ class TestCollation(unittest.TestCase):
self.db.test.find(collation=self.collation).count()
self.assertCollationInLastCommand()
@raisesConfigurationErrorForOldMongoDB
def test_count_documents(self):
self.db.test.count_documents({}, collation=self.collation)
self.assertCollationInLastCommand()
@raisesConfigurationErrorForOldMongoDB
def test_distinct(self):
self.db.test.distinct('foo', collation=self.collation)

View File

@ -1529,6 +1529,32 @@ class TestCollection(IntegrationTest):
self.assertEqual(
db.test.count({'foo': re.compile(r'ba.*')}), 2)
def test_count_documents(self):
db = self.db
db.drop_collection("test")
self.addCleanup(db.drop_collection, "test")
self.assertEqual(db.test.count_documents({}), 0)
db.wrong.insert_many([{}, {}])
self.assertEqual(db.test.count_documents({}), 0)
db.test.insert_many([{}, {}])
self.assertEqual(db.test.count_documents({}), 2)
db.test.insert_many([{'foo': 'bar'}, {'foo': 'baz'}])
self.assertEqual(db.test.count_documents({'foo': 'bar'}), 1)
self.assertEqual(
db.test.count_documents({'foo': re.compile(r'ba.*')}), 2)
def test_estimated_document_count(self):
db = self.db
db.drop_collection("test")
self.addCleanup(db.drop_collection, "test")
self.assertEqual(db.test.estimated_document_count(), 0)
db.wrong.insert_many([{}, {}])
self.assertEqual(db.test.estimated_document_count(), 0)
db.test.insert_many([{}, {}])
self.assertEqual(db.test.estimated_document_count(), 2)
def test_aggregate(self):
db = self.db
db.drop_collection("test")

View File

@ -436,6 +436,14 @@ class TestCommandAndReadPreference(TestReplicaSetClientBase):
def test_count(self):
self._test_coll_helper(True, self.c.pymongo_test.test, 'count')
def test_count_documents(self):
self._test_coll_helper(
True, self.c.pymongo_test.test, 'count_documents', {})
def test_estimated_document_count(self):
self._test_coll_helper(
True, self.c.pymongo_test.test, 'estimated_document_count')
def test_distinct(self):
self._test_coll_helper(True, self.c.pymongo_test.test, 'distinct', 'a')