[Feature]: Better handle the Globus token file

tests/integration/bash_tests/run_from_any/globus_auth.bash

@@ -16,7 +16,7 @@ check_log_does_not_have()
     local log_file="${2}"
     grep "${not_expected_grep}" ${log_file}
     if [ $? == 0 ]; then
-        echo "Not-expected grep '${expected_grep}' was found in ${log_file}. Test failed."
+        echo "Not-expected grep '${not_expected_grep}' was found in ${log_file}. Test failed."
         exit 2
     fi
 }
@@ -127,7 +127,8 @@ test_single_auth_code()
     check_log_has "INFO: Writing to empty ${INI_PATH}" ${case_name}.log
     check_log_has "INFO: Setting local_endpoint_id based on" ${case_name}.log
     # From get_transfer_client_with_auth
-    check_log_has "INFO: No stored tokens found - starting authentication" ${case_name}.log
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
     check_log_has "Please go to this URL and login:" ${case_name}.log # Our one expected authentication prompt
     # From save_tokens
     check_log_has "INFO: Tokens saved successfully" ${case_name}.log
@@ -150,7 +151,8 @@ test_single_auth_code()
     check_log_has "INFO: Setting local_endpoint_id based on ${INI_PATH}" ${case_name}.log # Differs from run1
     check_log_has "INFO: Setting local_endpoint_id based on" ${case_name}.log
     # From get_transfer_client_with_auth
-    check_log_has "INFO: Found stored refresh token - using it" ${case_name}.log # Differs from run1
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}.log # Updated to match new log message
+    check_log_has "using it" ${case_name}.log
     check_log_does_not_have "Please go to this URL and login:" ${case_name}.log # There should be no login prompts for run2!
     # From save_tokens
     check_log_does_not_have "INFO: Tokens saved successfully" ${case_name}.log # Differs from run1
@@ -191,16 +193,25 @@ test_different_endpoint1()
 
     case_name="different_endpoint1"
     setup ${case_name} "${src_dir}"
-    # Expecting to see exactly 1 authentication prompt
+    # With multi-token support, this should now work and prompt for auth
+    # Expecting to see exactly 1 authentication prompt for the new endpoint
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
-    check_log_has "INFO: Found stored refresh token - using it" ${case_name}.log
-    check_log_has "ERROR: One possible cause" ${case_name}.log
-    check_log_has "ERROR: Try deleting" ${case_name}.log
-    check_log_has "ERROR: Another possible cause" ${case_name}.log
-    check_log_has "try revoking consents before re-running" ${case_name}.log
-    check_log_has "ERROR: Exception: Insufficient Globus consents" ${case_name}.log
-
-    if ! confirm "Did you avoid having to paste any auth codes on this run?"; then
+    if [ $? != 0 ]; then
+        echo "${case_name} failed. Check ${case_name}.log for details."
+        exit 1
+    fi
+
+    # Should see token not found for this endpoint pair, then authentication
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+    check_log_has "INFO: Tokens saved successfully" ${case_name}.log
+
+    # Should NOT see any errors about insufficient consents
+    check_log_does_not_have "ERROR: One possible cause" ${case_name}.log
+    check_log_does_not_have "ERROR: Insufficient Globus consents" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code for this new endpoint?"; then
         echo "test_different_endpoint1 failed"
         exit 1
     fi
@@ -220,34 +231,29 @@ test_different_endpoint2()
     dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
     globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
 
-    echo "Reset Globus consents:"
-    echo "https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring > rescind all"
-    if ! confirm "Have you revoked Globus consents?"; then
-        exit 1
-    fi
-
     echo "Running test_different_endpoint2"
+    echo "This test verifies that we can switch back to a previously-used endpoint without re-authenticating"
     echo "Exit codes: 0 -- success, 1 -- failure"
 
-    case_name="different_endpoint2a"
-    setup ${case_name} "${src_dir}"
-    zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
-    check_log_has ".zstash_globus_tokens.json exists. We can try to load tokens from it." ${case_name}.log
-    check_log_has ".zstash_globus_tokens.json may be configured for a different Globus endpoint." ${case_name}.log
-    check_log_has "Try deleting" ${case_name}.log
-    check_log_has "globus_sdk.services.auth.errors.AuthAPIError: ('POST', 'https://auth.globus.org/v2/oauth2/token', None, 400, 'Error', 'Bad Request')" ${case_name}.log
-
-    rm -rf ~/.zstash_globus_tokens.json
-    case_name="different_endpoint2b"
+    case_name="different_endpoint2"
     setup ${case_name} "${src_dir}"
