make yield amount a param (#446)

codecov · May 13, 2024 · 966d2e8 · 966d2e8
1 parent 4c28604
commit 966d2e8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/tasks/backfill_existing_gh_app_installations.py b/tasks/backfill_existing_gh_app_installations.py
@@ -15,14 +15,16 @@
 
 log = logging.getLogger(__name__)
 
-YIELD_AMOUNT = 100
-
 
 class BackfillExistingGHAppInstallationsTask(
     BaseCodecovTask, name=backfill_existing_gh_app_installations_name
 ):
     def backfill_existing_gh_apps(
-        self, db_session: Session, owner_ids: List[int] = None, missed_owner_ids=[]
+        self,
+        db_session: Session,
+        owner_ids: List[int] = None,
+        missed_owner_ids=[],
+        yield_amount: int = 1000,
     ):
         # Get owners that have installations, and installations queries
         owners_query = (
@@ -42,7 +44,7 @@ def backfill_existing_gh_apps(
             )
 
         gh_app_installations: List[GithubAppInstallation] = (
-            gh_app_installations_query.yield_per(YIELD_AMOUNT)
+            gh_app_installations_query.yield_per(yield_amount)
         )
 
         for gh_app_installation in gh_app_installations:

diff --git a/tasks/backfill_owners_without_gh_app_installations.py b/tasks/backfill_owners_without_gh_app_installations.py
@@ -20,14 +20,16 @@
 
 log = logging.getLogger(__name__)
 
-YIELD_AMOUNT = 100
-
 
 class BackfillOwnersWithoutGHAppInstallations(
     BaseCodecovTask, name=backfill_owners_without_gh_app_installations_name
 ):
     def backfill_owners_with_integration_without_gh_app(
-        self, db_session: Session, owner_ids: List[int] = None, missed_owner_ids=[]
+        self,
+        db_session: Session,
+        owner_ids: List[int] = None,
+        missed_owner_ids=[],
+        yield_amount: int = 1000,
     ):
         owners_with_integration_id_without_gh_app_query = (
             db_session.query(Owner)
@@ -50,7 +52,7 @@ def backfill_owners_with_integration_without_gh_app(
             )
 
         owners: List[Owner] = owners_with_integration_id_without_gh_app_query.yield_per(
-            YIELD_AMOUNT
+            yield_amount
         )
 
         for owner in owners: