Store state with object

Author: forsyth2

zstash/create.py

@@ -6,15 +6,17 @@
 import os.path
 import sqlite3
 import sys
-from typing import Any, List, Tuple
+from typing import Any, List
 
 from six.moves.urllib.parse import urlparse
 
 from .globus import globus_activate, globus_finalize
 from .hpss import hpss_put
 from .hpss_utils import add_files
-from .settings import DEFAULT_CACHE, config, get_db_filename, logger
+from .settings import logger
 from .utils import (
+    CommandInfo,
+    HPSSType,
     create_tars_table,
     get_files_to_archive,
     run_command,
@@ -24,60 +26,46 @@
 
 
 def create():
-    cache: str
-    cache, args = setup_create()
-
-    # Check config fields
-    if config.path is not None:
-        path: str = config.path
-    else:
-        raise TypeError("Invalid config.path={}".format(config.path))
-    if config.hpss is not None:
-        hpss: str = config.hpss
-    else:
-        raise TypeError("Invalid config.hpss={}".format(config.hpss))
+    command_info = CommandInfo("create")
+    args = setup_create(command_info)
 
     # Start doing actual work
     logger.debug(f"{ts_utc()}: Running zstash create")
-    logger.debug("Local path : {}".format(path))
-    logger.debug("HPSS path  : {}".format(hpss))
-    logger.debug("Max size  : {}".format(config.maxsize))
-    logger.debug("Keep local tar files  : {}".format(args.keep))
+    logger.debug(f"Local path: {command_info.dir_to_archive_absolute}")
+    logger.debug(f"HPSS path: {command_info.hpss_path}")
+    logger.debug(f"Max size: {command_info.maxsize}")
+    logger.debug(f"Keep local tar files: {command_info.keep}")
 
     # Make sure input path exists and is a directory
     logger.debug("Making sure input path exists and is a directory")
-    if not os.path.isdir(path):
+    if not os.path.isdir(command_info.dir_to_archive_absolute):
         # Input path is not a directory
-        input_path_error_str: str = "Input path should be a directory: {}".format(path)
+        input_path_error_str: str = f"Input path should be a directory: {command_info.dir_to_archive_absolute}"
         logger.error(input_path_error_str)
         raise NotADirectoryError(input_path_error_str)
 
-    if hpss != "none":
-        url = urlparse(hpss)
-        if url.scheme == "globus":
-            # identify globus endpoints
-            logger.debug(f"{ts_utc()}:Calling globus_activate(hpss)")
-            globus_activate(hpss)
-        else:
-            # config.hpss is not "none", so we need to
-            # create target HPSS directory
-            logger.debug(f"{ts_utc()}: Creating target HPSS directory {hpss}")
-            mkdir_command: str = "hsi -q mkdir -p {}".format(hpss)
-            mkdir_error_str: str = "Could not create HPSS directory: {}".format(hpss)
-            run_command(mkdir_command, mkdir_error_str)
+    if command_info.hpss_type == HPSSType.GLOBUS:
+        # identify globus endpoints
+        logger.debug(f"{ts_utc()}: Calling globus_activate")
+        globus_activate(command_info.globus_info)
+    elif command_info.hpss_type == HPSSType.SAME_MACHINE_HPSS:
+        logger.debug(f"{ts_utc()}: Creating target HPSS directory {command_info.hpss_path}")
+        mkdir_command: str = f"hsi -q mkdir -p {command_info.hpss_path}"
+        mkdir_error_str: str = f"Could not create HPSS directory: {command_info.hpss_path}"
+        run_command(mkdir_command, mkdir_error_str)
 
-            # Make sure it is exists and is empty
-            logger.debug("Making sure target HPSS directory exists and is empty")
+        # Make sure it is exists and is empty
+        logger.debug("Making sure target HPSS directory exists and is empty")
 
-            ls_command: str = 'hsi -q "cd {}; ls -l"'.format(hpss)
-            ls_error_str: str = "Target HPSS directory is not empty"
-            run_command(ls_command, ls_error_str)
+        ls_command: str = f'hsi -q "cd {command_info.hpss_path}; ls -l"'
+        ls_error_str: str = "Target HPSS directory is not empty"
+        run_command(ls_command, ls_error_str)
 
     # Create cache directory
     logger.debug(f"{ts_utc()}: Creating local cache directory")
-    os.chdir(path)
+    os.chdir(command_info.dir_to_archive_absolute)
     try:
-        os.makedirs(cache)
+        os.makedirs(command_info.cache_dir)
     except OSError as exc:
         if exc.errno != errno.EEXIST:
             cache_error_str: str = "Cannot create local cache directory"
@@ -88,11 +76,12 @@ def create():
 
     # Create and set up the database
     logger.debug(f"{ts_utc()}: Calling create_database()")
-    failures: List[str] = create_database(cache, args)
+    failures: List[str] = create_database(command_info, args)
 
-    # Transfer to HPSS. Always keep a local copy.
-    logger.debug(f"{ts_utc()}: calling hpss_put() for {get_db_filename(cache)}")
-    hpss_put(hpss, get_db_filename(cache), cache, keep=args.keep, is_index=True)
+    # Transfer to HPSS. Always keep a local copy of the database.
+    logger.debug(f"{ts_utc()}: calling hpss_put() for {command_info.get_db_name()}")
+    # TODO: (A) Continue refactoring from here
+    hpss_put(command_info, command_info.get_db_name(), is_index=True)
 
     logger.debug(f"{ts_utc()}: calling globus_finalize()")
     globus_finalize(non_blocking=args.non_blocking)
@@ -104,7 +93,7 @@ def create():
             logger.error("Failed to archive {}".format(file_path))
 
 
-def setup_create() -> Tuple[str, argparse.Namespace]:
+def setup_create(ci: CommandInfo) -> argparse.Namespace:
     # Parser
     parser: argparse.ArgumentParser = argparse.ArgumentParser(
         usage="zstash create [<args>] path", description="Create a new zstash archive"
@@ -175,27 +164,25 @@ def setup_create() -> Tuple[str, argparse.Namespace]:
     if args.verbose:
         logger.setLevel(logging.DEBUG)
 
-    # Copy configuration
-    config.path = os.path.abspath(args.path)
-    config.hpss = args.hpss
-    config.maxsize = int(1024 * 1024 * 1024 * args.maxsize)
-    cache: str
     if args.cache:
-        cache = args.cache
-    else:
-        cache = DEFAULT_CACHE
+        ci.cache_dir = args.cache
+    ci.keep = args.keep
+    ci.set_dir_to_archive(args.path)
+    ci.set_maxsize(args.maxsize)
+    ci.set_hpss_parameters(args.hpss)
 
-    return cache, args
+    return args
 
 
-def create_database(cache: str, args: argparse.Namespace) -> List[str]:
+def create_database(command_info: CommandInfo, args: argparse.Namespace) -> List[str]:
     # Create new database
     logger.debug(f"{ts_utc()}:Creating index database")
-    if os.path.exists(get_db_filename(cache)):
+    db_name: str = command_info.get_db_name()
+    if os.path.exists(db_name):
         # Remove old database
-        os.remove(get_db_filename(cache))
+        os.remove(db_name)
     con: sqlite3.Connection = sqlite3.connect(
-        get_db_filename(cache), detect_types=sqlite3.PARSE_DECLTYPES
+        db_name, detect_types=sqlite3.PARSE_DECLTYPES
     )
     cur: sqlite3.Cursor = con.cursor()
 
@@ -233,8 +220,8 @@ def create_database(cache: str, args: argparse.Namespace) -> List[str]:
 
     # Store configuration in database
     # Loop through all attributes of config.
-    for attr in dir(config):
-        value: Any = getattr(config, attr)
+    for attr in dir(command_info.config):
+        value: Any = getattr(command_info.config, attr)
         if not callable(value) and not attr.startswith("__"):
             # config.{attr} is not a function.
             # The attribute name does not start with "__"
@@ -244,7 +231,7 @@ def create_database(cache: str, args: argparse.Namespace) -> List[str]:
             cur.execute("insert into config values (?,?)", (attr, value))
     con.commit()
 
-    files: List[str] = get_files_to_archive(cache, args.include, args.exclude)
+    files: List[str] = get_files_to_archive(command_info.cache_dir, args.include, args.exclude)
 
     failures: List[str]
     if args.follow_symlinks:
@@ -255,7 +242,7 @@ def create_database(cache: str, args: argparse.Namespace) -> List[str]:
                 con,
                 -1,
                 files,
-                cache,
+                command_info.cache_dir,
                 args.keep,
                 args.follow_symlinks,
                 skip_tars_md5=args.no_tars_md5,
@@ -270,7 +257,7 @@ def create_database(cache: str, args: argparse.Namespace) -> List[str]:
             con,
             -1,
             files,
-            cache,
+            command_info.cache_dir,
             args.keep,
             args.follow_symlinks,
             skip_tars_md5=args.no_tars_md5,

zstash/globus.py

@@ -13,15 +13,15 @@
 from six.moves.urllib.parse import urlparse
 
 from .settings import logger
-from .utils import ts_utc
+from .utils import GlobusInfo, ts_utc
 
-hpss_endpoint_map = {
+HPSS_ENDPOINT_MAP = {
     "ALCF": "de463ec4-6d04-11e5-ba46-22000b92c6ec",
     "NERSC": "9cd89cfd-6d04-11e5-ba46-22000b92c6ec",
 }
 
 # This is used if the `globus_endpoint_uuid` is not set in `~/.zstash.ini`
-regex_endpoint_map = {
+REGEX_ENDPOINT_MAP = {
     r"theta.*\.alcf\.anl\.gov": "08925f04-569f-11e7-bef8-22000b9a448b",
     r"blueslogin.*\.lcrc\.anl\.gov": "15288284-7006-4041-ba1a-6b52501e49f1",
     r"chrlogin.*\.lcrc\.anl\.gov": "15288284-7006-4041-ba1a-6b52501e49f1",
@@ -73,26 +73,19 @@ def submit_transfer_with_checks(transfer_data):
     return task
 
 
-def globus_activate(hpss: str):
+def globus_activate(globus_info: GlobusInfo):
     """
     Read the local globus endpoint UUID from ~/.zstash.ini.
     If the ini file does not exist, create an ini file with empty values,
     and try to find the local endpoint UUID based on the FQDN
     """
-    global transfer_client
-    global local_endpoint
-    global remote_endpoint
-
-    url = urlparse(hpss)
-    if url.scheme != "globus":
-        return
-    remote_endpoint = url.netloc
+    globus_info.remote_endpoint = globus_info.url.netloc
 
     ini_path = os.path.expanduser("~/.zstash.ini")
     ini = configparser.ConfigParser()
     if ini.read(ini_path):
         if "local" in ini.sections():
-            local_endpoint = ini["local"].get("globus_endpoint_uuid")
+            globus_info.local_endpoint = ini["local"].get("globus_endpoint_uuid")
     else:
         ini["local"] = {"globus_endpoint_uuid": ""}
         try:
@@ -101,33 +94,31 @@ def globus_activate(hpss: str):
         except Exception as e:
             logger.error(e)
             sys.exit(1)
-    if not local_endpoint:
+    if not globus_info.local_endpoint:
         fqdn = socket.getfqdn()
         if re.fullmatch(r"n.*\.local", fqdn) and os.getenv("HOSTNAME", "NA").startswith(
             "compy"
         ):
             fqdn = "compy.pnl.gov"
-        for pattern in regex_endpoint_map.keys():
+        for pattern in REGEX_ENDPOINT_MAP.keys():
             if re.fullmatch(pattern, fqdn):
-                local_endpoint = regex_endpoint_map.get(pattern)
+                globus_info.local_endpoint = REGEX_ENDPOINT_MAP.get(pattern)
                 break
     # FQDN is not set on Perlmutter at NERSC
-    if not local_endpoint:
+    if not globus_info.local_endpoint:
         nersc_hostname = os.environ.get("NERSC_HOST")
         if nersc_hostname and (
             nersc_hostname == "perlmutter" or nersc_hostname == "unknown"
         ):
-            local_endpoint = regex_endpoint_map.get(r"perlmutter.*\.nersc\.gov")
-    if not local_endpoint:
+            globus_info.local_endpoint = REGEX_ENDPOINT_MAP.get(r"perlmutter.*\.nersc\.gov")
+    if not globus_info.local_endpoint:
         logger.error(
-            "{} does not have the local Globus endpoint set nor could one be found in regex_endpoint_map.".format(
-                ini_path
-            )
+            f"{ini_path} does not have the local Globus endpoint set nor could one be found in REGEX_ENDPOINT_MAP."
         )
         sys.exit(1)
 
-    if remote_endpoint.upper() in hpss_endpoint_map.keys():
-        remote_endpoint = hpss_endpoint_map.get(remote_endpoint.upper())
+    if globus_info.remote_endpoint.upper() in HPSS_ENDPOINT_MAP.keys():
+        globus_info.remote_endpoint = HPSS_ENDPOINT_MAP.get(remote_endpoint.upper())
 
     native_client = NativeClient(
         client_id="6c1629cf-446c-49e7-af95-323c6412397f",
@@ -136,15 +127,13 @@ def globus_activate(hpss: str):
     )
     native_client.login(no_local_server=True, refresh_tokens=True)
     transfer_authorizer = native_client.get_authorizers().get("transfer.api.globus.org")
-    transfer_client = TransferClient(authorizer=transfer_authorizer)
+    globus_info.transfer_client = TransferClient(authorizer=transfer_authorizer)
 
-    for ep_id in [local_endpoint, remote_endpoint]:
-        r = transfer_client.endpoint_autoactivate(ep_id, if_expires_in=600)
+    for ep_id in [globus_info.local_endpoint, globus_info.remote_endpoint]:
+        r = globus_info.transfer_client.endpoint_autoactivate(ep_id, if_expires_in=600)
         if r.get("code") == "AutoActivationFailed":
             logger.error(
-                "The {} endpoint is not activated or the current activation expires soon. Please go to https://app.globus.org/file-manager/collections/{} and (re)activate the endpoint.".format(
-                    ep_id, ep_id
-                )
+                f"The {ep_id} endpoint is not activated or the current activation expires soon. Please go to https://app.globus.org/file-manager/collections/{ep_id} and (re)activate the endpoint."
             )
             sys.exit(1)