-    # Expecting to see exactly 1 authentication prompt
+    # This endpoint was already authenticated in test_different_endpoint1
+    # We should be able to use the stored token without prompting
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
     if [ $? != 0 ]; then
-        echo "${case_name} failed. Check ${case_name}_create.log for details."
+        echo "${case_name} failed. Check ${case_name}.log for details."
         exit 1
     fi
 
-    if ! confirm "Did you only have to paste an auth code once (for 2b, not 2a)?"; then
+    # Should find the stored token for this endpoint pair
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}.log
+    check_log_has "using it" ${case_name}.log
+
+    # Should NOT see authentication prompt
+    check_log_does_not_have "Please go to this URL and login:" ${case_name}.log
+    check_log_does_not_have "starting authentication" ${case_name}.log
+
+    if ! confirm "Did you NOT have to paste any auth codes for this run?"; then
         echo "test_different_endpoint2 failed"
         exit 1
     fi
@@ -268,27 +274,164 @@ test_different_endpoint3()
     globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
 
     echo "Running test_different_endpoint3"
+    echo "This test verifies a third different endpoint (also requires auth)"
     echo "Exit codes: 0 -- success, 1 -- failure"
 
-    rm -rf ~/.zstash_globus_tokens.json
     case_name="different_endpoint3"
     setup ${case_name} "${src_dir}"
+    # This is a third endpoint that hasn't been authenticated yet
     # Expecting to see exactly 1 authentication prompt
     zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
     if [ $? != 0 ]; then
-        echo "${case_name} failed. Check ${case_name}_create.log for details."
+        echo "${case_name} failed. Check ${case_name}.log for details."
         exit 1
     fi
 
-    if ! confirm "Did you only have to paste an auth code once?"; then
-        echo "test_different_endpoint2 failed"
+    # Should see no stored token, then authentication
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+    check_log_has "INFO: Tokens saved successfully" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code for this third endpoint?"; then
+        echo "test_different_endpoint3 failed"
         exit 1
     fi
     # Cleanup:
     cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
     rm -rf ${path_to_repo}/tests/utils/globus_auth
 }
 
