Skip to content

Commit

Permalink
MongoDB contrib : Next replaced with for loop to avoid propagation of…
Browse files Browse the repository at this point in the history
… StopIteration (spotify#2221)

* Replaced current logic to avoid return a StopIteration for the document count
  • Loading branch information
kritchie authored and dlstadther committed Sep 11, 2017
1 parent 65a6293 commit 1795eb1
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions luigi/contrib/mongodb.py
Original file line number Diff line number Diff line change
Expand Up @@ -212,8 +212,9 @@ def exists(self):

def read(self):
"""
Returns the count number of the target, -1 if the collection doesn't exist
Using the aggregate method to avoid inaccurate count if using a sharded cluster
https://docs.mongodb.com/manual/reference/method/db.collection.count/#behavior
"""
return next(self.get_collection().aggregate([{'$group': {'_id': None, 'count': {'$sum': 1}}}])).get('count')
for res in self.get_collection().aggregate([{'$group': {'_id': None, 'count': {'$sum': 1}}}]):
return res.get('count', None)
return None

0 comments on commit 1795eb1

Please sign in to comment.