MongoDB: cursor.count() deprecated in favour of coll.count_documents()

Note that `ManagerManager.owned_managers` changed return type because of
this change.

flamenco/auth.py

@@ -105,10 +105,6 @@ def current_user_may(self, action: Actions, project_id: bson.ObjectId) -> bool:
             return True
 
         managers_coll = current_flamenco.db('managers')
-        managers = managers_coll.find({'projects': project_id})
+        managers_count = managers_coll.count_documents({'projects': project_id})
 
-        if self._log.isEnabledFor(logging.DEBUG):
-            self._log.debug('User has access to the following managers for this project: %s',
-                            [m['_id'] for m in managers])
-
-        return managers.count() > 0
+        return managers_count > 0

flamenco/celery/job_archival.py

@@ -73,6 +73,7 @@ def archive_job(job_id: str):
     # Run each task log compression in a separate Celery task.
     tasks_coll = current_flamenco.db('tasks')
     tasks = tasks_coll.find({'job': job_oid}, {'_id': 1})
+    tasks_count = tasks_coll.count_documents({'job': job_oid})
 
     # The chain of everything except downloading tasks & logs. Celery can't handle empty
     # groups, so we have to be careful in constructing the download_tasks group.
@@ -82,7 +83,7 @@ def archive_job(job_id: str):
             cleanup.si(storage_path)
     )
 
-    if tasks.count():
+    if tasks_count:
         download_tasks = celery.group(*(
             download_task_and_log.si(storage_path, str(task['_id']))
             for task in tasks))
@@ -100,12 +101,13 @@ def resume_job_archiving():
     """
     age = current_app.config['FLAMENCO_RESUME_ARCHIVING_AGE']
     jobs_coll = current_flamenco.db('jobs')
-    archiving = jobs_coll.find({
+    query = {
         'status': 'archiving',
         '_updated': {'$lte': utcnow() - age},
-    })
+    }
+    archiving = jobs_coll.find(query)
 
-    log.info('Resume archiving %d jobs', archiving.count())
+    log.info('Resume archiving %d jobs', jobs_coll.count_documents(query))
     for job in archiving:
         log.debug('Resume archiving job %s', job['_id'])
         archive_job.delay(str(job['_id']))

flamenco/jobs/__init__.py

@@ -200,8 +200,8 @@ def __job_status_if_a_then_b(if_status: str, then_new_status: str):
 
         if new_task_status == 'failed':
             # Count the number of failed tasks. If it is more than 10%, fail the job.
-            total_count = tasks_coll.find({'job': job_id}).count()
-            fail_count = tasks_coll.find({'job': job_id, 'status': 'failed'}).count()
+            total_count = tasks_coll.count_documents({'job': job_id})
+            fail_count = tasks_coll.count_documents({'job': job_id, 'status': 'failed'})
             fail_perc = fail_count / float(total_count) * 100
             if fail_perc >= TASK_FAIL_JOB_PERCENTAGE:
                 msg = f'Failing job {job_id} because {fail_count} of its {total_count} tasks ' \
@@ -374,8 +374,8 @@ def _do_check_completion(self, job_id, new_status) -> str:
         """
 
         tasks_coll = current_flamenco.db('tasks')
