arvados.commands.get

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