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('-v', action='count', default=0,
  85                     help="""
  86 Once for verbose mode, twice for debug mode.
  87 """)
  88 group.add_argument('--skip-existing', action='store_true',
  89                    help="""
  90 Skip files that already exist. The default behavior is to refuse to
  91 write *anything* if any files exist that would have to be
  92 overwritten. This option causes even devices, sockets, and fifos to be
  93 skipped.
  94 """)
  95 group.add_argument('--strip-manifest', action='store_true', default=False,
  96                    help="""
  97 When getting a collection manifest, strip its access tokens before writing
  98 it.
  99 """)
 100
 101 def parse_arguments(arguments, stdout, stderr):
 102     args = parser.parse_args(arguments)
 103
 104     if args.locator[-1] == os.sep:
 105         args.r = True
 106     if (args.r and
 107         not args.n and
 108         not (args.destination and
 109              os.path.isdir(args.destination))):
 110         parser.error('Destination is not a directory.')
 111     if not args.r and (os.path.isdir(args.destination) or
 112                        args.destination[-1] == os.path.sep):
 113         args.destination = os.path.join(args.destination,
 114                                         os.path.basename(args.locator))
 115         logger.debug("Appended source file name to destination directory: %s",
 116                      args.destination)
 117
 118     if args.destination == '/dev/stdout':
 119         args.destination = "-"
 120
 121     if args.destination == '-':
 122         # Normally you have to use -f to write to a file (or device) that
 123         # already exists, but "-" and "/dev/stdout" are common enough to
 124         # merit a special exception.
 125         args.f = True
 126     else:
 127         args.destination = args.destination.rstrip(os.sep)
 128
 129     # Turn on --progress by default if stderr is a tty and output is
 130     # either going to a named file, or going (via stdout) to something
 131     # that isn't a tty.
 132     if (not (args.batch_progress or args.no_progress)
 133         and stderr.isatty()
 134         and (args.destination != '-'
 135              or not stdout.isatty())):
 136         args.progress = True
 137     return args
 138
 139 def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
 140     global api_client
 141
 142     if stdout is sys.stdout and hasattr(stdout, 'buffer'):
 143         # in Python 3, write to stdout as binary
 144         stdout = stdout.buffer
 145
 146     args = parse_arguments(arguments, stdout, stderr)
 147     logger.setLevel(logging.WARNING - 10 * args.v)
 148
 149     request_id = arvados.util.new_request_id()
 150     logger.info('X-Request-Id: '+request_id)
 151
 152     if api_client is None:
 153         api_client = arvados.api('v1', request_id=request_id)
 154
 155     r = re.search(r'^(.*?)(/.*)?$', args.locator)
 156     col_loc = r.group(1)
 157     get_prefix = r.group(2)
 158     if args.r and not get_prefix:
 159         get_prefix = os.sep
 160
 161     # User asked to download the collection's manifest
 162     if not get_prefix:
 163         if not args.n:
 164             open_flags = os.O_CREAT | os.O_WRONLY
 165             if not args.f:
 166                 open_flags |= os.O_EXCL
 167             try:
 168                 if args.destination == "-":
 169                     write_block_or_manifest(
 170                         dest=stdout, src=col_loc,
 171                         api_client=api_client, args=args)
 172                 else:
 173                     out_fd = os.open(args.destination, open_flags)
 174                     with os.fdopen(out_fd, 'wb') as out_file:
 175                         write_block_or_manifest(
 176                             dest=out_file, src=col_loc,
 177                             api_client=api_client, args=args)
 178             except (IOError, OSError) as error:
 179                 logger.error("can't write to '{}': {}".format(args.destination, error))
 180                 return 1
 181             except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:
 182                 logger.error("failed to download '{}': {}".format(col_loc, error))
 183                 return 1
 184             except arvados.errors.ArgumentError as error:
 185                 if 'Argument to CollectionReader' in str(error):
 186                     logger.error("error reading collection: {}".format(error))
 187                     return 1
 188                 else:
 189                     raise
 190         return 0
 191
 192     try:
 193         reader = arvados.CollectionReader(
 194             col_loc, api_client=api_client, num_retries=args.retries)
 195     except Exception as error:
 196         logger.error("failed to read collection: {}".format(error))
 197         return 1
 198
 199     # Scan the collection. Make an array of (stream, file, local
 200     # destination filename) tuples, and add up total size to extract.
 201     todo = []
 202     todo_bytes = 0
 203     try:
 204         if get_prefix == os.sep:
 205             item = reader
 206         else:
 207             item = reader.find('.' + get_prefix)
 208
 209         if isinstance(item, arvados.collection.Subcollection) or isinstance(item, arvados.collection.CollectionReader):
 210             # If the user asked for a file and we got a subcollection, error out.
 211             if get_prefix[-1] != os.sep:
 212                 logger.error("requested file '{}' is in fact a subcollection. Append a trailing '/' to download it.".format('.' + get_prefix))
 213                 return 1
 214             # If the user asked stdout as a destination, error out.
 215             elif args.destination == '-':
 216                 logger.error("cannot use 'stdout' as destination when downloading multiple files.")
 217                 return 1
 218             # User asked for a subcollection, and that's what was found. Add up total size
 219             # to download.
 220             for s, f in files_in_collection(item):
 221                 dest_path = os.path.join(
 222                     args.destination,
 223                     os.path.join(s.stream_name(), f.name)[len(get_prefix)+1:])
 224                 if (not (args.n or args.f or args.skip_existing) and
 225                     os.path.exists(dest_path)):
 226                     logger.error('Local file %s already exists.' % (dest_path,))
 227                     return 1
 228                 todo += [(s, f, dest_path)]
 229                 todo_bytes += f.size()
 230         elif isinstance(item, arvados.arvfile.ArvadosFile):
 231             todo += [(item.parent, item, args.destination)]
 232             todo_bytes += item.size()
 233         else:
 234             logger.error("'{}' not found.".format('.' + get_prefix))
 235             return 1
 236     except (IOError, arvados.errors.NotFoundError) as e:
 237         logger.error(e)
 238         return 1
 239
 240     out_bytes = 0
 241     for s, f, outfilename in todo:
 242         outfile = None
 243         digestor = None
 244         if not args.n:
 245             if outfilename == "-":
 246                 outfile = stdout
 247             else:
 248                 if args.skip_existing and os.path.exists(outfilename):
 249                     logger.debug('Local file %s exists. Skipping.', outfilename)
 250                     continue
 251                 elif not args.f and (os.path.isfile(outfilename) or
 252                                    os.path.isdir(outfilename)):
 253                     # Good thing we looked again: apparently this file wasn't
 254                     # here yet when we checked earlier.
 255                     logger.error('Local file %s already exists.' % (outfilename,))
 256                     return 1
 257                 if args.r:
 258                     arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
 259                 try:
 260                     outfile = open(outfilename, 'wb')
 261                 except Exception as error:
 262                     logger.error('Open(%s) failed: %s' % (outfilename, error))
 263                     return 1
 264         if args.hash:
 265             digestor = hashlib.new(args.hash)
 266         try:
 267             with s.open(f.name, 'rb') as file_reader:
 268                 for data in file_reader.readall():
 269                     if outfile:
 270                         outfile.write(data)
 271                     if digestor:
 272                         digestor.update(data)
 273                     out_bytes += len(data)
 274                     if args.progress:
 275                         stderr.write('\r%d MiB / %d MiB %.1f%%' %
 276                                      (out_bytes >> 20,
 277                                       todo_bytes >> 20,
 278                                       (100
 279                                        if todo_bytes==0
 280                                        else 100.0*out_bytes/todo_bytes)))
 281                     elif args.batch_progress:
 282                         stderr.write('%s %d read %d total\n' %
 283                                      (sys.argv[0], os.getpid(),
 284                                       out_bytes, todo_bytes))
 285             if digestor:
 286                 stderr.write("%s  %s/%s\n"
 287                              % (digestor.hexdigest(), s.stream_name(), f.name))
 288         except KeyboardInterrupt:
 289             if outfile and (outfile.fileno() > 2) and not outfile.closed:
 290                 os.unlink(outfile.name)
 291             break
 292         finally:
 293             if outfile != None and outfile != stdout:
 294                 outfile.close()
 295
 296     if args.progress:
 297         stderr.write('\n')
 298     return 0
 299
 300 def files_in_collection(c):
 301     # Sort first by file type, then alphabetically by file path.
 302     for i in sorted(list(c.keys()),
 303                     key=lambda k: (
 304                         isinstance(c[k], arvados.collection.Subcollection),
 305                         k.upper())):
 306         if isinstance(c[i], arvados.arvfile.ArvadosFile):
 307             yield (c, c[i])
 308         elif isinstance(c[i], arvados.collection.Subcollection):
 309             for s, f in files_in_collection(c[i]):
 310                 yield (s, f)
 311
 312 def write_block_or_manifest(dest, src, api_client, args):
 313     if '+A' in src:
 314         # block locator
 315         kc = arvados.keep.KeepClient(api_client=api_client)
 316         dest.write(kc.get(src, num_retries=args.retries))
 317     else:
 318         # collection UUID or portable data hash
 319         reader = arvados.CollectionReader(
 320             src, api_client=api_client, num_retries=args.retries)
 321         dest.write(reader.manifest_text(strip=args.strip_manifest).encode())