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 import arvados
  12 import arvados.commands._util as arv_cmd
  13
  14 logger = logging.getLogger('arvados.arv-get')
  15
  16 def abort(msg, code=1):
  17     print >>sys.stderr, "arv-get:", msg
  18     exit(code)
  19
  20 parser = argparse.ArgumentParser(
  21     description='Copy data from Keep to a local file or pipe.',
  22     parents=[arv_cmd.retry_opt])
  23 parser.add_argument('locator', type=str,
  24                     help="""
  25 Collection locator, optionally with a file path or prefix.
  26 """)
  27 parser.add_argument('destination', type=str, nargs='?', default='/dev/stdout',
  28                     help="""
  29 Local file or directory where the data is to be written. Default:
  30 /dev/stdout.
  31 """)
  32 group = parser.add_mutually_exclusive_group()
  33 group.add_argument('--progress', action='store_true',
  34                    help="""
  35 Display human-readable progress on stderr (bytes and, if possible,
  36 percentage of total data size). This is the default behavior when it
  37 is not expected to interfere with the output: specifically, stderr is
  38 a tty _and_ either stdout is not a tty, or output is being written to
  39 named files rather than stdout.
  40 """)
  41 group.add_argument('--no-progress', action='store_true',
  42                    help="""
  43 Do not display human-readable progress on stderr.
  44 """)
  45 group.add_argument('--batch-progress', action='store_true',
  46                    help="""
  47 Display machine-readable progress on stderr (bytes and, if known,
  48 total data size).
  49 """)
  50 group = parser.add_mutually_exclusive_group()
  51 group.add_argument('--hash',
  52                     help="""
  53 Display the hash of each file as it is read from Keep, using the given
  54 hash algorithm. Supported algorithms include md5, sha1, sha224,
  55 sha256, sha384, and sha512.
  56 """)
  57 group.add_argument('--md5sum', action='store_const',
  58                     dest='hash', const='md5',
  59                     help="""
  60 Display the MD5 hash of each file as it is read from Keep.
  61 """)
  62 parser.add_argument('-n', action='store_true',
  63                     help="""
  64 Do not write any data -- just read from Keep, and report md5sums if
  65 requested.
  66 """)
  67 parser.add_argument('-r', action='store_true',
  68                     help="""
  69 Retrieve all files in the specified collection/prefix. This is the
  70 default behavior if the "locator" argument ends with a forward slash.
  71 """)
  72 group = parser.add_mutually_exclusive_group()
  73 group.add_argument('-f', action='store_true',
  74                    help="""
  75 Overwrite existing files while writing. The default behavior is to
  76 refuse to write *anything* if any of the output files already
  77 exist. As a special case, -f is not needed to write to /dev/stdout.
  78 """)
  79 group.add_argument('--skip-existing', action='store_true',
  80                    help="""
  81 Skip files that already exist. The default behavior is to refuse to
  82 write *anything* if any files exist that would have to be
  83 overwritten. This option causes even devices, sockets, and fifos to be
  84 skipped.
  85 """)
  86
  87 args = parser.parse_args()
  88
  89 if args.locator[-1] == os.sep:
  90     args.r = True
  91 if (args.r and
  92     not args.n and
  93     not (args.destination and
  94          os.path.isdir(args.destination))):
  95     parser.error('Destination is not a directory.')
  96 if not args.r and (os.path.isdir(args.destination) or
  97                    args.destination[-1] == os.path.sep):
  98     args.destination = os.path.join(args.destination,
  99                                     os.path.basename(args.locator))
 100     logger.debug("Appended source file name to destination directory: %s",
 101                  args.destination)
 102
 103 if args.destination == '-':
 104     args.destination = '/dev/stdout'
 105 if args.destination == '/dev/stdout':
 106     # Normally you have to use -f to write to a file (or device) that
 107     # already exists, but "-" and "/dev/stdout" are common enough to
 108     # merit a special exception.
 109     args.f = True
 110 else:
 111     args.destination = args.destination.rstrip(os.sep)
 112
 113 # Turn on --progress by default if stderr is a tty and output is
 114 # either going to a named file, or going (via stdout) to something
 115 # that isn't a tty.
 116 if (not (args.batch_progress or args.no_progress)
 117     and sys.stderr.isatty()
 118     and (args.destination != '/dev/stdout'
 119          or not sys.stdout.isatty())):
 120     args.progress = True
 121
 122
 123 r = re.search(r'^(.*?)(/.*)?$', args.locator)
 124 collection = r.group(1)
 125 get_prefix = r.group(2)
 126 if args.r and not get_prefix:
 127     get_prefix = os.sep
 128 api_client = arvados.api('v1')
 129 reader = arvados.CollectionReader(collection, num_retries=args.retries)
 130
 131 if not get_prefix:
 132     if not args.n:
 133         open_flags = os.O_CREAT | os.O_WRONLY
 134         if not args.f:
 135             open_flags |= os.O_EXCL
 136         try:
 137             out_fd = os.open(args.destination, open_flags)
 138             with os.fdopen(out_fd, 'wb') as out_file:
 139                 out_file.write(reader.manifest_text())
 140         except (IOError, OSError) as error:
 141             abort("can't write to '{}': {}".format(args.destination, error))
 142         except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:
 143             abort("failed to download '{}': {}".format(collection, error))
 144     sys.exit(0)
 145
 146 reader.normalize()
 147
 148 # Scan the collection. Make an array of (stream, file, local
 149 # destination filename) tuples, and add up total size to extract.
 150 todo = []
 151 todo_bytes = 0
 152 try:
 153     for s in reader.all_streams():
 154         for f in s.all_files():
 155             if get_prefix and get_prefix[-1] == os.sep:
 156                 if 0 != string.find(os.path.join(s.name(), f.name()),
 157                                     '.' + get_prefix):
 158                     continue
 159                 dest_path = os.path.join(
 160                     args.destination,
 161                     os.path.join(s.name(), f.name())[len(get_prefix)+1:])
 162                 if (not (args.n or args.f or args.skip_existing) and
 163                     os.path.exists(dest_path)):
 164                     abort('Local file %s already exists.' % (dest_path,))
 165             else:
 166                 if os.path.join(s.name(), f.name()) != '.' + get_prefix:
 167                     continue
 168                 dest_path = args.destination
 169             todo += [(s, f, dest_path)]
 170             todo_bytes += f.size()
 171 except arvados.errors.NotFoundError as e:
 172     abort(e)
 173
 174 # Read data, and (if not -n) write to local file(s) or pipe.
 175
 176 out_bytes = 0
 177 for s,f,outfilename in todo:
 178     outfile = None
 179     digestor = None
 180     if not args.n:
 181         if args.skip_existing and os.path.exists(outfilename):
 182             logger.debug('Local file %s exists. Skipping.', outfilename)
 183             continue
 184         elif not args.f and (os.path.isfile(outfilename) or
 185                            os.path.isdir(outfilename)):
 186             # Good thing we looked again: apparently this file wasn't
 187             # here yet when we checked earlier.
 188             abort('Local file %s already exists.' % (outfilename,))
 189         if args.r:
 190             arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
 191         try:
 192             outfile = open(outfilename, 'wb')
 193         except Exception as error:
 194             abort('Open(%s) failed: %s' % (outfilename, error))
 195     if args.hash:
 196         digestor = hashlib.new(args.hash)
 197     try:
 198         for data in f.readall():
 199             if outfile:
 200                 outfile.write(data)
 201             if digestor:
 202                 digestor.update(data)
 203             out_bytes += len(data)
 204             if args.progress:
 205                 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %
 206                                  (out_bytes >> 20,
 207                                   todo_bytes >> 20,
 208                                   (100
 209                                    if todo_bytes==0
 210                                    else 100.0*out_bytes/todo_bytes)))
 211             elif args.batch_progress:
 212                 sys.stderr.write('%s %d read %d total\n' %
 213                                  (sys.argv[0], os.getpid(),
 214                                   out_bytes, todo_bytes))
 215         if digestor:
 216             sys.stderr.write("%s  %s/%s\n"
 217                              % (digestor.hexdigest(), s.name(), f.name()))
 218     except KeyboardInterrupt:
 219         if outfile and outfile != '/dev/stdout':
 220             os.unlink(outfilename)
 221         break
 222
 223 if args.progress:
 224     sys.stderr.write('\n')