Skip to content

Commit 5f421d7

Browse files
committed
PYTHON-2547 Change estimated_document_count() to use $collStats instead of count on 4.9+ (#606)
Fix CRUD v1 aggregate $out change for mongodb/specifications@3f3a3c2 PYTHON-2301 ValueError is an acceptable error for CRUD v2 error:true tests (cherry picked from commit 14ac9a3)
1 parent 88f95d5 commit 5f421d7

21 files changed

+2475
-130
lines changed

pymongo/collection.py

Lines changed: 49 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1648,38 +1648,49 @@ def parallel_scan(self, num_cursors, session=None, **kwargs):
16481648

16491649
return cursors
16501650

1651+
def _count_cmd(self, session, sock_info, slave_ok, cmd, collation):
1652+
"""Internal count command helper."""
1653+
# XXX: "ns missing" checks can be removed when we drop support for
1654+
# MongoDB 3.0, see SERVER-17051.
1655+
res = self._command(
1656+
sock_info,
1657+
cmd,
1658+
slave_ok,
1659+
allowable_errors=["ns missing"],
1660+
codec_options=self.__write_response_codec_options,
1661+
read_concern=self.read_concern,
1662+
collation=collation,
1663+
session=session)
1664+
if res.get("errmsg", "") == "ns missing":
1665+
return 0
1666+
return int(res["n"])
1667+
16511668
def _count(self, cmd, collation=None, session=None):
16521669
"""Internal count helper."""
16531670
# XXX: "ns missing" checks can be removed when we drop support for
16541671
# MongoDB 3.0, see SERVER-17051.
16551672
def _cmd(session, server, sock_info, slave_ok):
1656-
res = self._command(
1657-
sock_info,
1658-
cmd,
1659-
slave_ok,
1660-
allowable_errors=["ns missing"],
1661-
codec_options=self.__write_response_codec_options,
1662-
read_concern=self.read_concern,
1663-
collation=collation,
1664-
session=session)
1665-
if res.get("errmsg", "") == "ns missing":
1666-
return 0
1667-
return int(res["n"])
1673+
return self._count_cmd(
1674+
session, sock_info, slave_ok, cmd, collation)
16681675

16691676
return self.__database.client._retryable_read(
16701677
_cmd, self._read_preference_for(session), session)
16711678

16721679
def _aggregate_one_result(
1673-
self, sock_info, slave_ok, cmd, collation=None, session=None):
1680+
self, sock_info, slave_ok, cmd, collation, session):
16741681
"""Internal helper to run an aggregate that returns a single result."""
16751682
result = self._command(
16761683
sock_info,
16771684
cmd,
16781685
slave_ok,
1686+
allowable_errors=[26], # Ignore NamespaceNotFound.
16791687
codec_options=self.__write_response_codec_options,
16801688
read_concern=self.read_concern,
16811689
collation=collation,
16821690
session=session)
1691+
# cursor will not be present for NamespaceNotFound errors.
1692+
if 'cursor' not in result:
1693+
return None
16831694
batch = result['cursor']['firstBatch']
16841695
return batch[0] if batch else None
16851696

@@ -1704,9 +1715,31 @@ def estimated_document_count(self, **kwargs):
17041715
if 'session' in kwargs:
17051716
raise ConfigurationError(
17061717
'estimated_document_count does not support sessions')
1707-
cmd = SON([('count', self.__name)])
1708-
cmd.update(kwargs)
1709-
return self._count(cmd)
1718+
1719+
def _cmd(session, server, sock_info, slave_ok):
1720+
if sock_info.max_wire_version >= 12:
1721+
# MongoDB 4.9+
1722+
pipeline = [
1723+
{'$collStats': {'count': {}}},
1724+
{'$group': {'_id': 1, 'n': {'$sum': '$count'}}},
1725+
]
1726+
cmd = SON([('aggregate', self.__name),
1727+
('pipeline', pipeline),
1728+
('cursor', {})])
1729+
cmd.update(kwargs)
1730+
result = self._aggregate_one_result(
1731+
sock_info, slave_ok, cmd, collation=None, session=session)
1732+
if not result:
1733+
return 0
1734+
return int(result['n'])
1735+
else:
1736+
# MongoDB < 4.9
1737+
cmd = SON([('count', self.__name)])
1738+
cmd.update(kwargs)
1739+
return self._count_cmd(None, sock_info, slave_ok, cmd, None)
1740+
1741+
return self.__database.client._retryable_read(
1742+
_cmd, self.read_preference, None)
17101743

17111744
def count_documents(self, filter, session=None, **kwargs):
17121745
"""Count the number of documents in this collection.

0 commit comments

Comments
 (0)