sdk/python/arvados/commands/get.py

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