diff --git a/changelog/193.feature.rst b/changelog/193.feature.rst new file mode 100644 index 00000000..e8782ead --- /dev/null +++ b/changelog/193.feature.rst @@ -0,0 +1,3 @@ +Add a ``label=`` kwarg to `.FileManager.download` and +`.transfer_complete_datasets` allowing the user to completely customise the +Globus transfer task label. diff --git a/dkist/io/file_manager.py b/dkist/io/file_manager.py index f9c1d156..ed4aa831 100644 --- a/dkist/io/file_manager.py +++ b/dkist/io/file_manager.py @@ -313,7 +313,7 @@ def __init__(self, fits_loader: StripedExternalArray): # be populated with a reference to that Dataset instance. self._ndcube = None - def download(self, path=None, destination_endpoint=None, progress=True, wait=True): + def download(self, path=None, destination_endpoint=None, progress=True, wait=True, label=None): """ Start a Globus file transfer for all files in this Dataset. @@ -351,6 +351,9 @@ def download(self, path=None, destination_endpoint=None, progress=True, wait=Tru wait : `bool`, optional If `False` then the function will return while the Globus transfer task is still running. Setting ``wait=False`` implies ``progress=False``. + + label : `str` + Label for the Globus transfer. If `None` then a default will be used. """ if self._ndcube is None: raise ValueError( @@ -382,6 +385,7 @@ def download(self, path=None, destination_endpoint=None, progress=True, wait=Tru destination_endpoint=destination_endpoint, progress=progress, wait=wait, + label=label ) if is_local: diff --git a/dkist/io/tests/test_file_manager.py b/dkist/io/tests/test_file_manager.py index e6a6a4b9..dfb4244b 100644 --- a/dkist/io/tests/test_file_manager.py +++ b/dkist/io/tests/test_file_manager.py @@ -187,14 +187,16 @@ def test_download_default_keywords(dataset, orchestrate_transfer_mock): destination_endpoint=None, progress=True, wait=True, + label=None, ) @pytest.mark.parametrize("keywords", [ - {"progress": True, "wait": True, "destination_endpoint": None}, - {"progress": True, "wait": False, "destination_endpoint": None}, - {"progress": False, "wait": True, "destination_endpoint": None}, - {"progress": False, "wait": True, "destination_endpoint": "wibble"}, + {"progress": True, "wait": True, "destination_endpoint": None, "label": None}, + {"progress": True, "wait": False, "destination_endpoint": None, "label": None}, + {"progress": False, "wait": True, "destination_endpoint": None, "label": None}, + {"progress": False, "wait": True, "destination_endpoint": "wibble", "label": None}, + {"progress": False, "wait": True, "destination_endpoint": None, "label": "fibble"}, ]) def test_download_keywords(dataset, orchestrate_transfer_mock, keywords): """ @@ -238,6 +240,7 @@ def test_download_path_interpolation(dataset, orchestrate_transfer_mock): destination_endpoint=None, progress=True, wait=True, + label=None, ) assert dataset.files.basepath == Path("~/test_dataset").expanduser() diff --git a/dkist/net/globus/tests/test_transfer.py b/dkist/net/globus/tests/test_transfer.py index c69f361b..93b30c6e 100644 --- a/dkist/net/globus/tests/test_transfer.py +++ b/dkist/net/globus/tests/test_transfer.py @@ -195,7 +195,7 @@ def test_orchestrate_transfer(tfr_file_list, orchestrate_mocks): Path("/~/"), tfr_file_list, recursive=False, - label_suffix=None, + label=None, ) @@ -210,7 +210,7 @@ def test_orchestrate_transfer_no_progress(tfr_file_list, mocker, orchestrate_moc Path("/~/"), tfr_file_list, recursive=False, - label_suffix=None, + label=None, ) orchestrate_mocks.watch_transfer_progress.assert_not_called() @@ -231,7 +231,7 @@ def test_orchestrate_transfer_no_wait(tfr_file_list, mocker, orchestrate_mocks): Path("/~/"), tfr_file_list, recursive=False, - label_suffix=None, + label=None, ) orchestrate_mocks.watch_transfer_progress.assert_not_called() diff --git a/dkist/net/globus/transfer.py b/dkist/net/globus/transfer.py index f67cc396..8a1475d5 100644 --- a/dkist/net/globus/transfer.py +++ b/dkist/net/globus/transfer.py @@ -20,7 +20,7 @@ def start_transfer_from_file_list(src_endpoint, dst_endpoint, dst_base_path, file_list, - src_base_path=None, recursive=False, label_suffix=None): + src_base_path=None, recursive=False, label=None): """ Start a new transfer task for a list of files. @@ -55,8 +55,8 @@ def start_transfer_from_file_list(src_endpoint, dst_endpoint, dst_base_path, fil If you need to set this per-item in ``file_list`` it should be a `list` of `bool` of equal length as ``file_list``. - label_suffix : `str` - String to append to the transfer task label. + label : `str` + Label for the Globus transfer. If None then a default will be used. Returns ------- @@ -80,10 +80,10 @@ def start_transfer_from_file_list(src_endpoint, dst_endpoint, dst_base_path, fil dst_endpoint = get_endpoint_id(dst_endpoint, tc) auto_activate_endpoint(dst_endpoint, tc) - now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S") - label_suffix = f" {label_suffix}" if label_suffix is not None else "" + now = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M") + label = f"DKIST Python Tools - {now}" if label is None else label transfer_manifest = globus_sdk.TransferData(tc, src_endpoint, dst_endpoint, - label=f"DKIST Python Tools - {now}{label_suffix}", + label=label, sync_level="checksum", verify_checksum=True) @@ -271,7 +271,7 @@ def _orchestrate_transfer_task(file_list: List[PathLike], *, progress: Union[bool, Literal["verbose"]] = True, wait: bool = True, - label_suffix=None): + label=None): """ Transfer the files given in file_list to the path on ``destination_endpoint``. @@ -306,8 +306,8 @@ def _orchestrate_transfer_task(file_list: List[PathLike], If `False` then the function will return while the Globus transfer task is still running. Setting ``wait=False`` implies ``progress=False``. - label_suffix : `str` - String to append to the transfer task label. + label : `str` + Label for the Globus transfer. If `None` then a default will be used. Returns ------- @@ -325,7 +325,7 @@ def _orchestrate_transfer_task(file_list: List[PathLike], destination_path, file_list, recursive=recursive, - label_suffix=label_suffix) + label=label) tc = get_transfer_client() diff --git a/dkist/net/helpers.py b/dkist/net/helpers.py index 716f50a5..ad063977 100644 --- a/dkist/net/helpers.py +++ b/dkist/net/helpers.py @@ -1,8 +1,9 @@ """ Functions and classes for searching and downloading from the data center. """ +import datetime from os import PathLike -from typing import List, Union, Literal +from typing import List, Union, Literal, Optional from pathlib import Path from astropy import table @@ -27,7 +28,8 @@ def transfer_complete_datasets(datasets: Union[str, QueryResponseRow, DKISTQuery path: PathLike = "/~/", destination_endpoint: str = None, progress: Union[bool, Literal["verbose"]] = True, - wait: bool = True) -> Union[List[str], str]: + wait: bool = True, + label: Optional[str] = None) -> Union[List[str], str]: """ Transfer one or more complete datasets to a path on a globus endpoint. @@ -56,6 +58,9 @@ def transfer_complete_datasets(datasets: Union[str, QueryResponseRow, DKISTQuery function. To run multiple dataset transfer tasks in parallel (one task per dataset) specify ``wait=False``. + label : `str` + Label for the Globus transfer. If `None` then a default will be used. + Returns ------- The path to the directories containing the dataset(s) on the destination endpoint. @@ -81,7 +86,8 @@ def transfer_complete_datasets(datasets: Union[str, QueryResponseRow, DKISTQuery path=path, destination_endpoint=destination_endpoint, progress=progress, - wait=wait)) + wait=wait, + label=label)) return paths # At this point we only have one dataset @@ -98,12 +104,15 @@ def transfer_complete_datasets(datasets: Union[str, QueryResponseRow, DKISTQuery bucket=bucket ))] + now = datetime.datetime.now().strftime("%Y-%m-%dT%H:%M") + label = f"DKIST Python Tools - {now} {dataset_id}" if label is None else label + _orchestrate_transfer_task(file_list, recursive=True, destination_path=destination_path, destination_endpoint=destination_endpoint, progress=progress, wait=wait, - label_suffix=dataset_id) + label=label) return destination_path / dataset_id diff --git a/dkist/net/tests/test_helpers.py b/dkist/net/tests/test_helpers.py index b73e0737..d82bd84a 100644 --- a/dkist/net/tests/test_helpers.py +++ b/dkist/net/tests/test_helpers.py @@ -1,3 +1,4 @@ +import datetime from pathlib import Path import pytest @@ -16,10 +17,11 @@ def orchestrate_transfer_mock(mocker): @pytest.mark.parametrize( "keywords", [ - {"progress": True, "wait": True, "destination_endpoint": None}, - {"progress": True, "wait": False, "destination_endpoint": None}, - {"progress": False, "wait": True, "destination_endpoint": None}, - {"progress": False, "wait": True, "destination_endpoint": "wibble"}, + {"progress": True, "wait": True, "destination_endpoint": None, "label": None}, + {"progress": True, "wait": False, "destination_endpoint": None, "label": None}, + {"progress": False, "wait": True, "destination_endpoint": None, "label": None}, + {"progress": False, "wait": True, "destination_endpoint": "wibble", "label": None}, + {"progress": False, "wait": True, "destination_endpoint": None, "label": "fibble"}, ], ) def test_download_default_keywords(orchestrate_transfer_mock, keywords): @@ -32,11 +34,12 @@ def test_download_default_keywords(orchestrate_transfer_mock, keywords): **keywords ) + if keywords["label"] is None: + keywords["label"] = f"DKIST Python Tools - {datetime.datetime.now().strftime('%Y-%m-%dT%H:%M')} AAAA" orchestrate_transfer_mock.assert_called_once_with( [Path("/data/pm_1_10/AAAA")], recursive=True, destination_path=Path("/~/pm_1_10"), - label_suffix="AAAA", **keywords ) @@ -63,7 +66,7 @@ def test_transfer_from_dataset_id(mocker, orchestrate_transfer_mock): destination_endpoint=None, progress=True, wait=True, - label_suffix="AAAA", + label=f"DKIST Python Tools - {datetime.datetime.now().strftime('%Y-%m-%dT%H:%M')} AAAA" ) get_inv_mock.assert_called_once_with("AAAA") @@ -78,23 +81,21 @@ def test_transfer_from_table(orchestrate_transfer_mock, mocker): }, ) - transfer_complete_datasets(res) + transfer_complete_datasets(res, label="fibble") - kwargs = {"progress": True, "wait": True, "destination_endpoint": None} + kwargs = {"progress": True, "wait": True, "destination_endpoint": None, "label": "fibble"} orchestrate_transfer_mock.assert_has_calls( [ mocker.call( [Path("/data/pm_1_10/A")], recursive=True, destination_path=Path("/~/pm_1_10"), - label_suffix="A", **kwargs ), mocker.call( [Path("/data/pm_2_20/B")], recursive=True, destination_path=Path("/~/pm_2_20"), - label_suffix="B", **kwargs ), ] @@ -120,23 +121,21 @@ def test_transfer_from_UnifiedResponse(orchestrate_transfer_mock, mocker): ) res._list[0].client = res._list[1].client = DKISTClient() - transfer_complete_datasets(res) + transfer_complete_datasets(res, label="fibble") - kwargs = {"progress": True, "wait": True, "destination_endpoint": None} + kwargs = {"progress": True, "wait": True, "destination_endpoint": None, "label": "fibble"} orchestrate_transfer_mock.assert_has_calls( [ mocker.call( [Path("/data/pm_1_10/A")], recursive=True, destination_path=Path("/~/pm_1_10"), - label_suffix="A", **kwargs ), mocker.call( [Path("/data/pm_2_20/B")], recursive=True, destination_path=Path("/~/pm_2_20"), - label_suffix="B", **kwargs ), ]