Add download script as artifacts (#250)

Author: phwissmann

backend/archiver/flows/archive_datasets_flow.py

@@ -181,7 +181,7 @@ def update_progress(progress):
         prefix=prefix,
         bucket=Bucket.staging_bucket(),
         source_folder=datablocks_scratch_folder,
-        ext=".gz",
+        ext=".tar",
         progress_callback=update_progress,
     )
 
@@ -225,7 +225,7 @@ def calculate_checksum(dataset_id: str, datablock: DataBlock):
 
 @task(task_run_name=generate_task_name_dataset, tags=[ConcurrencyLimits().LTS_WRITE_TAG])
 def move_data_to_LTS(dataset_id: str, datablock: DataBlock):
-    """Prefect task to move a datablock (.tar.gz file) to the LTS. Concurrency of this task is limited to 2 instances
+    """Prefect task to move a datablock (.tar file) to the LTS. Concurrency of this task is limited to 2 instances
     at the same time.
     """
     datablocks_operations.move_data_to_LTS(dataset_id, datablock)
@@ -238,7 +238,7 @@ def move_data_to_LTS(dataset_id: str, datablock: DataBlock):
     retry_delay_seconds=[60, 120, 240, 480, 960],
 )
 def copy_datablock_from_LTS(dataset_id: str, datablock: DataBlock):
-    """Prefect task to move a datablock (.tar.gz file) to the LTS. Concurrency of this task is limited to 2 instances
+    """Prefect task to move a datablock (.tar file) to the LTS. Concurrency of this task is limited to 2 instances
     at the same time.
     """
     datablocks_operations.copy_file_from_LTS(dataset_id, datablock)
@@ -312,7 +312,7 @@ def move_datablocks_to_lts_flow(dataset_id: str, datablocks: List[DataBlock]):
 
 @flow(name="create_datablocks", flow_run_name=generate_subflow_run_name_job_id_dataset_id)
 def create_datablocks_flow(dataset_id: str) -> List[DataBlock]:
-    """Prefect (sub-)flow to create datablocks (.tar.gz files) for files of a dataset and register them in Scicat.
+    """Prefect (sub-)flow to create datablocks (.tar files) for files of a dataset and register them in Scicat.
 
     Args:
         dataset_id (str): Dataset id

backend/archiver/flows/mock_flows.py

@@ -340,7 +340,7 @@ def verify_data_from_minio(dataset_pid, datablock_name, datablock_url):
     # Verify retrieved datablock in MINIO
     retrieved_datablock = s3_client.stat_object(
         bucket=Bucket("retrieval"),
-        filename=f"openem-network/datasets/{dataset_pid}/datablocks/{dataset_pid}_0.tar.gz",
+        filename=f"openem-network/datasets/{dataset_pid}/datablocks/{dataset_pid}_0.tar",
     )
     assert retrieved_datablock is not None
     assert retrieved_datablock.Size > 80 * 1024 * 1024

backend/archiver/flows/retrieve_datasets_flow.py

@@ -224,16 +224,14 @@ async def retrieve_datasets_flow(job_id: UUID):
         else:
             await wait_for_retrieval_flow(existing_run_id)
 
-    job_results_future = create_job_result_object_task.submit(dataset_ids=dataset_ids)
-    job_results = job_results_future.result()
-    job_results_object = JobResultObject(result=job_results)
+    job_results_object = create_job_result_object_task.submit(dataset_ids=dataset_ids)
 
-    access_token = get_scicat_access_token.submit(wait_for=[job_results_future])
+    access_token = get_scicat_access_token.submit(wait_for=[job_results_object])
 
     update_scicat_retrieval_job_status.submit(
         job_id=job_id,
         status_code=SciCatClient.JOBSTATUSCODE.FINISHED_SUCCESSFULLY,
         status_message=SciCatClient.JOBSTATUSMESSAGE.JOB_FINISHED,
-        jobResultObject=job_results_object,
+        jobResultObject=job_results_object.result(),
         token=access_token,
     ).wait()

backend/archiver/flows/tests/conftest.py

@@ -14,6 +14,7 @@ def aws_and_s3_credentials():
     os.environ["AWS_DEFAULT_REGION"] = "eu-west-1"
     os.environ["MINIO_REGION"] = "eu-west-1"
     os.environ["MINIO_ENDPOINT"] = "endpoint:9000"
