Merge pull request #12 from E3SM-Project/check

Check

Author: zshaheen

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="zstash",
-    version="0.0.2",
+    version="0.1.0",
     author="Chris Golaz, Zeshawn Shaheen",
     author_email="[email protected], [email protected]",
     description="Long term HPSS archiving software for E3SM",

tests/test.py

@@ -12,7 +12,7 @@ def write_file(name, contents):
 
 def run_cmd(cmd):
     """
-    Run a command while printing and returning the stdout and stderr
+    Run a command while printing and returning the stdout and stderr.
     """
     print('+ {}'.format(cmd))
     if isinstance(cmd, str):
@@ -29,9 +29,9 @@ def str_not_in(output, msg):
     """
     if msg in output:
         print('*'*40)
-        print('This was not supposed to be found: {}',format(msg))
+        print('This was not supposed to be found: {}'.format(msg))
         print('*'*40)
-        exit()
+        stop()
 
 def str_in(output, msg):
     """
@@ -41,7 +41,7 @@ def str_in(output, msg):
         print('*'*40)
         print('This was supposed to be found, but was not: {}'.format(msg))
         print('*'*40)
-        exit()
+        stop()
 
 def cleanup():
     """
@@ -55,15 +55,15 @@ def cleanup():
     cmd = 'hsi rm -R {}'.format(HPSS_PATH)
     run_cmd(cmd)
 
-def exit():
+def stop():
     """
-    Cleanup and stop running this script
+    Cleanup and stop running this script.
     """
     cleanup()
     sys.exit()
 
 
-# TODO: Change the hpss directory to a dir that's accessable to everyone
+# TODO: Change the hpss directory to a dir that's accessable to everyone.
 HPSS_PATH='/home/z/zshaheen/zstash_test'
 
 # Create files and directories
@@ -96,28 +96,32 @@ def exit():
 cmd = 'zstash create --hpss={} zstash_test'.format(HPSS_PATH)
 output, err = run_cmd(cmd)
 str_in(output+err, 'Transferring file to HPSS')
+str_not_in(output+err, 'ERROR')
 
 print('Testing chgrp')
 GROUP = 'acme'
 print('First, make sure that the files are not already in the {} group'.format(GROUP))
 cmd = 'hsi ls -l {}'.format(HPSS_PATH)
 output, err = run_cmd(cmd)
 str_not_in(output+err, GROUP)
+str_not_in(output+err, 'ERROR')
 print('Running zstash chgrp')
 cmd = 'zstash chgrp -R {} {}'.format(GROUP, HPSS_PATH)
 output, err = run_cmd(cmd)
+str_not_in(output+err, 'ERROR')
 print('Now check that the files are in the {} group'.format(GROUP))
 cmd = 'hsi ls -l {}'.format(HPSS_PATH)
 output, err = run_cmd(cmd)
 str_in(output+err, 'acme')
+str_not_in(output+err, 'ERROR')
 
 print('Running update on the newly created directory, nothing should happen')
 os.chdir('zstash_test')
 cmd = 'zstash update --hpss={}'.format(HPSS_PATH)
 output, err = run_cmd(cmd)
 os.chdir('../')
 str_in(output+err, 'Nothing to update')
-
+str_not_in(output+err, 'ERROR')
 
 print('Testing update with an actual change')
 if not os.path.exists('zstash_test/dir2'):
@@ -130,10 +134,17 @@ def exit():
 output, err = run_cmd(cmd)
 os.chdir('../')
 str_in(output+err, 'Transferring file to HPSS')
+str_not_in(output+err, 'ERROR')
 # Make sure none of the old files are moved
 str_not_in(output+err, 'file0')
 str_not_in(output+err, 'file_empty')
 str_not_in(output+err, 'empty_dir')
+str_not_in(output+err, 'ERROR')
+
+print('Testing the checking functionality')
+cmd = 'zstash check --hpss={}'.format(HPSS_PATH)
+output, err = run_cmd(cmd)
+str_not_in(output+err, 'ERROR')
 
 print('Testing the extract functionality')
 os.rename('zstash_test', 'zstash_test_backup')
@@ -150,13 +161,15 @@ def exit():
 str_in(output+err, 'Extracting dir/file1.txt')
 str_in(output+err, 'Extracting empty_dir')
 str_in(output+err, 'Extracting dir2/file2.txt')
+str_not_in(output+err, 'ERROR')
 
 print('Running update on the newly extracted directory, nothing should happen')
 os.chdir('zstash_test')
 cmd = 'zstash update --hpss={}'.format(HPSS_PATH)
 output, err = run_cmd(cmd)
 os.chdir('../')
 str_in(output+err, 'Nothing to update')
+str_not_in(output+err, 'ERROR')
 
 print('Verifying the data from database with the actual files')
 # Checksums from HPSS

zstash/check.py

@@ -0,0 +1,10 @@
+import logging
+import extract
+
+def check():
+    """
+    Check that the files in a given HPSS archive are valid.
+    """
+    # This basically just goes through the process of extracting the files,
+    # but doesn't actually save the output.
+    extract.extract(keep_files=False)

zstash/extract.py

@@ -11,11 +11,14 @@
 from settings import config, CACHE, BLOCK_SIZE, DB_FILENAME
 
 
-def extract():
+def extract(keep_files=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).
+    """
     parser = argparse.ArgumentParser(
         usage='zstash extract [<args>] [files]',
         description='Extract files from existing archive')
-    required = parser.add_argument_group('required named arguments')
     optional = parser.add_argument_group('optional named arguments')
     optional.add_argument('--hpss', type=str, help='path to HPSS storage')
     parser.add_argument('files', nargs='*', default=['*'])
@@ -50,7 +53,8 @@ def extract():
         config.hpss = args.hpss
 
     # Start doing actual work
-    logging.debug('Running zstash extract')
+    cmd = 'extract' if keep_files else 'check'
+    logging.debug('Running zstash ' + cmd)
     logging.debug('Local path : %s' % (config.path))
     logging.debug('HPSS path  : %s' % (config.hpss))
     logging.debug('Max size  : %i' % (config.maxsize))
@@ -59,7 +63,7 @@ def extract():
     # Find matching files
     matches = []
     for file in args.files:
-        cur.execute(u"select * from files where name GLOB ?", (file,))
+        cur.execute(u"select * from files where name GLOB ? or tar GLOB ?", (file, file))
         matches = matches + cur.fetchall()
 
     # Remove duplicates
@@ -69,15 +73,32 @@ def extract():
     matches = sorted(matches, key=lambda x: (x[5], x[6]))
 
     # Retrieve from tapes
-    extractFiles(matches)
+    failures = extractFiles(matches, keep_files)
 
     # Close database
     logging.debug('Closing index database')
     con.close()
 
+    if failures:
+        logging.error('Encountered an error for files:')
+        for fail in failures:
+            logging.error('{} in {}'.format(fail[1], fail[5]))
 
-def extractFiles(files):
+        broken_tars = set(sorted([f[5] for f in failures]))
+        logging.error('The following tar archives had errors:')
+        for tar in broken_tars:
+            logging.error(tar)
 
+
+def extractFiles(files, keep_files):
+    """
+    Given a list of database rows, extract the files from the
+    tar archives to the current location on disk.
+
+    If keep_files is False, the files are not extracted.
+    This is used for when checking if the files in an HPSS
+    repository are valid.
+    """
     failures = []
     tfname = None
     newtar = True
@@ -97,7 +118,8 @@ def extractFiles(files):
             tar = tarfile.open(tfname, "r")
 
         # Extract file
-        logging.info('Extracting %s' % (file[1]))
+        cmd = 'Extracting' if keep_files else 'Checking'
+        logging.info(cmd + ' %s' % (file[1]))
         try:
 
             # Seek file position
@@ -108,39 +130,48 @@ def extractFiles(files):
 
             if tarinfo.isfile():
                 # fileobj to extract
-                fin = tar.extractfile(tarinfo)
-                fname = tarinfo.name
-                path, name = os.path.split(fname)
-                if path != '':
-                    if not os.path.isdir(path):
-                        os.makedirs(path)
-                fout = open(fname, 'w')
-                hash_md5 = hashlib.md5()
-                while True:
-                    s = fin.read(BLOCK_SIZE)
-                    if len(s) > 0:
-                        fout.write(s)
-                        hash_md5.update(s)
-                    if len(s) < BLOCK_SIZE:
-                        break
-                fin.close()
-                fout.close()
+                try:
+                    fin = tar.extractfile(tarinfo)
+                    fname = tarinfo.name
+                    path, name = os.path.split(fname)
+                    if path != '':
+                        if not os.path.isdir(path):
+                            os.makedirs(path)
+                    if keep_files:
+                        fout = open(fname, 'w')
+
+                    hash_md5 = hashlib.md5()
+                    while True:
+                        s = fin.read(BLOCK_SIZE)
+                        if len(s) > 0:
+                            hash_md5.update(s)
+                            if keep_files:
+                                fout.write(s)
+                        if len(s) < BLOCK_SIZE:
+                            break
+                finally:
+                    fin.close()
+                    if keep_files:
+                        fout.close()
+
                 md5 = hash_md5.hexdigest()
-                tar.chown(tarinfo, fname)
-                tar.chmod(tarinfo, fname)
-                tar.utime(tarinfo, fname)
-                # Verify size
-                if os.path.getsize(fname) != file[2]:
-                    logging.error('size mismatch for: %s' % (fname))
+                if keep_files:
+                    tar.chown(tarinfo, fname)
+                    tar.chmod(tarinfo, fname)
+                    tar.utime(tarinfo, fname)
+                    # Verify size
+                    if os.path.getsize(fname) != file[2]:
+                        logging.error('size mismatch for: %s' % (fname))
                 # Verify md5 checksum
                 if md5 != file[4]:
                     logging.error('md5 mismatch for: %s' % (fname))
                     logging.error('md5 of extracted file: %s' % (md5))
                     logging.error('md5 of original file:  %s' % (file[4]))
+                    failures.append(file)
                 else:
                     logging.debug('Valid md5: %s %s' % (md5, fname))
 
-            else:
+            elif keep_files:
                 tar.extract(tarinfo)
                 # Note: tar.extract() will not restore time stamps of symbolic
                 # links. Could not find a Python-way to restore it either, so
@@ -155,6 +186,7 @@ def extractFiles(files):
         except:
             traceback.print_exc()
             logging.error('Retrieving %s' % (file[1]))
+            failures.append(file)
 
         # Close current archive?
         if (i == nfiles-1 or files[i][5] != files[i+1][5]):
@@ -165,3 +197,5 @@ def extractFiles(files):
 
             # Open new archive next time
             newtar = True
+
+    return failures

zstash/main.py

@@ -10,6 +10,7 @@
 from update import update
 from extract import extract
 from chgrp import chgrp
+from check import check
 
 
 # -----------------------------------------------------------------------------
@@ -24,6 +25,8 @@ def main():
   create     create new archive
   update     update existing archive
   extract    extract files from archive
+  chgrp      change the group of an archive
+  check      check the integrity of the files in the archive
 
 For help with a specific command
   zstash command --help
@@ -42,10 +45,12 @@ def main():
         extract()
     elif args.command == 'chgrp':
         chgrp()
+    elif args.command == 'check':
+        check()
     else:
         print 'Unrecognized command'
         parser.print_help()
-        exit(1)
+        sys.exit(1)
 
 
 # -----------------------------------------------------------------------------