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