+    os.environ["MINIO_EXTERNAL_ENDPOINT"] = "endpoint:9000"
 
 
 @pytest.fixture(scope="function")

backend/archiver/flows/tests/helpers.py

@@ -1,3 +1,4 @@
+import base64
 from typing import List, Dict, Any
 
 from pydantic import SecretStr
@@ -13,6 +14,7 @@
 from pathlib import Path
 
 from utils.s3_storage_interface import S3Storage
+from utils.script_generation import generate_download_script
 
 
 def mock_s3client() -> S3Storage:
@@ -113,7 +115,7 @@ def expected_datablocks(dataset_id: str, idx: int):
 
     return DataBlock(
         id=f"Block_{idx}",
-        archiveId=f"/path/to/archived/Block_{idx}.tar.gz",
+        archiveId=f"/path/to/archived/Block_{idx}.tar",
         size=size_per_file * 10,
         packedSize=size_per_file * 10,
         version=str(1),
@@ -133,7 +135,15 @@ def expected_jobresultsobject(dataset_id: str, datablocks: List[DataBlock]):
             )
         )
 
-    return JobResultObject(result=results).model_dump(exclude_none=True)
+    dataset_to_datablocks = {}
+
+    for result in results:
+        dataset_to_datablocks.setdefault(dataset_id, []).append({"name" : Path(result.archiveId).name, "url" : result.url})
+
+    script =  generate_download_script(dataset_to_datablocks)
+    script_b64 = base64.b64encode(bytes(script, 'utf-8'))
+
+    return JobResultObject(result=results, downloadScript=script_b64).model_dump(exclude_none=True)
 
 
 def mock_create_datablock_entries(
@@ -143,7 +153,7 @@ def mock_create_datablock_entries(
     for o in origDataBlocks:
         d = DataBlock(
             id=o.id,
-            archiveId=f"/path/to/archived/{o.id}.tar.gz",
+            archiveId=f"/path/to/archived/{o.id}.tar",
             size=o.size,
             packedSize=o.size,
             version=str(1),

backend/archiver/scicat/scicat_tasks.py

@@ -1,3 +1,4 @@
+import base64
 from typing import List
 from prefect import task
 from uuid import UUID
@@ -18,7 +19,8 @@
 from utils.s3_storage_interface import Bucket, S3Storage, get_s3_client
 
 
-from prefect.artifacts import create_link_artifact
+from prefect.artifacts import create_link_artifact, create_markdown_artifact
+from utils.script_generation import generate_download_script
 
 scicat_instance: SciCatClient | None = None
 
@@ -197,7 +199,7 @@ def reset_dataset(dataset_id: str, token: SecretStr):
 
 
 @task
-def create_job_result_object_task(dataset_ids: List[str]) -> List[JobResultEntry]:
+def create_job_result_object_task(dataset_ids: List[str]) -> JobResultObject:
     access_token = get_scicat_access_token.submit()
     access_token.wait()
 
@@ -209,28 +211,52 @@ def create_job_result_object_task(dataset_ids: List[str]) -> List[JobResultEntry
         datablocks_future.wait()
         datablocks = datablocks_future.result()
 
-        dataset_job_results = create_job_result_object(dataset_id, datablocks)
+        dataset_job_results = create_job_result_entries(dataset_id, datablocks)
         job_results = job_results + dataset_job_results
 
-    return job_results
+    job_results_object = JobResultObject(result=job_results)
 
+    script = create_download_script(job_results);
+    job_results_object.downloadScript = base64.b64encode(bytes(script, 'utf-8'))
+
+    markdown = f"""Download script for all datablocks in this job\n```bash\n{script}\n```\n"""
+
+    create_markdown_artifact(
+        key=f"script", markdown=markdown)
+
+    return job_results_object
+
+def create_download_script(job_result_entries: List[JobResultEntry]) -> str:
+
+    dataset_to_datablocks = {}
+
+    for result in job_result_entries:
+        dataset_to_datablocks.setdefault(result.datasetId, []).append({"name" : Path(result.archiveId).name, "url" : result.url})
+
+    return generate_download_script(dataset_to_datablocks)
+    
+    
 
 def create_presigned_url(client: S3Storage, datablock: DataBlock):
     url = client.get_presigned_url(Bucket.retrieval_bucket(), datablock.archiveId)
     return url
 
+def sanitize_name(name: str) -> str:
+    invalid_chars = ["/", ".", "_"]
+    sanitized_name = ""
+    for c in invalid_chars:
+        sanitized_name = name.replace(c, "-")
+    return sanitized_name
 
 @log
-def create_job_result_object(dataset_id: str, datablocks: List[DataBlock]) -> List[JobResultEntry]:
+def create_job_result_entries(dataset_id: str, datablocks: List[DataBlock]) -> List[JobResultEntry]:
     s3_client = get_s3_client()
     job_result_entries: List[JobResultEntry] = []
     for datablock in datablocks:
         url = create_presigned_url(s3_client, datablock)
 
-        invalid_chars = ["/", ".", "_"]
-        sanitized_name = str(Path(datablock.archiveId).name)
-        for c in invalid_chars:
-            sanitized_name = sanitized_name.replace(c, "-")
+        sanitized_name = sanitize_name(str(Path(datablock.archiveId).stem))
+
         create_link_artifact(
             key=sanitized_name,
             link=url,

backend/archiver/tests/test_e2e.py

@@ -341,7 +341,7 @@ async def test_end_to_end(scicat_token_setup, set_env, s3_client):
     # Verify retrieved datablock in MINIO
     retrieved_datablock = s3_client.stat_object(
         bucket=Bucket("retrieval"),
-        filename=f"openem-network/datasets/{dataset_pid}/datablocks/{dataset_pid}_0.tar.gz",
+        filename=f"openem-network/datasets/{dataset_pid}/datablocks/{dataset_pid}_0.tar",
     )
     assert retrieved_datablock is not None
     assert retrieved_datablock.Size > 80 * 1024 * 1024

backend/archiver/utils/datablocks.py

@@ -93,7 +93,7 @@ def create_tarfiles(
     target_size: int,
     progress_callback: Callable[[float], None] = None,
 ) -> List[ArchiveInfo]:
-    """Create datablocks, i.e. .tar.gz files, from all files in a folder. Folder structures are kept and symlnks not resolved.
+    """Create datablocks, i.e. .tar files, from all files in a folder. Folder structures are kept and symlnks not resolved.
     The created tar files will be named according to the dataset they belong to.
 
     Args:
@@ -120,7 +120,7 @@ def create_tar(idx: int, files: List) -> ArchiveInfo:
         current_tar_info = ArchiveInfo(
             unpackedSize=0,
             packedSize=0,
-            path=Path(dst_folder / Path(f"{tar_name}_{idx}.tar.gz")),
+            path=Path(dst_folder / Path(f"{tar_name}_{idx}.tar")),
             fileCount=len(files),
         )
         current_tarfile: tarfile.TarFile = tarfile.open(current_tar_info.path, "w")

backend/archiver/utils/model.py

@@ -20,6 +20,7 @@ class JobResultEntry(BaseModel):
 
 class JobResultObject(BaseModel):
     result: Optional[List[JobResultEntry]]
+    downloadScript: Optional[str] = None
 
 
 class Job(BaseModel):

backend/archiver/utils/s3_storage_interface.py

@@ -76,6 +76,15 @@ def __init__(self, url: str, user: str, password: SecretStr, region: str):
             region_name=self._REGION,
             config=Config(signature_version="s3v4", max_pool_connections=32),
         )
+        
+        self._external_minio = boto3.client(
+            "s3",
+            endpoint_url=f"https://{Variables().MINIO_EXTERNAL_ENDPOINT}",
+            aws_access_key_id=self._USER.strip(),
+            aws_secret_access_key=self._PASSWORD.get_secret_value().strip(),
+            region_name=self._REGION,
+            config=Config(signature_version="s3v4", max_pool_connections=32),
+        )
         self._resource = boto3.resource(
             "s3",
             endpoint_url=f"https://{self._URL}" if self._URL is not None and self._URL != "" else None,
@@ -97,7 +106,9 @@ def url(self):
     @log_debug
     def get_presigned_url(self, bucket: Bucket, filename: str) -> str:
         days_to_seconds = 60 * 60 * 24
-        presigned_url = self._minio.generate_presigned_url(
+
+
+        presigned_url = self._external_minio.generate_presigned_url(
             "get_object",
             Params={"Bucket": bucket.name, "Key": filename},
             ExpiresIn=Variables().MINIO_URL_EXPIRATION_DAYS