sdk/python/arvados/commands/get.py

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