sdk/python/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 group = parser.add_mutually_exclusive_group()
  63 group.add_argument('-f', action='store_true',
  64                    help="""
  65 Overwrite existing files while writing. The default behavior is to
  66 refuse to write *anything* if any of the output files already
  67 exist. As a special case, -f is not needed to write to /dev/stdout.
  68 """)
  69 group.add_argument('--skip-existing', action='store_true',
  70                    help="""
  71 Skip files that already exist. The default behavior is to refuse to
  72 write *anything* if any files exist that would have to be
  73 overwritten. This option causes even devices, sockets, and fifos to be
  74 skipped.
  75 """)
  76
  77 args = parser.parse_args()
  78
  79 if args.locator[-1] == os.sep:
  80     args.r = True
  81 if (args.r and
  82     not args.n and
  83     not (args.destination and
  84          os.path.isdir(args.destination))):
  85     parser.error('Destination is not a directory.')
  86 if not args.r and (os.path.isdir(args.destination) or
  87                    args.destination[-1] == os.path.sep):
  88     args.destination = os.path.join(args.destination,
  89                                     os.path.basename(args.locator))
  90     logger.debug("Appended source file name to destination directory: %s" %
  91                  args.destination)
  92
  93 # Turn on --progress by default if stderr is a tty and stdout isn't.
  94 if (not (args.batch_progress or args.no_progress)
  95     and os.isatty(sys.stderr.fileno())
  96     and not os.isatty(sys.stdout.fileno())):
  97     args.progress = True
  98
  99 if args.destination == '-':
 100     args.destination = '/dev/stdout'
 101 if args.destination == '/dev/stdout':
 102     # Normally you have to use -f to write to a file (or device) that
 103     # already exists, but "-" and "/dev/stdout" are common enough to
 104     # merit a special exception.
 105     args.f = True
 106 else:
 107     args.destination = args.destination.rstrip(os.sep)
 108
 109
 110 import arvados
 111
 112 r = re.search(r'^(.*?)(/.*)?$', args.locator)
 113 collection = r.group(1)
 114 get_prefix = r.group(2)
 115 if args.r and not get_prefix:
 116     get_prefix = os.sep
 117
 118 todo = []
 119 todo_bytes = 0
 120 if not get_prefix:
 121     try:
 122         if not args.n:
 123             if not args.f and os.path.exists(args.destination):
 124                 logger.error('Local file %s already exists' % args.destination)
 125                 sys.exit(1)
 126             with open(args.destination, 'wb') as f:
 127                 f.write(arvados.Keep.get(collection))
 128         sys.exit(0)
 129     except arvados.errors.NotFoundError as e:
 130         logger.error(e)
 131         sys.exit(1)
 132
 133 reader = arvados.CollectionReader(collection)
 134
 135 # Scan the collection. Make an array of (stream, file, local
 136 # destination filename) tuples, and add up total size to extract.
 137
 138 try:
 139     for s in reader.all_streams():
 140         for f in s.all_files():
 141             if get_prefix and get_prefix[-1] == os.sep:
 142                 if 0 != string.find(os.path.join(s.name(), f.name()),
 143                                     '.' + get_prefix):
 144                     continue
 145                 dest_path = os.path.join(
 146                     args.destination,
 147                     os.path.join(s.name(), f.name())[len(get_prefix)+1:])
 148                 if (not (args.f or args.skip_existing) and
 149                     os.path.exists(dest_path)):
 150                     logger.error('Local file %s already exists' % dest_path)
 151                     sys.exit(1)
 152             else:
 153                 if os.path.join(s.name(), f.name()) != '.' + get_prefix:
 154                     continue
 155                 dest_path = args.destination
 156             todo += [(s, f, dest_path)]
 157             todo_bytes += f.size()
 158 except arvados.errors.NotFoundError as e:
 159     logger.error(e)
 160     sys.exit(1)
 161
 162 # Read data, and (if not -n) write to local file(s) or pipe.
 163
 164 out_bytes = 0
 165 for s,f,outfilename in todo:
 166     outfile = None
 167     digestor = None
 168     if not args.n:
 169         if args.skip_existing and os.path.exists(outfilename):
 170             logger.debug('Local file %s exists. Skipping.' % outfilename)
 171             continue
 172         elif not args.f and (os.path.isfile(outfilename) or
 173                            os.path.isdir(outfilename)):
 174             # Good thing we looked again: apparently this file wasn't
 175             # here yet when we checked earlier.
 176             logger.error('Local file %s already exists' % outfilename)
 177             sys.exit(1)
 178         if args.r:
 179             arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
 180         try:
 181             outfile = open(outfilename, 'wb')
 182         except Exception as e:
 183             logger.error('Open(%s) failed: %s' % (outfilename, e))
 184             sys.exit(1)
 185     if args.hash:
 186         digestor = hashlib.new(args.hash)
 187     try:
 188         for data in f.readall():
 189             if outfile:
 190                 outfile.write(data)
 191             if digestor:
 192                 digestor.update(data)
 193             out_bytes += len(data)
 194             if args.progress:
 195                 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %
 196                                  (out_bytes >> 20,
 197                                   todo_bytes >> 20,
 198                                   (100
 199                                    if todo_bytes==0
 200                                    else 100.0*out_bytes/todo_bytes)))
 201             elif args.batch_progress:
 202                 sys.stderr.write('%s %d read %d total\n' %
 203                                  (sys.argv[0], os.getpid(),
 204                                   out_bytes, todo_bytes))
 205         if digestor:
 206             sys.stderr.write("%s  %s/%s\n"
 207                              % (digestor.hexdigest(), s.name(), f.name()))
 208     except KeyboardInterrupt:
 209         if outfile:
 210             os.unlink(outfilename)
 211         break
 212
 213 if args.progress:
 214     sys.stderr.write('\n')