fix(grz-cli): Submission id upload tracking

packages/grz-common/src/grz_common/progress/states.py

@@ -11,6 +11,7 @@ class State(TypedDict, total=False):
     """
 
     errors: list[str]
+    submission_id: str
 
 
 class ValidationState(State):

packages/grz-common/src/grz_common/workers/download.py

@@ -164,6 +164,7 @@ def download_file(
         s3_object_id: str,
         progress_logger: FileProgressLogger[DownloadState],
         file_metadata: SubmissionFileMetadata,
+        submission_id: str,
     ):
         """
         Download a single file from S3 to the specified local_file_path.
@@ -179,7 +180,11 @@ def download_file(
             self._download_with_progress(str(local_file_path), s3_object_id)
 
             self.__log.info(f"Download complete for {str(local_file_path)}.")
-            progress_logger.set_state(local_file_path, file_metadata, state=DownloadState(download_successful=True))
+            progress_logger.set_state(
+                local_file_path,
+                file_metadata,
+                state=DownloadState(download_successful=True, submission_id=submission_id),
+            )
 
         except botocore.exceptions.ClientError as e:
             if e.response.get("Error", {}).get("Code") == "404":
@@ -190,13 +195,17 @@ def download_file(
                 exc = e  # type: ignore[assignment]
             self.__log.error(error_msg)
             progress_logger.set_state(
-                local_file_path, file_metadata, state=DownloadState(download_successful=False, errors=[str(exc)])
+                local_file_path,
+                file_metadata,
+                state=DownloadState(download_successful=False, errors=[str(exc)], submission_id=submission_id),
             )
             raise exc from e
         except Exception as e:
             self.__log.error("Download failed for '%s': %s", str(local_file_path), e)
             progress_logger.set_state(
-                local_file_path, file_metadata, state=DownloadState(download_successful=False, errors=[str(e)])
+                local_file_path,
+                file_metadata,
+                state=DownloadState(download_successful=False, errors=[str(e)], submission_id=submission_id),
             )
             raise e
 
@@ -217,7 +226,11 @@ def download(self, submission_id: str, encrypted_submission: EncryptedSubmission
             file_key = f"{submission_id}/files/{relative_encrypted_path}"
 
             logged_state = progress_logger.get_state(local_file_path, file_metadata)
-            if logged_state and logged_state.get("download_successful"):
+            if (
+                logged_state
+                and logged_state.get("download_successful")
+                and logged_state.get("submission_id") == submission_id
+            ):
                 self.__log.info(
                     "File '%s' already downloaded (at '%s'), skipping.",
                     file_key,
@@ -226,7 +239,7 @@ def download(self, submission_id: str, encrypted_submission: EncryptedSubmission
                 continue
 
             self.__log.info("Downloading file: '%s' -> '%s'", file_key, str(local_file_path))
-            self.download_file(local_file_path, file_key, progress_logger, file_metadata)
+            self.download_file(local_file_path, file_key, progress_logger, file_metadata, submission_id)
 
 
 class InboxSubmissionState(enum.StrEnum):

packages/grz-common/src/grz_common/workers/submission.py

@@ -259,6 +259,10 @@ def __init__(self, metadata_dir: str | PathLike, files_dir: str | PathLike):
 
         self.metadata = SubmissionMetadata(self.metadata_dir / "metadata.json")
 
+    @property
+    def submission_id(self) -> str:
+        return self.metadata.content.submission_id
+
     @property
     def files(self) -> dict[Path, SubmissionFileMetadata]:
         """
@@ -312,14 +316,19 @@ def validate_files(  # noqa: C901, PLR0912, PLR0915
         tasks = []
 
         def should_check_file(file_path: Path, file_metadata: SubmissionFileMetadata) -> bool:
-            # Check against both logs. If either is missing a "pass", re-check.
             checksum_state = checksum_progress_logger.get_state(file_path, file_metadata)
             seq_data_state = seq_data_progress_logger.get_state(file_path, file_metadata)
             is_seq_file = file_metadata.file_type in ("fastq", "bam")
-
-            checksum_passed = checksum_state and checksum_state.get("validation_passed")
-            seq_data_passed = seq_data_state and seq_data_state.get("validation_passed")
-
+            checksum_passed = (
+                checksum_state
+                and checksum_state.get("validation_passed")
+                and checksum_state.get("submission_id") == self.submission_id
+            )
+            seq_data_passed = (
+                seq_data_state
+                and seq_data_state.get("validation_passed")
+                and seq_data_state.get("submission_id") == self.submission_id
+            )
             if is_seq_file:
                 return not (checksum_passed and seq_data_passed)
             return not checksum_passed
@@ -448,11 +457,15 @@ def _execute_task(task_type, paths, metas, kwargs, pbar):
                         checksum_issues.append("File not found for size check.")
 
                     checksum_passed = not checksum_issues
-                    checksum_state = ValidationState(errors=checksum_issues, validation_passed=checksum_passed)
+                    checksum_state = ValidationState(
+                        errors=checksum_issues, validation_passed=checksum_passed, submission_id=self.submission_id
+                    )
                     checksum_progress_logger.set_state(file_path, file_metadata, checksum_state)
 
                     if file_metadata.file_type in ("fastq", "bam"):
-                        seq_data_state = ValidationState(errors=report.errors, validation_passed=report.is_valid)
+                        seq_data_state = ValidationState(
+                            errors=report.errors, validation_passed=report.is_valid, submission_id=self.submission_id
+                        )
                         seq_data_progress_logger.set_state(file_path, file_metadata, seq_data_state)
 
                 if not no_mmap:
@@ -524,6 +537,7 @@ def encrypt(
                 (logged_state is None)
                 or not logged_state.get("encryption_successful", False)
                 or not encrypted_file_path.is_file()
+                or logged_state.get("submission_id") != self.submission_id
             ):
                 self.__log.info(
                     "Encrypting file: '%s' -> '%s'",
@@ -543,15 +557,17 @@ def encrypt(
                     progress_logger.set_state(
                         file_path,
                         file_metadata,
-                        state=EncryptionState(encryption_successful=True),
+                        state=EncryptionState(encryption_successful=True, submission_id=self.submission_id),
                     )
                 except Exception as e:
                     self.__log.error("Encryption failed for '%s'", str(file_path))
 
                     progress_logger.set_state(
                         file_path,
                         file_metadata,
-                        state=EncryptionState(encryption_successful=False, errors=[str(e)]),
+                        state=EncryptionState(
+                            encryption_successful=False, errors=[str(e)], submission_id=self.submission_id
+                        ),
                     )
 
                     raise e
@@ -704,6 +720,7 @@ def decrypt(
             if (
                 (logged_state is None)
                 or not logged_state.get("decryption_successful", False)
+                or logged_state.get("submission_id") != self.submission_id
                 or not decrypted_file_path.is_file()
             ):
                 self.__log.info(
@@ -719,15 +736,17 @@ def decrypt(
                     progress_logger.set_state(
                         encrypted_file_path,
                         file_metadata,
-                        state=DecryptionState(decryption_successful=True),
+                        state=DecryptionState(decryption_successful=True, submission_id=self.submission_id),
                     )
                 except Exception as e:
                     self.__log.error("Decryption failed for '%s'", str(encrypted_file_path))
 
                     progress_logger.set_state(
                         encrypted_file_path,
                         file_metadata,
-                        state=DecryptionState(decryption_successful=False, errors=[str(e)]),
+                        state=DecryptionState(
+                            decryption_successful=False, errors=[str(e)], submission_id=self.submission_id
+                        ),
                     )
 
                     raise e

packages/grz-common/src/grz_common/workers/upload.py

@@ -168,8 +168,11 @@ def _upload_logged_files(self, encrypted_submission, progress_logger, files_to_u
             self.__log.debug("state for %s: %s", file_path, logged_state)
 
             s3_object_id = files_to_upload[file_path]
-
-            if (logged_state is None) or not logged_state.get("upload_successful", False):
+            if (
+                (logged_state is None)
+                or not logged_state.get("upload_successful", False)
+                or logged_state.get("submission_id") != encrypted_submission.submission_id
+            ):
                 self.__log.info(
                     "Uploading file: '%s' -> '%s'",
                     str(file_path),
@@ -183,22 +186,25 @@ def _upload_logged_files(self, encrypted_submission, progress_logger, files_to_u
                     progress_logger.set_state(
                         file_path,
                         file_metadata,
-                        state=UploadState(upload_successful=True),
+                        state=UploadState(upload_successful=True, submission_id=encrypted_submission.submission_id),
                     )
                 except Exception as e:
                     self.__log.error("Upload failed for '%s'", str(file_path))
 
                     progress_logger.set_state(
                         file_path,
                         file_metadata,
-                        state=UploadState(upload_successful=False, errors=[str(e)]),
+                        state=UploadState(
+                            upload_successful=False, errors=[str(e)], submission_id=encrypted_submission.submission_id
+                        ),
                     )
 
                     raise e
             else:
                 self.__log.info(
-                    "File '%s' already uploaded (at '%s')",
+                    "File '%s' already uploaded for submission '%s' (at '%s')",
                     str(file_path),
+                    encrypted_submission.submission_id,
                     str(s3_object_id),
                 )
 

tests/test_download.py

@@ -67,3 +67,102 @@ def test_boto_download(
     assert calculate_sha256(files_dir / "small_test_file.txt") == temp_small_file_sha256sum, (
         "Text file SHA256 mismatch."
     )
+
+
+def test_download_skips_file_already_downloaded_for_same_submission(
+    s3_config_model,
+    remote_bucket,
+    temp_download_log_file_path,
+    encrypted_submission,
+    mocker,
+):
+    """Files with download_successful=True and matching submission_id should be skipped."""
+    from grz_common.progress.progress_logging import FileProgressLogger
+    from grz_common.progress.states import DownloadState
+
+    download_worker = S3BotoDownloadWorker(
+        s3_options=s3_config_model.s3,
+        status_file_path=temp_download_log_file_path,
+    )
+
+    progress_logger = FileProgressLogger[DownloadState](temp_download_log_file_path)
+    for file_path, file_metadata in encrypted_submission.encrypted_files.items():
+        progress_logger.set_state(
+            file_path,
+            file_metadata,
+            state=DownloadState(download_successful=True, submission_id=encrypted_submission.submission_id),
+        )
+
+    download_spy = mocker.spy(download_worker, "download_file")
+    download_worker.download(encrypted_submission.submission_id, encrypted_submission)
+
+    assert download_spy.call_count == 0, (
+        f"Expected all files to be skipped, but {download_spy.call_count} were downloaded"
+    )
+
+
+def test_download_redownloads_file_with_different_submission_id(
+    s3_config_model,
+    remote_bucket,
+    temp_download_log_file_path,
+    encrypted_submission,
+    mocker,
+):
+    """Files logged as download_successful=True for a different submission_id must be re-downloaded."""
+    from grz_common.progress.progress_logging import FileProgressLogger
+    from grz_common.progress.states import DownloadState
+
+    download_worker = S3BotoDownloadWorker(
+        s3_options=s3_config_model.s3,
+        status_file_path=temp_download_log_file_path,
+    )
+
+    progress_logger = FileProgressLogger[DownloadState](temp_download_log_file_path)
+    for file_path, file_metadata in encrypted_submission.encrypted_files.items():
+        progress_logger.set_state(
+            file_path,
+            file_metadata,
+            state=DownloadState(download_successful=True, submission_id="different-submission-id-9999"),
+        )
+
+    mock_download = mocker.patch.object(download_worker, "download_file")
+    download_worker.download(encrypted_submission.submission_id, encrypted_submission)
+
+    expected = len(encrypted_submission.encrypted_files)
+    assert mock_download.call_count == expected, (
+        f"Expected {expected} files to be re-downloaded for a different submission_id, "
+        f"but only {mock_download.call_count} were downloaded"
+    )
+
+
+def test_download_redownloads_file_after_failed_download(
+    s3_config_model,
+    remote_bucket,
+    temp_download_log_file_path,
+    encrypted_submission,
+    mocker,
+):
+    """Files logged as download_successful=False must be retried even with matching submission_id."""
+    from grz_common.progress.progress_logging import FileProgressLogger
+    from grz_common.progress.states import DownloadState
+
+    download_worker = S3BotoDownloadWorker(
+        s3_options=s3_config_model.s3,
+        status_file_path=temp_download_log_file_path,
+    )
+
+    progress_logger = FileProgressLogger[DownloadState](temp_download_log_file_path)
+    for file_path, file_metadata in encrypted_submission.encrypted_files.items():
+        progress_logger.set_state(
+            file_path,
+            file_metadata,
+            state=DownloadState(download_successful=False, submission_id=encrypted_submission.submission_id),
+        )
+
+    mock_download = mocker.patch.object(download_worker, "download_file")
+    download_worker.download(encrypted_submission.submission_id, encrypted_submission)
+
+    expected = len(encrypted_submission.encrypted_files)
+    assert mock_download.call_count == expected, (
+        f"Expected {expected} failed files to be retried, but only {mock_download.call_count} were downloaded"
+    )