projects / arvados.git / blobdiff
summary | shortlog | log | commit | commitdiff | tree
raw | inline | side by side
Merge branch 'master' into 12033-multisite-search
[arvados.git] / sdk / python / arvados / commands / get.py
diff --git a/sdk/python/arvados/commands/get.py b/sdk/python/arvados/commands/get.py
index 67f38c4cbd635c64d4a0c7821ae50310b307de0c..56d471888d65a8f200bb6ae7af39991286be8aab 100755 (executable)
--- a/sdk/python/arvados/commands/get.py
+++ b/sdk/python/arvados/commands/get.py
@@ -1,4 +1,7 @@
 #!/usr/bin/env python
 #!/usr/bin/env python
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
 
 import argparse
 import hashlib
 
 import argparse
 import hashlib
@@ -10,6 +13,7 @@ import logging
 
 import arvados
 import arvados.commands._util as arv_cmd
 
 import arvados
 import arvados.commands._util as arv_cmd
+import arvados.util as util
 
 from arvados._version import __version__
 
 
 from arvados._version import __version__
 
@@ -84,6 +88,11 @@ write *anything* if any files exist that would have to be
 overwritten. This option causes even devices, sockets, and fifos to be
 skipped.
 """)
 overwritten. This option causes even devices, sockets, and fifos to be
 skipped.
 """)
+group.add_argument('--strip-manifest', action='store_true', default=False,
+                   help="""
+When getting a collection manifest, strip its access tokens before writing
+it.
+""")
 
 def parse_arguments(arguments, stdout, stderr):
     args = parser.parse_args(arguments)
 
 def parse_arguments(arguments, stdout, stderr):
     args = parser.parse_args(arguments)
@@ -125,18 +134,27 @@ def parse_arguments(arguments, stdout, stderr):
 
 def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
     global api_client
 
 def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
     global api_client
-    
+
+    if stdout is sys.stdout and hasattr(stdout, 'buffer'):
+        # in Python 3, write to stdout as binary
+        stdout = stdout.buffer
+
     args = parse_arguments(arguments, stdout, stderr)
     if api_client is None:
         api_client = arvados.api('v1')
 
     r = re.search(r'^(.*?)(/.*)?$', args.locator)
     args = parse_arguments(arguments, stdout, stderr)
     if api_client is None:
         api_client = arvados.api('v1')
 
     r = re.search(r'^(.*?)(/.*)?$', args.locator)
-    collection = r.group(1)
+    col_loc = r.group(1)
     get_prefix = r.group(2)
     if args.r and not get_prefix:
         get_prefix = os.sep
     get_prefix = r.group(2)
     if args.r and not get_prefix:
         get_prefix = os.sep
-    reader = arvados.CollectionReader(collection, num_retries=args.retries)
+    try:
+        reader = arvados.CollectionReader(col_loc, num_retries=args.retries)
+    except Exception as error:
+        logger.error("failed to read collection: {}".format(error))
+        return 1
 
 
+    # User asked to download the collection's manifest
     if not get_prefix:
         if not args.n:
             open_flags = os.O_CREAT | os.O_WRONLY
     if not get_prefix:
         if not args.n:
             open_flags = os.O_CREAT | os.O_WRONLY
@@ -144,16 +162,16 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
                 open_flags |= os.O_EXCL
             try:
                 if args.destination == "-":
                 open_flags |= os.O_EXCL
             try:
                 if args.destination == "-":
-                    stdout.write(reader.manifest_text())
+                    stdout.write(reader.manifest_text(strip=args.strip_manifest).encode())
                 else:
                     out_fd = os.open(args.destination, open_flags)
                     with os.fdopen(out_fd, 'wb') as out_file:
                 else:
                     out_fd = os.open(args.destination, open_flags)
                     with os.fdopen(out_fd, 'wb') as out_file:
-                        out_file.write(reader.manifest_text())
+                        out_file.write(reader.manifest_text(strip=args.strip_manifest).encode())
             except (IOError, OSError) as error:
                 logger.error("can't write to '{}': {}".format(args.destination, error))
                 return 1
             except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:
             except (IOError, OSError) as error:
                 logger.error("can't write to '{}': {}".format(args.destination, error))
                 return 1
             except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:
