add FlareCleanupTask (#947)

Author: nora-shap

celery_config.py

@@ -10,6 +10,7 @@
 from shared.celery_config import (
     BaseCeleryConfig,
     brolly_stats_rollup_task_name,
+    # flare_cleanup_task_name,
     gh_app_webhook_check_task_name,
     health_check_task_name,
     profiling_finding_task_name,
@@ -88,12 +89,18 @@ def _beat_schedule():
         },
         "trial_expiration_cron": {
             "task": trial_expiration_cron_task_name,
-            # 4 UTC is 12am EDT
-            "schedule": crontab(minute="0", hour="4"),
+            "schedule": crontab(minute="0", hour="4"),  # 4 UTC is 12am EDT
             "kwargs": {
                 "cron_task_generation_time_iso": BeatLazyFunc(get_utc_now_as_iso_format)
             },
         },
+        # "flare_cleanup": {
+        #     "task": flare_cleanup_task_name,
+        #     "schedule": crontab(minute="0", hour="5"),  # every day, 5am UTC (10pm PDT)
+        #     "kwargs": {
+        #         "cron_task_generation_time_iso": BeatLazyFunc(get_utc_now_as_iso_format)
+        #     },
+        # },
     }
 
     if get_config("setup", "find_uncollected_profilings", "enabled", default=True):

conftest.py

@@ -276,6 +276,18 @@ def mock_storage(mocker):
     return storage_server
 
 
+@pytest.fixture
+def mock_archive_storage(mocker):
+    m = mocker.patch("shared.api_archive.archive.StorageService")
+    use_archive = mocker.patch(
+        "shared.django_apps.core.models.should_write_data_to_storage_config_check"
+    )
+    use_archive.return_value = True
+    storage_server = MemoryStorageService({})
+    m.return_value = storage_server
+    return storage_server
+
+
 @pytest.fixture
 def mock_smtp(mocker):
     m = mocker.patch("services.smtp.SMTPService")

database/models/core.py

@@ -380,7 +380,7 @@ def get_repository(self):
     def get_commitid(self):
         return self.commitid
 
-    def should_write_to_storage(self) -> bool:
+    def should_write_to_storage(self: object) -> bool:
         if self.repository is None or self.repository.owner is None:
             return False
         is_codecov_repo = self.repository.owner.username == "codecov"
@@ -447,7 +447,6 @@ class Pull(CodecovBaseModel):
     commentid = Column(types.Text)
     bundle_analysis_commentid = Column(types.Text)
     diff = Column(postgresql.JSON)
-    flare = Column(postgresql.JSON)
     author_id = Column("author", types.Integer, ForeignKey("owners.ownerid"))
     behind_by = Column(types.Integer)
     behind_by_commit = Column(types.Text)
@@ -457,6 +456,22 @@ class Pull(CodecovBaseModel):
         Repository, backref=backref("pulls", cascade="delete", lazy="dynamic")
     )
 
+    def should_write_to_storage(self: object) -> bool:
+        if self.repository is None or self.repository.owner is None:
+            return False
+        is_codecov_repo = self.repository.owner.username == "codecov"
+        return should_write_data_to_storage_config_check(
+            master_switch_key="pull_flare",
+            is_codecov_repo=is_codecov_repo,
+            repoid=self.repository.repoid,
+        )
+
+    _flare = Column("flare", postgresql.JSON)
+    _flare_storage_path = Column("flare_storage_path", types.Text, nullable=True)
+    flare = ArchiveField(
+        should_write_to_storage_fn=should_write_to_storage, default_value_class=dict
+    )
+
     __table_args__ = (Index("pulls_repoid_pullid", "repoid", "pullid", unique=True),)
 
     def __repr__(self):
@@ -503,16 +518,6 @@ def external_id(self):
     def id(self):
         return self.id_
 
