Skip to content

Commit

Permalink
Exposure Globus "label" directly to all top-level download functions (#…
Browse files Browse the repository at this point in the history
…193)

* Exposure Globus "label" directly to all top-level download functions

This resulted in removing the label_suffix from some intermediate functions. I
think this way is better because it gives total control to the user. The default
values have only been changed very slightly.

* Apply suggestions from code review

* Add changelog

Co-authored-by: Stuart Mumford <stuart@cadair.com>
  • Loading branch information
eigenbrot and Cadair committed Aug 4, 2022
1 parent a75bf4e commit 6d9d91d
Show file tree
Hide file tree
Showing 7 changed files with 54 additions and 36 deletions.
3 changes: 3 additions & 0 deletions changelog/193.feature.rst
Original file line number Diff line number Diff line change
@@ -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.
6 changes: 5 additions & 1 deletion dkist/io/file_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
11 changes: 7 additions & 4 deletions dkist/io/tests/test_file_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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()
6 changes: 3 additions & 3 deletions dkist/net/globus/tests/test_transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ def test_orchestrate_transfer(tfr_file_list, orchestrate_mocks):
Path("/~/"),
tfr_file_list,
recursive=False,
label_suffix=None,
label=None,
)


Expand All @@ -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()
Expand All @@ -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()
Expand Down
20 changes: 10 additions & 10 deletions dkist/net/globus/transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
-------
Expand All @@ -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)

Expand Down Expand Up @@ -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``.
Expand Down Expand Up @@ -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
-------
Expand All @@ -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()

Expand Down
17 changes: 13 additions & 4 deletions dkist/net/helpers.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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
Expand All @@ -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
27 changes: 13 additions & 14 deletions dkist/net/tests/test_helpers.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import datetime
from pathlib import Path

import pytest
Expand All @@ -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):
Expand All @@ -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
)

Expand All @@ -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")
Expand All @@ -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
),
]
Expand All @@ -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
),
]
Expand Down

0 comments on commit 6d9d91d

Please sign in to comment.