-                logger.error("failed to download '{}': {}".format(collection, error))
+                logger.error("failed to download '{}': {}".format(col_loc, error))
                 return 1
         return 0
 
                 return 1
         return 0
 
@@ -162,28 +180,39 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
     todo = []
     todo_bytes = 0
     try:
     todo = []
     todo_bytes = 0
     try:
-        for s, f in files_in_collection(reader):
-            if get_prefix and get_prefix[-1] == os.sep:
-                if 0 != string.find(os.path.join(s.stream_name(), f.name),
-                                    '.' + get_prefix):
-                    continue
-                if args.destination == "-":
-                    dest_path = "-"
-                else:
-                    dest_path = os.path.join(
-                        args.destination,
-                        os.path.join(s.stream_name(), f.name)[len(get_prefix)+1:])
-                    if (not (args.n or args.f or args.skip_existing) and
-                        os.path.exists(dest_path)):
-                        logger.error('Local file %s already exists.' % (dest_path,))
-                        return 1
-            else:
-                if os.path.join(s.stream_name(), f.name) != '.' + get_prefix:
-                    continue
-                dest_path = args.destination
-            todo += [(s, f, dest_path)]
-            todo_bytes += f.size()
-    except arvados.errors.NotFoundError as e:
+        if get_prefix == os.sep:
+            item = reader
+        else:
+            item = reader.find('.' + get_prefix)
+
+        if isinstance(item, arvados.collection.Subcollection) or isinstance(item, arvados.collection.CollectionReader):
+            # If the user asked for a file and we got a subcollection, error out.
+            if get_prefix[-1] != os.sep:
+                logger.error("requested file '{}' is in fact a subcollection. Append a trailing '/' to download it.".format('.' + get_prefix))
+                return 1
+            # If the user asked stdout as a destination, error out.
+            elif args.destination == '-':
+                logger.error("cannot use 'stdout' as destination when downloading multiple files.")
+                return 1
+            # User asked for a subcollection, and that's what was found. Add up total size
+            # to download.
+            for s, f in files_in_collection(item):
+                dest_path = os.path.join(
+                    args.destination,
+                    os.path.join(s.stream_name(), f.name)[len(get_prefix)+1:])
+                if (not (args.n or args.f or args.skip_existing) and
+                    os.path.exists(dest_path)):
+                    logger.error('Local file %s already exists.' % (dest_path,))
+                    return 1
+                todo += [(s, f, dest_path)]
+                todo_bytes += f.size()
+        elif isinstance(item, arvados.arvfile.ArvadosFile):
+            todo += [(item.parent, item, args.destination)]
+            todo_bytes += item.size()
+        else:
+            logger.error("'{}' not found.".format('.' + get_prefix))
+            return 1
+    except (IOError, arvados.errors.NotFoundError) as e:
         logger.error(e)
         return 1
 
         logger.error(e)
         return 1
 
@@ -214,7 +243,7 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
         if args.hash:
             digestor = hashlib.new(args.hash)
         try:
         if args.hash:
             digestor = hashlib.new(args.hash)
         try:
-            with s.open(f.name, 'r') as file_reader:
+            with s.open(f.name, 'rb') as file_reader:
                 for data in file_reader.readall():
                     if outfile:
                         outfile.write(data)
                 for data in file_reader.readall():
                     if outfile:
                         outfile.write(data)
@@ -239,13 +268,17 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr):
             if outfile and (outfile.fileno() > 2) and not outfile.closed:
                 os.unlink(outfile.name)
             break
             if outfile and (outfile.fileno() > 2) and not outfile.closed:
                 os.unlink(outfile.name)
             break
+        finally:
+            if outfile != None and outfile != stdout:
+                outfile.close()
 
     if args.progress:
         stderr.write('\n')
 
     if args.progress:
         stderr.write('\n')
+    return 0
 
 def files_in_collection(c):
     # Sort first by file type, then alphabetically by file path.
 
 def files_in_collection(c):
     # Sort first by file type, then alphabetically by file path.
-    for i in sorted(c.keys(),
+    for i in sorted(list(c.keys()),
                     key=lambda k: (
                         isinstance(c[k], arvados.collection.Subcollection),
                         k.upper())):
                     key=lambda k: (
                         isinstance(c[k], arvados.collection.Subcollection),
                         k.upper())):