-    def should_write_to_storage(self) -> bool:
-        if self.repository is None or self.repository.owner is None:
-            return False
-        is_codecov_repo = self.repository.owner.username == "codecov"
-        return should_write_data_to_storage_config_check(
-            master_switch_key="pull_flare",
-            is_codecov_repo=is_codecov_repo,
-            repoid=self.repository.repoid,
-        )
-
     @cached_property
     def is_first_coverage_pull(self):
         """
@@ -536,12 +541,6 @@ def is_first_coverage_pull(self):
             return first_pull_with_coverage.id_ == self.id_
         return True
 
-    _flare = Column("flare", postgresql.JSON)
-    _flare_storage_path = Column("flare_storage_path", types.Text, nullable=True)
-    flare = ArchiveField(
-        should_write_to_storage_fn=should_write_to_storage, default_value_class=dict
-    )
-
 
 class CommitNotification(CodecovBaseModel):
     __tablename__ = "commit_notifications"

requirements.in

@@ -1,5 +1,5 @@
 https://github.com/codecov/test-results-parser/archive/c840502d1b4dd7d05b2efc2c1328affaf2acd27c.tar.gz#egg=test-results-parser
-https://github.com/codecov/shared/archive/2674ae99811767e63151590906691aed4c5ce1f9.tar.gz#egg=shared
+https://github.com/codecov/shared/archive/efe48352e172f658c21465371453dcefc98f6793.tar.gz#egg=shared
 https://github.com/codecov/timestring/archive/d37ceacc5954dff3b5bd2f887936a98a668dda42.tar.gz#egg=timestring
 asgiref>=3.7.2
 analytics-python==1.3.0b1

requirements.txt

@@ -336,7 +336,7 @@ sentry-sdk==2.13.0
     #   shared
 setuptools==75.6.0
     # via nodeenv
-shared @ https://github.com/codecov/shared/archive/2674ae99811767e63151590906691aed4c5ce1f9.tar.gz#egg=shared
+shared @ https://github.com/codecov/shared/archive/efe48352e172f658c21465371453dcefc98f6793.tar.gz#egg=shared
     # via -r requirements.in
 six==1.16.0
     # via

tasks/flare_cleanup.py

@@ -0,0 +1,100 @@
+import logging
+
+from shared.api_archive.archive import ArchiveService
+from shared.celery_config import flare_cleanup_task_name
+from shared.django_apps.core.models import Pull, PullStates
+
+from app import celery_app
+from tasks.crontasks import CodecovCronTask
+
+log = logging.getLogger(__name__)
+
+
+class FlareCleanupTask(CodecovCronTask, name=flare_cleanup_task_name):
+    """
+    Flare is a field on a Pull object.
+    Flare is used to draw static graphs (see GraphHandler view in api) and can be large.
+    The majority of flare graphs are used in pr comments, so we keep the (maybe large) flare "available"
+    in either the db or Archive storage while the pull is OPEN.
+    If the pull is not OPEN, we dump the flare to save space.
+    If we need to generate a flare graph for a non-OPEN pull, we build_report_from_commit
+    and generate fresh flare from that report (see GraphHandler view in api).
+    """
+
+    @classmethod
+    def get_min_seconds_interval_between_executions(cls):
+        return 72000  # 20h
+
+    def run_cron_task(self, db_session, batch_size=1000, limit=10000, *args, **kwargs):
+        # for any Pull that is not OPEN, clear the flare field(s), targeting older data
+        non_open_pulls = Pull.objects.exclude(state=PullStates.OPEN.value).order_by(
+            "updatestamp"
+        )
+
+        log.info("Starting FlareCleanupTask")
+
+        # clear in db
+        non_open_pulls_with_flare_in_db = non_open_pulls.filter(
+            _flare__isnull=False
+        ).exclude(_flare={})
+
+        # Process in batches
+        total_updated = 0
+        start = 0
+        while start < limit:
+            stop = start + batch_size if start + batch_size < limit else limit
+            batch = non_open_pulls_with_flare_in_db.values_list("id", flat=True)[
+                start:stop
+            ]
+            if not batch:
+                break
+            n_updated = non_open_pulls_with_flare_in_db.filter(id__in=batch).update(
+                _flare=None
+            )
+            total_updated += n_updated
+            start = stop
+
+        log.info(f"FlareCleanupTask cleared {total_updated} database flares")
+
+        # clear in Archive
+        non_open_pulls_with_flare_in_archive = non_open_pulls.filter(
+            _flare_storage_path__isnull=False
+        )
+
+        # Process archive deletions in batches
+        total_updated = 0
+        start = 0
+        while start < limit:
+            stop = start + batch_size if start + batch_size < limit else limit
+            batch = non_open_pulls_with_flare_in_archive.values_list("id", flat=True)[
+                start:stop
+            ]
+            if not batch:
+                break
+            flare_paths_from_batch = Pull.objects.filter(id__in=batch).values_list(
+                "_flare_storage_path", flat=True
+            )
+            try:
+                archive_service = ArchiveService(repository=None)
+                archive_service.delete_files(flare_paths_from_batch)
+            except Exception as e:
+                # if something fails with deleting from archive, leave the _flare_storage_path on the pull object.
+                # only delete _flare_storage_path if the deletion from archive was successful.
+                log.error(f"FlareCleanupTask failed to delete archive files: {e}")
+                continue
+
+            # Update the _flare_storage_path field for successfully processed pulls
+            n_updated = Pull.objects.filter(id__in=batch).update(
+                _flare_storage_path=None
+            )
+            total_updated += n_updated
+            start = stop
+
+        log.info(f"FlareCleanupTask cleared {total_updated} Archive flares")
+
+    def manual_run(self, db_session=None, limit=1000, *args, **kwargs):
+        self.run_cron_task(db_session, limit=limit, *args, **kwargs)
+
+
+RegisteredFlareCleanupTask = celery_app.register_task(FlareCleanupTask())
+flare_cleanup_task = celery_app.tasks[RegisteredFlareCleanupTask.name]