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='-',
  28                     help="""
  29 Local file or directory where the data is to be written. Default: stdout.
  30 """)
  31 group = parser.add_mutually_exclusive_group()
  32 group.add_argument('--progress', action='store_true',
  33                    help="""
  34 Display human-readable progress on stderr (bytes and, if possible,
  35 percentage of total data size). This is the default behavior when it
  36 is not expected to interfere with the output: specifically, stderr is
  37 a tty _and_ either stdout is not a tty, or output is being written to
  38 named files rather than stdout.
  39 """)
  40 group.add_argument('--no-progress', action='store_true',
  41                    help="""
  42 Do not display human-readable progress on stderr.
  43 """)
  44 group.add_argument('--batch-progress', action='store_true',
  45                    help="""
  46 Display machine-readable progress on stderr (bytes and, if known,
  47 total data size).
  48 """)
  49 group = parser.add_mutually_exclusive_group()
  50 group.add_argument('--hash',
  51                     help="""
  52 Display the hash of each file as it is read from Keep, using the given
  53 hash algorithm. Supported algorithms include md5, sha1, sha224,
  54 sha256, sha384, and sha512.
  55 """)
  56 group.add_argument('--md5sum', action='store_const',
  57                     dest='hash', const='md5',
  58                     help="""
  59 Display the MD5 hash of each file as it is read from Keep.
  60 """)
  61 parser.add_argument('-n', action='store_true',
  62                     help="""
  63 Do not write any data -- just read from Keep, and report md5sums if
  64 requested.
  65 """)
  66 parser.add_argument('-r', action='store_true',
  67                     help="""
  68 Retrieve all files in the specified collection/prefix. This is the
  69 default behavior if the "locator" argument ends with a forward slash.
  70 """)
  71 group = parser.add_mutually_exclusive_group()
  72 group.add_argument('-f', action='store_true',
  73                    help="""
  74 Overwrite existing files while writing. The default behavior is to
  75 refuse to write *anything* if any of the output files already
  76 exist. As a special case, -f is not needed to write to stdout.
  77 """)
  78 group.add_argument('--skip-existing', action='store_true',
  79                    help="""
  80 Skip files that already exist. The default behavior is to refuse to
  81 write *anything* if any files exist that would have to be
  82 overwritten. This option causes even devices, sockets, and fifos to be
  83 skipped.
  84 """)
  85
  86 args = parser.parse_args()
  87
  88 if args.locator[-1] == os.sep:
  89     args.r = True
  90 if (args.r and
  91     not args.n and
  92     not (args.destination and
  93          os.path.isdir(args.destination))):
  94     parser.error('Destination is not a directory.')
  95 if not args.r and (os.path.isdir(args.destination) or
  96                    args.destination[-1] == os.path.sep):
  97     args.destination = os.path.join(args.destination,
  98                                     os.path.basename(args.locator))
  99     logger.debug("Appended source file name to destination directory: %s",
 100                  args.destination)
 101
 102 if args.destination == '/dev/stdout':
 103     args.destination = "-"
 104
 105 if args.destination == '-':
 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 != '-'
 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             if args.destination == "-":
 138                 sys.stdout.write(reader.manifest_text())
 139             else:
 140                 out_fd = os.open(args.destination, open_flags)
 141                 with os.fdopen(out_fd, 'wb') as out_file:
 142                     out_file.write(reader.manifest_text())
 143         except (IOError, OSError) as error:
 144             abort("can't write to '{}': {}".format(args.destination, error))
 145         except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:
 146             abort("failed to download '{}': {}".format(collection, error))
 147     sys.exit(0)
 148
 149 reader.normalize()
 150
 151 # Scan the collection. Make an array of (stream, file, local
 152 # destination filename) tuples, and add up total size to extract.
 153 todo = []
 154 todo_bytes = 0
 155 try:
 156     for s in reader.all_streams():
 157         for f in s.all_files():
 158             if get_prefix and get_prefix[-1] == os.sep:
 159                 if 0 != string.find(os.path.join(s.name(), f.name()),
 160                                     '.' + get_prefix):
 161                     continue
 162                 if args.destination == "-":
 163                     dest_path = "-"
 164                 else:
 165                     dest_path = os.path.join(
 166                         args.destination,
 167                         os.path.join(s.name(), f.name())[len(get_prefix)+1:])
 168                     if (not (args.n or args.f or args.skip_existing) and
 169                         os.path.exists(dest_path)):
 170                         abort('Local file %s already exists.' % (dest_path,))
 171             else:
 172                 if os.path.join(s.name(), f.name()) != '.' + get_prefix:
 173                     continue
 174                 dest_path = args.destination
 175             todo += [(s, f, dest_path)]
 176             todo_bytes += f.size()
 177 except arvados.errors.NotFoundError as e:
 178     abort(e)
 179
 180 # Read data, and (if not -n) write to local file(s) or pipe.
 181
 182 out_bytes = 0
 183 for s,f,outfilename in todo:
 184     outfile = None
 185     digestor = None
 186     if not args.n:
 187         if outfilename == "-":
 188             outfile = sys.stdout
 189         else:
 190             if args.skip_existing and os.path.exists(outfilename):
 191                 logger.debug('Local file %s exists. Skipping.', outfilename)
 192                 continue
 193             elif not args.f and (os.path.isfile(outfilename) or
 194                                os.path.isdir(outfilename)):
 195                 # Good thing we looked again: apparently this file wasn't
 196                 # here yet when we checked earlier.
 197                 abort('Local file %s already exists.' % (outfilename,))
 198             if args.r:
 199                 arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
 200             try:
 201                 outfile = open(outfilename, 'wb')
 202             except Exception as error:
 203                 abort('Open(%s) failed: %s' % (outfilename, error))
 204     if args.hash:
 205         digestor = hashlib.new(args.hash)
 206     try:
 207         for data in f.readall():
 208             if outfile:
 209                 outfile.write(data)
 210             if digestor:
 211                 digestor.update(data)
 212             out_bytes += len(data)
 213             if args.progress:
 214                 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %
 215                                  (out_bytes >> 20,
 216                                   todo_bytes >> 20,
 217                                   (100
 218                                    if todo_bytes==0
 219                                    else 100.0*out_bytes/todo_bytes)))
 220             elif args.batch_progress:
 221                 sys.stderr.write('%s %d read %d total\n' %
 222                                  (sys.argv[0], os.getpid(),
 223                                   out_bytes, todo_bytes))
 224         if digestor:
 225             sys.stderr.write("%s  %s/%s\n"
 226                              % (digestor.hexdigest(), s.name(), f.name()))
 227     except KeyboardInterrupt:
 228         if outfile and (outfile.fileno() > 2) and not outfile.closed:
 229             os.unlink(outfile.name)
 230         break
 231
 232 if args.progress:
 233     sys.stderr.write('\n')