Improve handling of TransferData

Author: forsyth2

zstash/globus.py

@@ -10,10 +10,12 @@
 
 from .globus_utils import (
     HPSS_ENDPOINT_MAP,
+    add_file_to_TransferData,
     check_state_files,
+    create_TransferData,
+    get_label,
     get_local_endpoint_id,
     get_transfer_client_with_auth,
-    set_up_TransferData,
     submit_transfer_with_checks,
 )
 from .settings import logger
@@ -122,14 +124,38 @@ def globus_transfer(  # noqa: C901
         raise RuntimeError(
             "The transfer manager should always have at least one batch by the time globus_transfer is called, however, the batch list is empty."
         )
-    transfer_data = set_up_TransferData(
+
+    if transfer_manager.globus_config.local_endpoint:
+        local_endpoint: str = transfer_manager.globus_config.local_endpoint
+    else:
+        raise ValueError("Local endpoint ID is not set.")
+    if transfer_manager.globus_config.remote_endpoint:
+        remote_endpoint: str = transfer_manager.globus_config.remote_endpoint
+    else:
+        raise ValueError("Remote endpoint ID is not set.")
+    label = get_label(remote_path, name)
+
+    transfer_data: TransferData
+    if mrb.transfer_data:
+        # We already have a TransferData for this batch.
+        transfer_data = mrb.transfer_data
+    else:
+        # We need to create a new TransferData for this batch.
+        transfer_data = create_TransferData(
+            transfer_type,
+            local_endpoint,
+            remote_endpoint,
+            transfer_manager.globus_config.transfer_client,
+            label,
+        )
+    add_file_to_TransferData(
         transfer_type,
-        transfer_manager.globus_config.local_endpoint,
-        transfer_manager.globus_config.remote_endpoint,
+        local_endpoint,
+        remote_endpoint,
         remote_path,
         name,
-        transfer_manager.globus_config.transfer_client,
-        mrb.transfer_data,
+        transfer_data,
+        label,
     )
 
     task: GlobusHTTPResponse
@@ -215,9 +241,6 @@ def globus_transfer(  # noqa: C901
             error_str = "transfer_manager has no batches"
             logger.error(error_str)
             raise RuntimeError(error_str)
-
-        # Nullify the submitted transfer data structure so that a new one will be created on next call.
-        transfer_data = None
     except TransferAPIError as e:
         if e.code == "NoCredException":
             logger.error(

zstash/globus_utils.py

@@ -217,15 +217,48 @@ def save_tokens(token_response):
 
 
 # Primarily used by globus_transfer ###########################################
-def set_up_TransferData(
+def get_label(remote_path: str, name: str) -> str:
+    subdir = os.path.basename(os.path.normpath(remote_path))
+    subdir_label = re.sub("[^A-Za-z0-9_ -]", "", subdir)
+    filename = name.split(".")[0]
+    label = subdir_label + " " + filename
+    return label
+
+
+def create_TransferData(
+    transfer_type: str,
+    local_endpoint: str,
+    remote_endpoint: str,
+    transfer_client: TransferClient,
+    label: str,
+) -> TransferData:
+    if transfer_type == "get":
+        src_ep = remote_endpoint
+        dst_ep = local_endpoint
+    else:
+        src_ep = local_endpoint
+        dst_ep = remote_endpoint
+    transfer_data = TransferData(
+        transfer_client,
+        src_ep,
+        dst_ep,
+        label=label,
+        verify_checksum=True,
+        preserve_timestamp=True,
+        fail_on_quota_errors=True,
+    )
+    return transfer_data
+
+
+def add_file_to_TransferData(
     transfer_type: str,
     local_endpoint: Optional[str],
     remote_endpoint: Optional[str],
     remote_path: str,
     name: str,
-    transfer_client: TransferClient,
-    transfer_data: Optional[TransferData] = None,
-) -> TransferData:
+    transfer_data: TransferData,
+    label: str,
+):
     """
     Set up the TransferData object, creating one if not provided.
     """
@@ -234,34 +267,14 @@ def set_up_TransferData(
     if not remote_endpoint:
         raise ValueError("Remote endpoint ID is not set.")
     if transfer_type == "get":
-        src_ep = remote_endpoint
         src_path = os.path.join(remote_path, name)
-        dst_ep = local_endpoint
         dst_path = os.path.join(os.getcwd(), name)
     else:
-        src_ep = local_endpoint
         src_path = os.path.join(os.getcwd(), name)
-        dst_ep = remote_endpoint
         dst_path = os.path.join(remote_path, name)
 
-    subdir = os.path.basename(os.path.normpath(remote_path))
-    subdir_label = re.sub("[^A-Za-z0-9_ -]", "", subdir)
-    filename = name.split(".")[0]
-    label = subdir_label + " " + filename
-
-    if not transfer_data:
-        transfer_data = TransferData(
-            transfer_client,
-            src_ep,
-            dst_ep,
-            label=label,
-            verify_checksum=True,
-            preserve_timestamp=True,
-            fail_on_quota_errors=True,
-        )
     transfer_data.add_item(src_path, dst_path)
     transfer_data["label"] = label
-    return transfer_data
 
 
 def submit_transfer_with_checks(transfer_client, transfer_data) -> GlobusHTTPResponse: