sdk/cli/bin/arv-get

   1 #!/usr/bin/env python
   2
   3 import argparse
   4 import hashlib
   5 import os
   6 import re
   7 import string
   8 import sys
   9 import logging
  10
  11 logger = logging.getLogger(os.path.basename(sys.argv[0]))
  12
  13 parser = argparse.ArgumentParser(
  14     description='Copy data from Keep to a local file or pipe.')
  15 parser.add_argument('locator', type=str,
  16                     help="""
  17 Collection locator, optionally with a file path or prefix.
  18 """)
  19 parser.add_argument('destination', type=str, nargs='?', default='/dev/stdout',
  20                     help="""
  21 Local file or directory where the data is to be written. Default:
  22 /dev/stdout.
  23 """)
  24 group = parser.add_mutually_exclusive_group()
  25 group.add_argument('--progress', action='store_true',
  26                    help="""
  27 Display human-readable progress on stderr (bytes and, if possible,
  28 percentage of total data size). This is the default behavior when
  29 stderr is a tty and stdout is not a tty.
  30 """)
  31 group.add_argument('--no-progress', action='store_true',
  32                    help="""
  33 Do not display human-readable progress on stderr.
  34 """)
  35 group.add_argument('--batch-progress', action='store_true',
  36                    help="""
  37 Display machine-readable progress on stderr (bytes and, if known,
  38 total data size).
  39 """)
  40 group = parser.add_mutually_exclusive_group()
  41 group.add_argument('--hash',
  42                     help="""
  43 Display the hash of each file as it is read from Keep, using the given
  44 hash algorithm. Supported algorithms include md5, sha1, sha224,
  45 sha256, sha384, and sha512.
  46 """)
  47 group.add_argument('--md5sum', action='store_const',
  48                     dest='hash', const='md5',
  49                     help="""
  50 Display the MD5 hash of each file as it is read from Keep.
  51 """)
  52 parser.add_argument('-n', action='store_true',
  53                     help="""
  54 Do not write any data -- just read from Keep, and report md5sums if
  55 requested.
  56 """)
  57 parser.add_argument('-r', action='store_true',
  58                     help="""
  59 Retrieve all files in the specified collection/prefix. This is the
  60 default behavior if the "locator" argument ends with a forward slash.
  61 """)
  62
  63 args = parser.parse_args()
  64
  65 if args.locator[-1] == os.sep:
  66     args.r = True
  67 if (args.r and
  68     not args.n and
  69     not (args.destination and
  70          os.path.isdir(args.destination))):
  71     parser.error('Destination is not a directory.')
  72 if not args.r and (os.path.isdir(args.destination) or
  73                    args.destination[-1] == os.path.sep):
  74     args.destination = os.path.join(args.destination,
  75                                     os.path.basename(args.locator))
  76     logger.debug("Appended source file name to destination directory: %s" %
  77                  args.destination)
  78
  79 # Turn on --progress by default if stderr is a tty and stdout isn't.
  80 if (not (args.batch_progress or args.no_progress)
  81     and os.isatty(sys.stderr.fileno())
  82     and not os.isatty(sys.stdout.fileno())):
  83     args.progress = True
  84
  85 if args.destination == '-':
  86     args.destination = '/dev/stdout'
  87 args.destination = args.destination.rstrip(os.sep)
  88
  89
  90 import arvados
  91
  92 r = re.search(r'^(.*?)(/.*)?$', args.locator)
  93 collection = r.group(1)
  94 get_prefix = r.group(2)
  95 if args.r and not get_prefix:
  96     get_prefix = os.sep
  97
  98 todo = []
  99 todo_bytes = 0
 100 if not get_prefix:
 101     try:
 102         if not args.n:
 103             with open(args.destination, 'wb') as f:
 104                 f.write(arvados.Keep.get(collection))
 105         sys.exit(0)
 106     except arvados.errors.NotFoundError as e:
 107         logger.error(e)
 108         sys.exit(1)
 109
 110 reader = arvados.CollectionReader(collection)
 111
 112 # Scan the collection. Make an array of (stream, file, local
 113 # destination filename) tuples, and add up total size to extract.
 114
 115 try:
 116     for s in reader.all_streams():
 117         for f in s.all_files():
 118             if get_prefix and get_prefix[-1] == os.sep:
 119                 if 0 != string.find(os.path.join(s.name(), f.name()),
 120                                     '.' + get_prefix):
 121                     continue
 122                 dest_path = os.path.join(
 123                     args.destination,
 124                     os.path.join(s.name(), f.name())[len(get_prefix)+1:])
 125             else:
 126                 if os.path.join(s.name(), f.name()) != '.' + get_prefix:
 127                     continue
 128                 dest_path = args.destination
 129             todo += [(s, f, dest_path)]
 130             todo_bytes += f.size()
 131 except arvados.errors.NotFoundError as e:
 132     logger.error(e)
 133     sys.exit(1)
 134
 135 # Read data, and (if not -n) write to local file(s) or pipe.
 136
 137 out_bytes = 0
 138 for s,f,outfilename in todo:
 139     outfile = None
 140     digestor = None
 141     if not args.n:
 142         if args.r:
 143             arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
 144         try:
 145             outfile = open(outfilename, 'wb')
 146         except Exception as e:
 147             logger.error('Open(%s) failed: %s' % (outfilename, e))
 148     if args.hash:
 149         digestor = hashlib.new(args.hash)
 150     try:
 151         for data in f.readall():
 152             if outfile:
 153                 outfile.write(data)
 154             if digestor:
 155                 digestor.update(data)
 156             out_bytes += len(data)
 157             if args.progress:
 158                 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %
 159                                  (out_bytes >> 20,
 160                                   todo_bytes >> 20,
 161                                   (100
 162                                    if todo_bytes==0
 163                                    else 100.0*out_bytes/todo_bytes)))
 164             elif args.batch_progress:
 165                 sys.stderr.write('%s %d read %d total\n' %
 166                                  (sys.argv[0], os.getpid(),
 167                                   out_bytes, todo_bytes))
 168         if digestor:
 169             sys.stderr.write("%s  %s/%s\n"
 170                              % (digestor.hexdigest(), s.name(), f.name()))
 171     except KeyboardInterrupt:
 172         if outfile:
 173             os.unlink(outfilename)
 174         break
 175
 176 if args.progress:
 177     sys.stderr.write('\n')