-        total_tasks = tasks_coll.find({'job': job_id}).count()
-        completed_tasks = tasks_coll.find({'job': job_id, 'status': 'completed'}).count()
+        total_tasks = tasks_coll.count_documents({'job': job_id})
+        completed_tasks = tasks_coll.count_documents({'job': job_id, 'status': 'completed'})
         if completed_tasks < total_tasks:
             # Not yet completed, so just stay at current status.
             self._log.debug('Job %s has %d of %d tasks completed, staying at status %r',

flamenco/managers/__init__.py

@@ -374,8 +374,8 @@ def share_unshare_manager(self, manager_id: bson.ObjectId, share_action: ShareAc
 
         # Check that there is at least one user left in the group.
         users_coll = current_app.db('users')
-        owners = users_coll.find({'groups': owner_gid})
-        if share_action == ShareAction.unshare and owners.count() < 2:
+        owner_count = users_coll.count_documents({'groups': owner_gid})
+        if share_action == ShareAction.unshare and owner_count < 2:
             self._log.warning('User %s tried to make Manager %s ownerless',
                               current_user.user_id, manager_id)
             raise ValueError('Manager cannot become ownerless.')
@@ -404,19 +404,24 @@ def managers_for_project(self, project_id: bson.ObjectId) -> typing.List[bson.Ob
         return [m['_id'] for m in managers]
 
     def owned_managers(self, user_group_ids: typing.List[bson.ObjectId],
-                       projection: typing.Optional[dict] = None) -> pymongo.cursor.Cursor:
+                       projection: typing.Optional[dict] = None) \
+            -> typing.Tuple[pymongo.cursor.Cursor, int]:
         """Returns a Mongo cursor of Manager object IDs owned by the given user.
 
         :param user_group_ids: list of the group IDs of the user.
         :param projection: When not None, it is used instead of the default {'_id': 1}.
+        :return: tuple (cursor, manager count)
         """
 
         if projection is None:
             projection = {'_id': 1}
 
         managers_coll = current_flamenco.db('managers')
-        managers = managers_coll.find({'owner': {'$in': user_group_ids}}, projection)
-        return managers
+        query = {'owner': {'$in': user_group_ids}}
+        manager_cursor = managers_coll.find(query, projection)
+        manager_count = managers_coll.count_documents(query)
+
+        return manager_cursor, manager_count
 
     def queue_task_log_request(self, manager_id: bson.ObjectId, job_id: bson.ObjectId,
                                task_id: bson.ObjectId):

flamenco/managers/linking_routes.py

@@ -44,10 +44,10 @@ def index():
 
     # Fetch available Managers.
     man_man = current_flamenco.manager_manager
-    managers = list(man_man.owned_managers(
-        current_user.group_ids, {'_id': 1, 'name': 1}))
-    manager_limit_reached = not current_user.has_cap('admin') and \
-                            len(managers) >= flamenco.auth.MAX_MANAGERS_PER_USER
+    manager_cursor, manager_count = man_man.owned_managers(
+        current_user.group_ids, {'_id': 1, 'name': 1})
+    manager_limit_reached = (not current_user.has_cap('admin')) and \
+                            manager_count >= flamenco.auth.MAX_MANAGERS_PER_USER
 
     # Get the query arguments
     identifier: str = request.args.get('identifier', '')
@@ -115,5 +115,5 @@ def index():
         return redirect(direct_to, 307)
 
     return render_template('flamenco/managers/linking/choose_manager.html',
-                           managers=managers,
+                           managers=list(manager_cursor),
                            can_create_manager=not manager_limit_reached)

flamenco/routes.py

@@ -187,20 +187,20 @@ def setup_for_flamenco(project: pillarsdk.Project):
 
     # Find the Managers available to this user, so we can auto-assign if there is exactly one.
     man_man = current_flamenco.manager_manager
-    managers = man_man.owned_managers([bson.ObjectId(gid) for gid in current_user.groups])
-    manager_count = managers.count()
+    managers, managers_count = man_man.owned_managers(
+        [bson.ObjectId(gid) for gid in current_user.groups])
 
     project_oid = str2id(project_id)
     user_id = current_user_id()
 
-    if manager_count == 0:
+    if managers_count == 0:
         _, mngr_doc, _ = man_man.create_new_manager('My Manager', '', user_id)
         assign_man_oid = mngr_doc['_id']
         log.info('Created and auto-assigning Manager %s to project %s upon setup for Flamenco.',
                  assign_man_oid, project_oid)
         man_man.api_assign_to_project(assign_man_oid, project_oid, 'assign')
 
-    elif manager_count == 1:
+    elif managers_count == 1:
         assign_manager = managers.next()
         assign_man_oid = str2id(assign_manager['_id'])
         log.info('Auto-assigning Manager %s to project %s upon setup for Flamenco.',

flamenco/setup.py

@@ -36,13 +36,12 @@ def create_manager(owner_email, name, description):
     :returns: tuple (mngr_doc, account, token)
     """
 
-    from pymongo.cursor import Cursor
-
     # Find the owner, to add it to the owner group afterward.
-    possible_owners: Cursor = current_app.db('users').find(
+    possible_owners = list(current_app.db('users').find(
         {'email': owner_email},
-        {'_id': 1, 'full_name': 1})
-    owner_count = possible_owners.count()
+        {'_id': 1, 'full_name': 1}))
+
+    owner_count = len(possible_owners)
     if owner_count == 0:
         raise ValueError(f'No user found with email address {owner_email}; '
                          'cannot assign ownership of Manager')

tests/abstract_flamenco_test.py

@@ -93,7 +93,7 @@ def create_manager_service_account(
 
             # Make sure there is an owner for this manager.
             users_coll = self.app.db('users')
-            count = users_coll.find({'email': owner_email}).count()
+            count = users_coll.count_documents({'email': owner_email})
             if count == 0:
                 self.create_user(user_id=ObjectId(), email=owner_email)
             elif count > 1:

tests/test_job_patch.py

@@ -351,7 +351,7 @@ def test_rna_overrides_new(self):
             render_tasks = list(tasks_coll.find({'job': self.job_id,
                                                  'task_type': 'blender-render',
                                                  'parents': [rm_tree_task['_id']]}))
-            task_count = tasks_coll.count({'job': self.job_id})
+            task_count = tasks_coll.count_documents({'job': self.job_id})
 
         # Just checking some assumptions this test relies on.
         for task in render_tasks:

tests/test_manager_linking.py

@@ -43,7 +43,7 @@ def test_exchange_secret_key_malformed_key(self):
 
         with self.app.app_context():
             coll = self.flamenco.db('manager_linking_keys')
-            self.assertEqual(0, coll.count())
+            self.assertEqual(0, coll.count_documents({}))
 
     def test_exchange_secret_key_no_key(self):
         self.post('/api/flamenco/managers/link/exchange',
@@ -52,7 +52,7 @@ def test_exchange_secret_key_no_key(self):
 
         with self.app.app_context():
             coll = self.flamenco.db('manager_linking_keys')
-            self.assertEqual(0, coll.count())
+            self.assertEqual(0, coll.count_documents({}))
 
     def test_reset_auth_token_happy(self):
         import secrets
@@ -68,7 +68,7 @@ def test_reset_auth_token_happy(self):
             self.post('/api/flamenco/managers/link/exchange', json={'key': 'aabbccddeeff'})
 
             coll = self.flamenco.db('manager_linking_keys')
-            self.assertEqual(2, coll.count())
+            self.assertEqual(2, coll.count_documents({}))
 
             # Bind them to the same Manager
             coll.update_many({}, {'$set': {'manager_id': manager_id}})

tests/test_manager_management.py

@@ -185,14 +185,14 @@ def test_gen_new_auth_token(self):
             service.generate_auth_token(service_account_id)
             service.generate_auth_token(service_account_id)
 
-            all_tokens = tokens_coll.find({'user': service_account_id})
-            self.assertEqual(all_tokens.count(), 4)
+            token_count = tokens_coll.count_documents({'user': service_account_id})
+            self.assertEqual(token_count, 4)
 
             token = self.flamenco.manager_manager.gen_new_auth_token(self.mngr_id)
 
             # There can be only one, rest should have been deleted.
-            all_tokens = tokens_coll.find({'user': service_account_id})
-            self.assertEqual(all_tokens.count(), 1)
+            token_count = tokens_coll.count_documents({'user': service_account_id})
+            self.assertEqual(token_count, 1)
 
         self.assertNotEqual(token.token, self.mngr_token)
         self.assertTrue(token.token.startswith('SRV'))
@@ -210,14 +210,14 @@ def test_revoke_auth_token(self):
             service.generate_auth_token(service_account_id)
             service.generate_auth_token(service_account_id)
 
-            all_tokens = tokens_coll.find({'user': service_account_id})
-            self.assertEqual(all_tokens.count(), 4)
+            token_count = tokens_coll.count_documents({'user': service_account_id})
+            self.assertEqual(token_count, 4)
 
             self.flamenco.manager_manager.revoke_auth_token(self.mngr_id)
 
             # All should have been deleted.
-            all_tokens = tokens_coll.find({'user': service_account_id})
-            self.assertEqual(all_tokens.count(), 0)
+            token_count = tokens_coll.count_documents({'user': service_account_id})
+            self.assertEqual(token_count, 0)
 
     def test_share(self):