zstash extract refactored

Author: forsyth2

zstash/extract.py

@@ -26,23 +26,19 @@
     TIME_TOL,
     FilesRow,
     TupleFilesRow,
-    config,
-    get_db_filename,
     logger,
 )
-from .utils import tars_table_exists, update_config
+from .utils import CommandInfo, HPSSType, tars_table_exists
 
 
 def extract(keep_files: bool = True):
     """
     Given an HPSS path in the zstash database or passed via the command line,
     extract the archived data based on the file pattern (if given).
     """
-    args: argparse.Namespace
-    cache: str
-    args, cache = setup_extract()
-
-    failures: List[FilesRow] = extract_database(args, cache, keep_files)
+    command_info = CommandInfo("extract")
+    args: argparse.Namespace = setup_extract(command_info)
+    failures: List[FilesRow] = extract_database(command_info, args, keep_files)
 
     if failures:
         logger.error("Encountered an error for files:")
@@ -64,7 +60,7 @@ def extract(keep_files: bool = True):
         )
 
 
-def setup_extract() -> Tuple[argparse.Namespace, str]:
+def setup_extract(command_info: CommandInfo) -> Tuple[argparse.Namespace, str]:
     parser: argparse.ArgumentParser = argparse.ArgumentParser(
         usage="zstash extract [<args>] [files]",
         description="Extract files from existing archive",
@@ -103,19 +99,22 @@ def setup_extract() -> Tuple[argparse.Namespace, str]:
     )
     parser.add_argument("files", nargs="*", default=["*"])
     args: argparse.Namespace = parser.parse_args(sys.argv[2:])
-    if args.hpss and args.hpss.lower() == "none":
+
+    if args.hpss and (args.hpss.lower() == "none"):
         args.hpss = "none"
-    if args.cache:
-        cache = args.cache
-    else:
-        cache = DEFAULT_CACHE
     # Note: setting logging level to anything other than DEBUG doesn't work with
     # multiple workers. This must have someting to do with the custom logger
     # implemented for multiple workers.
     if args.verbose or args.workers > 1:
         logger.setLevel(logging.DEBUG)
 
-    return args, cache
+    if args.cache:
+        command_info.cache_dir = args.cache
+    command_info.keep = args.keep
+    command_info.set_dir_to_archive(os.getcwd())
+    command_info.set_hpss_parameters(args.hpss, null_hpss_allowed=True)
+
+    return args
 
 
 def parse_tars_option(tars: str, first_tar: str, last_tar: str) -> List[str]:
@@ -156,58 +155,43 @@ def parse_tars_option(tars: str, first_tar: str, last_tar: str) -> List[str]:
     return tar_list
 
 
-def extract_database(
-    args: argparse.Namespace, cache: str, keep_files: bool
+def extract_database(command_info: CommandInfo,
+    args: argparse.Namespace, keep_files: bool
 ) -> List[FilesRow]:
 
     # Open database
     logger.debug("Opening index database")
-    if not os.path.exists(get_db_filename(cache)):
+    if not os.path.exists(command_info.get_db_name()):
         # Will need to retrieve from HPSS
-        if args.hpss is not None:
-            config.hpss = args.hpss
-            if config.hpss is not None:
-                hpss: str = config.hpss
-            else:
-                raise TypeError("Invalid config.hpss={}".format(config.hpss))
-            hpss_get(hpss, get_db_filename(cache), cache)
+        if command_info.hpss_type != HPSSType.UNDEFINED:
+            hpss_get(command_info.config.hpss, command_info.get_db_name())
         else:
             error_str: str = (
                 "--hpss argument is required when local copy of database is unavailable"
             )
             logger.error(error_str)
-
             raise ValueError(error_str)
     con: sqlite3.Connection = sqlite3.connect(
-        get_db_filename(cache), detect_types=sqlite3.PARSE_DECLTYPES
+        command_info.get_db_name(), detect_types=sqlite3.PARSE_DECLTYPES
     )
     cur: sqlite3.Cursor = con.cursor()
 
-    update_config(cur)
-    if config.maxsize is not None:
-        maxsize = config.maxsize
-    else:
-        raise TypeError("Invalid config.maxsize={}".format(config.maxsize))
-    config.maxsize = int(maxsize)
-
-    # The command line arg should always have precedence
-    if args.hpss is not None:
-        config.hpss = args.hpss
-    keep: bool
-    if config.hpss == "none":
-        # If no HPSS is available, always keep the files.
-        keep = True
-    else:
-        keep = args.keep
+    command_info.update_config_using_db(cur)
+    command_info.validate_maxsize()
+
+    if command_info.hpss_type == HPSSType.NO_HPSS:
+        # If not using HPSS, always keep the files.
+        command_info.keep = True
+    # else: keep command_info.keep set to args.keep
 
     # Start doing actual work
     cmd: str = "extract" if keep_files else "check"
 
     logger.debug("Running zstash " + cmd)
-    logger.debug("Local path : {}".format(config.path))
-    logger.debug("HPSS path  : {}".format(config.hpss))
-    logger.debug("Max size  : {}".format(config.maxsize))
-    logger.debug("Keep local tar files : {}".format(keep))
+    logger.debug(f"Local path : {command_info.config.path}")
+    logger.debug(f"HPSS path  : {command_info.config.hpss}")
+    logger.debug(f"Max size  : {command_info.config.maxsize}")
+    logger.debug(f"Keep local tar files : {command_info.keep}")
 
     matches_: List[TupleFilesRow] = []
     if args.tars is not None:
@@ -277,10 +261,10 @@ def extract_database(
     if args.workers > 1:
         logger.debug("Running zstash {} with multiprocessing".format(cmd))
         failures = multiprocess_extract(
-            args.workers, matches, keep_files, keep, cache, cur, args
+            args.workers, matches, keep_files, command_info.keep, command_info.cache_dir, cur, args
         )
     else:
-        failures = extractFiles(matches, keep_files, keep, cache, cur, args)
+        failures = extractFiles(matches, keep_files, command_info.keep, command_info.cache_dir, cur, args)
 
     # Close database
     logger.debug("Closing index database")