(Archiving) Clean retrieved springs (#2815) (minor)

### Added

- CLI command for deleting retrieved Spring files.

Author: islean

cg/apps/housekeeper/hk.py

@@ -1,13 +1,17 @@
 """ Module to decouple cg code from Housekeeper code """
-import datetime as dt
 import logging
 import os
+from datetime import datetime
 from pathlib import Path
 
 from housekeeper.include import checksum as hk_checksum
 from housekeeper.include import include_version
 from housekeeper.store import Store, models
-from housekeeper.store.database import create_all_tables, drop_all_tables, initialize_database
+from housekeeper.store.database import (
+    create_all_tables,
+    drop_all_tables,
+    initialize_database,
+)
 from housekeeper.store.models import Archive, Bundle, File, Version
 from sqlalchemy.orm import Query
 
@@ -33,7 +37,7 @@ def __getattr__(self, name):
         LOG.warning(f"Called undefined {name} on {self.__class__.__name__}, please wrap")
         return getattr(self._store, name)
 
-    def new_bundle(self, name: str, created_at: dt.datetime = None) -> Bundle:
+    def new_bundle(self, name: str, created_at: datetime = None) -> Bundle:
         """Create a new file bundle."""
         return self._store.new_bundle(name, created_at)
 
@@ -237,11 +241,11 @@ def include_file(self, file_obj: File, version_obj: Version) -> File:
         file_obj.path = str(new_path).replace(f"{global_root_dir}/", "", 1)
         return file_obj
 
-    def new_version(self, created_at: dt.datetime, expires_at: dt.datetime = None) -> Version:
+    def new_version(self, created_at: datetime, expires_at: datetime = None) -> Version:
         """Create a new bundle version."""
         return self._store.new_version(created_at, expires_at)
 
-    def version(self, bundle: str, date: dt.datetime) -> Version:
+    def version(self, bundle: str, date: datetime) -> Version:
         """Fetch a version."""
         LOG.debug(f"Return version: {date}, from {bundle}")
         return self._store.get_version_by_date_and_bundle_name(
@@ -286,7 +290,7 @@ def get_create_version(self, bundle_name: str) -> Version:
             bundle_result: tuple[Bundle, Version] = self.add_bundle(
                 bundle_data={
                     "name": bundle_name,
-                    "created_at": dt.datetime.now(),
+                    "created_at": datetime.now(),
                     "expires_at": None,
                     "files": [],
                 }
@@ -316,7 +320,7 @@ def get_tag_names_from_file(file: File) -> list[str]:
     def include(self, version_obj: Version):
         """Call the include version function to import related assets."""
         include_version(self.get_root_dir(), version_obj)
-        version_obj.included_at = dt.datetime.now()
+        version_obj.included_at = datetime.now()
 
     def add_commit(self, obj):
         """Wrap method in Housekeeper Store."""
@@ -376,7 +380,7 @@ def include_files_to_latest_version(self, bundle_name: str) -> None:
             LOG.warning(
                 f"File is already included in Housekeeper for bundle: {bundle_name}, version: {bundle_version}"
             )
-        bundle_version.included_at = dt.datetime.now()
+        bundle_version.included_at = datetime.now()
         self.commit()
 
     def get_file_from_latest_version(self, bundle_name: str, tags: set[str]) -> File | None:
@@ -649,3 +653,14 @@ def update_archive_retrieved_at(
         for archive in archives_to_update:
             archive.retrieval_task_id = new_retrieval_job_id
         self._store.session.commit()
+
+    def get_spring_files_retrieved_before(self, date: datetime):
+        return self._store.get_files_retrieved_before(date, tag_names=[SequencingFileTag.SPRING])
+
+    def reset_retrieved_archive_data(self, files_to_reset: list[File]):
+        """Resets 'retrieval_task_id' and 'retrieved_at' for all files' corresponding archive entries"""
+        for file in files_to_reset:
+            LOG.debug(f"Resetting retrieval data for file {file.path}")
+            file.archive.retrieval_task_id = None
+            file.archive.retrieved_at = None
+        self.commit()

cg/cli/clean.py

@@ -27,6 +27,7 @@
 from cg.exc import CleanFlowCellFailedError, FlowCellError
 from cg.meta.clean.api import CleanAPI
 from cg.meta.clean.clean_flow_cells import CleanFlowCellAPI
+from cg.meta.clean.clean_retrieved_spring_files import CleanRetrievedSpringFilesAPI
 from cg.models.cg_config import CGConfig
 from cg.store import Store
 from cg.store.models import Analysis
@@ -277,6 +278,24 @@ def clean_flow_cells(context: CGConfig, dry_run: bool):
             continue
 
 
+@clean.command("retrieved-spring-files")
+@click.option(
+    "--age-limit",
+    type=int,
+    default=7,
+    help="Clean all Spring files which were retrieved more than given amount of days ago.",
+    show_default=True,
+)
+@DRY_RUN
+@click.pass_obj
+def clean_retrieved_spring_files(context: CGConfig, age_limit: int, dry_run: bool):
+    """Clean Spring files which were retrieved more than given amount of days ago."""
+    clean_retrieved_spring_files_api = CleanRetrievedSpringFilesAPI(
+        housekeeper_api=context.housekeeper_api, dry_run=dry_run
+    )
+    clean_retrieved_spring_files_api.clean_retrieved_spring_files(age_limit)
+
+
 def _get_confirm_question(bundle, file_obj) -> str:
     """Return confirmation question."""
     return (

cg/meta/clean/clean_retrieved_spring_files.py

@@ -0,0 +1,46 @@
+import logging
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from housekeeper.store.models import File
+
+from cg.apps.housekeeper.hk import HousekeeperAPI
+
+LOG = logging.getLogger(__name__)
+
+
+class CleanRetrievedSpringFilesAPI:
+    """API for cleaning archived Spring files which have been retrieved."""
+
+    def __init__(self, housekeeper_api: HousekeeperAPI, dry_run: bool):
+        self.housekeeper_api: HousekeeperAPI = housekeeper_api
+        self.dry_run = dry_run
+
+    def _get_files_to_remove(self, age_limit: int) -> list[File]:
+        """Returns all Spring files which were retrieved more than given amount of days ago."""
+        return self.housekeeper_api.get_spring_files_retrieved_before(
+            date=datetime.now() - timedelta(days=age_limit)
+        )
+
+    def _unlink_files(self, files_to_unlink: list[File]) -> None:
+        for file in files_to_unlink:
+            file_path: str = file.full_path
+            if self.dry_run:
+                LOG.info(f"Dry run - would have unlinked {file_path}")
+            LOG.info(f"Unlinking {file_path}")
+            Path(file_path).unlink(missing_ok=True)
+
+    def clean_retrieved_spring_files(self, age_limit: int):
+        """Removes Spring files retrieved more than given amount of days ago from Hasta,
+        and resets retrieval data in Housekeeper."""
+
+        LOG.info("Starting cleaning of retrieved Spring files.")
+        files_to_remove: list[File] = self._get_files_to_remove(age_limit)
+        if not files_to_remove:
+            LOG.info("No retrieved Spring files to remove - exiting.")
+            return
+        self._unlink_files(files_to_remove)
+        if not self.dry_run:
+            self.housekeeper_api.reset_retrieved_archive_data(files_to_remove)
+        else:
+            LOG.info("Would have reset the files' retrieval data")

poetry.lock

@@ -69,7 +69,7 @@ urllib3 = "*"
 
 # Apps
 genologics = "*"
-housekeeper = "*"
+housekeeper = ">=4.11.1"
 
 
 [tool.poetry.dev-dependencies]

pyproject.toml

@@ -8,6 +8,7 @@
 from cg.constants import SequencingFileTag
 from cg.constants.subject import Sex
 from cg.meta.clean.clean_flow_cells import CleanFlowCellAPI
+from cg.meta.clean.clean_retrieved_spring_files import CleanRetrievedSpringFilesAPI
 from cg.models.flow_cell.flow_cell import FlowCellDirectoryData
 from cg.store import Store
 from cg.store.models import Flowcell, Sample
@@ -234,3 +235,119 @@ def hk_sample_bundle_for_flow_cell_not_to_clean(
         "expires": timestamp_yesterday,
         "files": [],
     }
+
+
+@pytest.fixture
+def clean_retrieved_spring_files_api(
+    real_housekeeper_api: HousekeeperAPI, tmp_path
+) -> CleanRetrievedSpringFilesAPI:
+    """Returns a CleanRetrievedSpringFilesAPI."""
+    real_housekeeper_api.root_dir = tmp_path
+    return CleanRetrievedSpringFilesAPI(housekeeper_api=real_housekeeper_api, dry_run=False)
+
+
+@pytest.fixture
+def path_to_old_retrieved_spring_file() -> str:
+    return Path("path", "to", "old_retrieved_spring_file").as_posix()
+
+
+@pytest.fixture
+def path_to_newly_retrieved_spring_file() -> str:
+    return Path("path", "to", "newly_retrieved_spring_file").as_posix()
+
+
+@pytest.fixture
+def path_to_archived_but_not_retrieved_spring_file() -> str:
+    return Path("path", "to", "archived_spring_file").as_posix()
+
+
+@pytest.fixture
+def path_to_fastq_file() -> str:
+    return Path("path", "to", "fastq_file").as_posix()
+
+
+@pytest.fixture
+def paths_for_populated_clean_retrieved_spring_files_api(
+    path_to_old_retrieved_spring_file: str,
+    path_to_newly_retrieved_spring_file: str,
+    path_to_archived_but_not_retrieved_spring_file: str,
+    path_to_fastq_file: str,
+) -> list[str]:
+    return [
+        path_to_old_retrieved_spring_file,
+        path_to_newly_retrieved_spring_file,
+        path_to_archived_but_not_retrieved_spring_file,
+        path_to_fastq_file,
+    ]
+
+
+@pytest.fixture
+def retrieved_test_bundle_name() -> str:
+    return "retrieved_test_bundle"
+
+
+@pytest.fixture
+def path_to_old_spring_file_in_housekeeper(
+    retrieved_test_bundle_name: str, path_to_old_retrieved_spring_file
+) -> str:
+    return Path(
+        retrieved_test_bundle_name,
+        str(datetime.today().date()),
+        Path(path_to_old_retrieved_spring_file).name,
+    ).as_posix()
+
+
+@pytest.fixture
+def populated_clean_retrieved_spring_files_api(
+    clean_retrieved_spring_files_api: CleanRetrievedSpringFilesAPI,
+    paths_for_populated_clean_retrieved_spring_files_api: list[str],
+    retrieved_test_bundle_name: str,
+    archival_job_id_miria: int,
+    retrieval_job_id_miria: int,
+    timestamp: datetime,
+    timestamp_yesterday: datetime,
+    old_timestamp: datetime,
+    tmp_path,
+) -> CleanRetrievedSpringFilesAPI:
+    """
+    Returns a populated CleanRetrievedSpringFilesAPI, containing a bundle with one version and the following files:
+        - an archived Spring file which has not been retrieved
+        - an archived Spring file which was retrieved 1 day ago
+        - an archived Spring file which was retrieved in the year 1900
+        - a Fastq file
+    """
+    clean_retrieved_spring_files_api.housekeeper_api.add_bundle_and_version_if_non_existent(
+        bundle_name=retrieved_test_bundle_name
+    )
+    clean_retrieved_spring_files_api.housekeeper_api.commit()
+    for path in paths_for_populated_clean_retrieved_spring_files_api:
+        tags: list[str] = (
+            [SequencingFileTag.SPRING]
+            if SequencingFileTag.SPRING in path
+            else [SequencingFileTag.FASTQ]
+        )
+        Path(tmp_path, path).parent.mkdir(parents=True, exist_ok=True)
+        file_to_add = Path(tmp_path, path)
+        file_to_add.touch()
+        clean_retrieved_spring_files_api.housekeeper_api.add_and_include_file_to_latest_version(
+            file=file_to_add, bundle_name=retrieved_test_bundle_name, tags=tags
+        )
+    for file in clean_retrieved_spring_files_api.housekeeper_api.get_files(
+        bundle=retrieved_test_bundle_name
+    ):
+        Path(file.full_path).parent.mkdir(parents=True, exist_ok=True)
+        Path(file.full_path).touch()
+        if "spring" in file.path:
+            clean_retrieved_spring_files_api.housekeeper_api.add_archives(
+                files=[file], archive_task_id=archival_job_id_miria
+            )
+            file.archive.archived_at = timestamp
+            if "retrieved" in file.path:
+                file.archive.retrieval_task_id = retrieval_job_id_miria
+                file.archive.retrieved_at = (
+                    old_timestamp if "old" in file.path else timestamp_yesterday
+                )
+            clean_retrieved_spring_files_api.housekeeper_api.add_commit(file.archive)
+
+    clean_retrieved_spring_files_api.housekeeper_api.commit()
+    return clean_retrieved_spring_files_api

tests/meta/clean/conftest.py

@@ -0,0 +1,57 @@
+import logging
+from pathlib import Path
+
+from housekeeper.store.models import File
+
+from cg.meta.clean.clean_retrieved_spring_files import CleanRetrievedSpringFilesAPI
+
+
+def test_get_files_to_remove(
+    populated_clean_retrieved_spring_files_api: CleanRetrievedSpringFilesAPI,
+    path_to_old_retrieved_spring_file: str,
+    path_to_old_spring_file_in_housekeeper: str,
+):
+    """Tests that only old retrieved files are cleaned with clean_retrieved_spring_files. With the provided populated
+    api, this should not return a newly retrieved spring file, a fastq file nor an archived spring file which
+    has not been retrieved."""
+
+    # GIVEN a CleanRetrievedSpringFilesAPI with a populated Housekeeper database
+
+    # WHEN getting files to remove when cleaning retrieved spring files
+    files_to_remove: list[File] = populated_clean_retrieved_spring_files_api._get_files_to_remove(
+        age_limit=7
+    )
+
+    # THEN only the file with an old enough 'retrieved_at' should be returned
+    assert [file.path for file in files_to_remove] == [path_to_old_spring_file_in_housekeeper]
+
+
+def test_clean_retrieved_spring_files_dry_run(
+    populated_clean_retrieved_spring_files_api: CleanRetrievedSpringFilesAPI,
+    path_to_old_retrieved_spring_file: str,
+    retrieved_test_bundle_name: str,
+    caplog,
+    path_to_old_spring_file_in_housekeeper: str,
+):
+    """Tests that only the Spring file with an old enough 'retrieved_at' would be removed when cleaning retrieved
+    Spring files."""
+
+    caplog.set_level(logging.INFO)
+
+    # GIVEN a CleanRetrievedSpringFilesAPI with a populated Housekeeper database
+    # GIVEN that an old retrieved file exists
+    files = populated_clean_retrieved_spring_files_api.housekeeper_api.files(
+        bundle=retrieved_test_bundle_name
+    ).all()
+    for file in files:
+        assert Path(file.full_path).exists()
+
+    # WHEN running 'clean_retrieved_spring_files'
+    populated_clean_retrieved_spring_files_api.clean_retrieved_spring_files(age_limit=7)
+
+    # THEN only the file with an old enough 'retrieved_at' should have been removed
+    for file in files:
+        if file.path == path_to_old_spring_file_in_housekeeper:
+            assert not Path(file.full_path).exists()
+        else:
+            assert Path(file.full_path).exists()