+test_legacy_token_migration()
+{
+    local path_to_repo=$1
+    local dst_endpoint=$2
+    local dst_dir=$3
+
+    src_dir=${path_to_repo}/tests/utils/globus_auth
+    mkdir -p ${src_dir}
+    dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
+    globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
+
+    TOKEN_FILE=${HOME}/.zstash_globus_tokens.json
+
+    echo "Running test_legacy_token_migration"
+    echo "This test verifies that legacy single-token format is handled gracefully"
+    echo "Exit codes: 0 -- success, 1 -- failure"
+
+    # Create a mock legacy token file (just the structure, not valid tokens)
+    cat > ${TOKEN_FILE} << 'EOF'
+{
+  "transfer.api.globus.org": {
+    "access_token": "fake_legacy_token",
+    "refresh_token": "fake_legacy_refresh",
+    "expires_at": 1234567890
+  }
+}
+EOF
+
+    case_name="legacy_migration"
+    setup ${case_name} "${src_dir}"
+
+    # Should detect legacy format and require re-authentication
+    zstash create --hpss=${globus_path}/${case_name} zstash_demo 2>&1 | tee ${case_name}.log
+    if [ $? != 0 ]; then
+        echo "${case_name} failed. Check ${case_name}.log for details."
+        exit 1
+    fi
+
+    # Should see migration message and re-authentication
+    check_log_has "INFO: Detected legacy token format, migrating to new format" ${case_name}.log
+    check_log_has "INFO: No stored tokens found for endpoints" ${case_name}.log
+    check_log_has "starting authentication" ${case_name}.log
+    check_log_has "Please go to this URL and login:" ${case_name}.log
+
+    if ! confirm "Did you have to paste an auth code to migrate from legacy format?"; then
+        echo "test_legacy_token_migration failed"
+        exit 1
+    fi
+
+    # Cleanup:
+    cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
+    rm -rf ${path_to_repo}/tests/utils/globus_auth
+}
+
+test_custom_token_file()
+{
+    local path_to_repo=$1
+    local dst_endpoint=$2
+    local dst_dir=$3
+
+    src_dir=${path_to_repo}/tests/utils/globus_auth
+    mkdir -p ${src_dir}
+    dst_endpoint_uuid=$(get_endpoint ${dst_endpoint})
+    globus_path=globus://${dst_endpoint_uuid}/${dst_dir}
+
+    echo "Running test_custom_token_file"
+    echo "This test verifies that --globus-token-file parameter works"
+    echo "Exit codes: 0 -- success, 1 -- failure"
+
+    # Create a temporary directory for test tokens
+    TEMP_TOKEN_DIR=$(mktemp -d)
+    CUSTOM_TOKEN_FILE="${TEMP_TOKEN_DIR}/custom_tokens.json"
+
+    case_name="custom_token_file"
+    setup ${case_name} "${src_dir}"
+
+    # First run - should authenticate and save to custom location
+    zstash create --hpss=${globus_path}/${case_name}_run1 --globus-token-file="${CUSTOM_TOKEN_FILE}" zstash_demo 2>&1 | tee ${case_name}_run1.log
+    if [ $? != 0 ]; then
+        echo "${case_name} run1 failed. Check ${case_name}_run1.log for details."
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Verify custom token file was created
+    if [ ! -f "${CUSTOM_TOKEN_FILE}" ]; then
+        echo "ERROR: Custom token file ${CUSTOM_TOKEN_FILE} was not created!"
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Verify token was saved to custom location and custom path was used
+    check_log_has "INFO: Using custom token file: ${CUSTOM_TOKEN_FILE}" ${case_name}_run1.log
+    check_log_has "INFO: Tokens saved successfully to ${CUSTOM_TOKEN_FILE}" ${case_name}_run1.log
+    check_log_has "INFO: Token file ${CUSTOM_TOKEN_FILE} exists" ${case_name}_run1.log
+
+    # Second run - should use token from custom location
+    zstash create --hpss=${globus_path}/${case_name}_run2 --globus-token-file="${CUSTOM_TOKEN_FILE}" zstash_demo 2>&1 | tee ${case_name}_run2.log
+    if [ $? != 0 ]; then
+        echo "${case_name} run2 failed. Check ${case_name}_run2.log for details."
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Should have used the stored token
+    check_log_has "INFO: Using custom token file: ${CUSTOM_TOKEN_FILE}" ${case_name}_run2.log
+    check_log_has "INFO: Found stored refresh token for endpoints" ${case_name}_run2.log
+    check_log_does_not_have "Please go to this URL and login:" ${case_name}_run2.log
+
+    # Verify default token file was NOT created
+    DEFAULT_TOKEN_FILE=${HOME}/.zstash_globus_tokens.json
+    if [ -f "${DEFAULT_TOKEN_FILE}" ]; then
+        echo "WARNING: Default token file ${DEFAULT_TOKEN_FILE} exists when it shouldn't"
+        echo "This suggests the custom token file setting may not be working correctly"
+    fi
+
+    if ! confirm "Did you only have to paste an auth code once (for run1, not run2)?"; then
+        echo "test_custom_token_file failed"
+        rm -rf ${TEMP_TOKEN_DIR}
+        exit 1
+    fi
+
+    # Cleanup
+    rm -rf ${TEMP_TOKEN_DIR}
+    cd ${path_to_repo}/tests/integration/bash_tests/run_from_any
+    rm -rf ${path_to_repo}/tests/utils/globus_auth
+
+    echo "test_custom_token_file completed successfully"
+}
+
 # Follow these directions #####################################################
 
 # Example usage:
@@ -371,16 +514,21 @@ test_single_auth_code ${path_to_repo} NERSC_HPSS_ENDPOINT ${hpss_dst_dir}
 echo "Testing transfer to pic#compy-dtn ######################################"
 test_single_auth_code ${path_to_repo} PIC_COMPY_DTN_ENDPOINT ${compy_dst_dir}
 
-echo "Follow-up tests: behavior when switching to different endpoints"
-echo "NOTE: if you commented out tests above, and your last endpoint used was NERSC_PERLMUTTER_ENDPOINT, the following test will not work properly."
-echo "Test 1: What if we switch to a different endpoint? #####################"
+echo "Follow-up tests: multi-endpoint token storage behavior"
+echo "Test 1: Switch to a different endpoint (should prompt for new auth) ####"
 test_different_endpoint1 ${path_to_repo} ${dst_endpoint_switch1} ${dst_dir_switch1}
-echo "Test 2: What if we try a) revoking consents and then b) removing the token file? ###"
+echo "Test 2: Switch back to a previous endpoint (should use stored token) ###"
 test_different_endpoint2 ${path_to_repo} ${dst_endpoint_switch1} ${dst_dir_switch1}
-echo "Test 3: What if we switch to a different endpoint again, but first remove the token file? ###"
+echo "Test 3: Switch to a third different endpoint (should prompt for auth) ##"
 test_different_endpoint3 ${path_to_repo} ${dst_endpoint_switch2} ${dst_dir_switch2}
-echo "Check https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring: you should have *two* consents there now."
-if ! confirm "Does https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring show *two* consents?"; then
+echo "Test 4: Verify legacy token format migration ###########################"
+test_legacy_token_migration ${path_to_repo} LCRC_IMPROV_DTN_ENDPOINT ${chrysalis_dst_dir}
+echo "Test 5: Verify custom token file via environment variable ##############"
+test_custom_token_file ${path_to_repo} LCRC_IMPROV_DTN_ENDPOINT ${chrysalis_dst_dir}
+
+echo "Check https://auth.globus.org/v2/web/consents > Globus Endpoint Performance Monitoring"
+echo "You should now have multiple consents (one for each endpoint pair authenticated)."
+if ! confirm "Does https://auth.globus.org/v2/web/consents show multiple consents as expected?"; then
     exit 1
 fi
 echo "All globus_auth tests completed successfully."

zstash/create.py

@@ -57,7 +57,7 @@ def create():
         if url.scheme == "globus":
             # identify globus endpoints
             logger.debug(f"{ts_utc()}:Calling globus_activate(hpss)")
-            globus_activate(hpss)
+            globus_activate(hpss, args.globus_token_file)
         else:
             # config.hpss is not "none", so we need to
             # create target HPSS directory
@@ -176,6 +176,11 @@ def setup_create() -> Tuple[str, argparse.Namespace]:
         action="store_true",
         help="FOR ADVANCED USERS ONLY: If a duplicate tar is encountered, overwrite the existing database record with the new one (i.e., it will assume the latest tar is the correct one). If this flag is not set, zstash will permit multiple entries for the same tar in its database.",
     )
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "--for-developers-force-database-corruption",
         type=str,

zstash/extract.py

@@ -19,6 +19,7 @@
 import _io
 
 from . import parallel
+from .globus import globus_activate
 from .hpss import hpss_get
 from .settings import (
     BLOCK_SIZE,
@@ -98,6 +99,11 @@ def setup_extract() -> Tuple[argparse.Namespace, str]:
         "--retries", type=int, default=1, help="number of times to retry an hsi command"
     )
     optional.add_argument("--tars", type=str, help="specify which tars to process")
+    optional.add_argument(
+        "--globus-token-file",
+        type=str,
+        help="Path to custom Globus token file. If not specified, uses ~/.zstash_globus_tokens.json",
+    )
     optional.add_argument(
         "-v", "--verbose", action="store_true", help="increase output verbosity"
     )
@@ -175,6 +181,7 @@ def extract_database(
                 hpss: str = config.hpss
             else:
                 raise TypeError("Invalid config.hpss={}".format(config.hpss))
+            globus_activate(hpss, args.globus_token_file)
             hpss_get(hpss, get_db_filename(cache), cache)
         else:
             error_str: str = (

zstash/globus.py

@@ -13,6 +13,7 @@
     check_state_files,
     get_local_endpoint_id,
     get_transfer_client_with_auth,
+    set_token_file_path,
     set_up_TransferData,
     submit_transfer_with_checks,
 )
@@ -27,7 +28,7 @@
 archive_directory_listing: IterableTransferResponse = None
 
 
-def globus_activate(hpss: str):
+def globus_activate(hpss: str, token_file: Optional[str] = None):
     """
     Read the local globus endpoint UUID from ~/.zstash.ini.
     If the ini file does not exist, create an ini file with empty values,
@@ -40,6 +41,11 @@ def globus_activate(hpss: str):
     url = urlparse(hpss)
     if url.scheme != "globus":
         return
+
+    # Set the token file path if provided
+    if token_file:
+        set_token_file_path(token_file)
+
     check_state_files()
     remote_endpoint = url.netloc
     local_endpoint = get_local_endpoint_id(local_endpoint)