Merge branch 'master' into 2449-keep-flags
[arvados.git] / sdk / python / bin / arv-get
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 logger = logging.getLogger(os.path.basename(sys.argv[0]))
12
13 parser = argparse.ArgumentParser(
14     description='Copy data from Keep to a local file or pipe.')
15 parser.add_argument('locator', type=str,
16                     help="""
17 Collection locator, optionally with a file path or prefix.
18 """)
19 parser.add_argument('destination', type=str, nargs='?', default='/dev/stdout',
20                     help="""
21 Local file or directory where the data is to be written. Default:
22 /dev/stdout.
23 """)
24 group = parser.add_mutually_exclusive_group()
25 group.add_argument('--progress', action='store_true',
26                    help="""
27 Display human-readable progress on stderr (bytes and, if possible,
28 percentage of total data size). This is the default behavior when
29 stderr is a tty and stdout is not a tty.
30 """)
31 group.add_argument('--no-progress', action='store_true',
32                    help="""
33 Do not display human-readable progress on stderr.
34 """)
35 group.add_argument('--batch-progress', action='store_true',
36                    help="""
37 Display machine-readable progress on stderr (bytes and, if known,
38 total data size).
39 """)
40 group = parser.add_mutually_exclusive_group()
41 group.add_argument('--hash',
42                     help="""
43 Display the hash of each file as it is read from Keep, using the given
44 hash algorithm. Supported algorithms include md5, sha1, sha224,
45 sha256, sha384, and sha512.
46 """)
47 group.add_argument('--md5sum', action='store_const',
48                     dest='hash', const='md5',
49                     help="""
50 Display the MD5 hash of each file as it is read from Keep.
51 """)
52 parser.add_argument('-n', action='store_true',
53                     help="""
54 Do not write any data -- just read from Keep, and report md5sums if
55 requested.
56 """)
57 parser.add_argument('-r', action='store_true',
58                     help="""
59 Retrieve all files in the specified collection/prefix. This is the
60 default behavior if the "locator" argument ends with a forward slash.
61 """)
62 group = parser.add_mutually_exclusive_group()
63 group.add_argument('-f', action='store_true',
64                    help="""
65 Overwrite existing files while writing. The default behavior is to
66 refuse to write *anything* if any of the output files already
67 exist. As a special case, -f is not needed to write to /dev/stdout.
68 """)
69 group.add_argument('--skip-existing', action='store_true',
70                    help="""
71 Skip files that already exist. The default behavior is to refuse to
72 write *anything* if any files exist that would have to be
73 overwritten. This option causes even devices, sockets, and fifos to be
74 skipped.
75 """)
76
77 args = parser.parse_args()
78
79 if args.locator[-1] == os.sep:
80     args.r = True
81 if (args.r and
82     not args.n and
83     not (args.destination and
84          os.path.isdir(args.destination))):
85     parser.error('Destination is not a directory.')
86 if not args.r and (os.path.isdir(args.destination) or
87                    args.destination[-1] == os.path.sep):
88     args.destination = os.path.join(args.destination,
89                                     os.path.basename(args.locator))
90     logger.debug("Appended source file name to destination directory: %s" %
91                  args.destination)
92
93 # Turn on --progress by default if stderr is a tty and stdout isn't.
94 if (not (args.batch_progress or args.no_progress)
95     and os.isatty(sys.stderr.fileno())
96     and not os.isatty(sys.stdout.fileno())):
97     args.progress = True
98
99 if args.destination == '-':
100     args.destination = '/dev/stdout'
101 if args.destination == '/dev/stdout':
102     # Normally you have to use -f to write to a file (or device) that
103     # already exists, but "-" and "/dev/stdout" are common enough to
104     # merit a special exception.
105     args.f = True
106 else:
107     args.destination = args.destination.rstrip(os.sep)
108
109
110 import arvados
111
112 r = re.search(r'^(.*?)(/.*)?$', args.locator)
113 collection = r.group(1)
114 get_prefix = r.group(2)
115 if args.r and not get_prefix:
116     get_prefix = os.sep
117
118 todo = []
119 todo_bytes = 0
120 if not get_prefix:
121     try:
122         if not args.n:
123             if not args.f and os.path.exists(args.destination):
124                 logger.error('Local file %s already exists' % args.destination)
125                 sys.exit(1)
126             with open(args.destination, 'wb') as f:
127                 try:
128                     c = arvados.api('v1').collections().get(
129                         uuid=collection).execute()
130                     manifest = c['manifest_text']
131                 except Exception as e:
132                     logging.warning(
133                         "API lookup failed for collection %s (%s: %s)" %
134                         (collection, type(e), str(e)))
135                     manifest = arvados.Keep.get(collection)
136                 f.write(manifest)
137         sys.exit(0)
138     except arvados.errors.NotFoundError as e:
139         logger.error(e)
140         sys.exit(1)
141
142 reader = arvados.CollectionReader(collection)
143
144 # Scan the collection. Make an array of (stream, file, local
145 # destination filename) tuples, and add up total size to extract.
146
147 try:
148     for s in reader.all_streams():
149         for f in s.all_files():
150             if get_prefix and get_prefix[-1] == os.sep:
151                 if 0 != string.find(os.path.join(s.name(), f.name()),
152                                     '.' + get_prefix):
153                     continue
154                 dest_path = os.path.join(
155                     args.destination,
156                     os.path.join(s.name(), f.name())[len(get_prefix)+1:])
157                 if (not (args.n or args.f or args.skip_existing) and
158                     os.path.exists(dest_path)):
159                     logger.error('Local file %s already exists' % dest_path)
160                     sys.exit(1)
161             else:
162                 if os.path.join(s.name(), f.name()) != '.' + get_prefix:
163                     continue
164                 dest_path = args.destination
165             todo += [(s, f, dest_path)]
166             todo_bytes += f.size()
167 except arvados.errors.NotFoundError as e:
168     logger.error(e)
169     sys.exit(1)
170
171 # Read data, and (if not -n) write to local file(s) or pipe.
172
173 out_bytes = 0
174 for s,f,outfilename in todo:
175     outfile = None
176     digestor = None
177     if not args.n:
178         if args.skip_existing and os.path.exists(outfilename):
179             logger.debug('Local file %s exists. Skipping.' % outfilename)
180             continue
181         elif not args.f and (os.path.isfile(outfilename) or
182                            os.path.isdir(outfilename)):
183             # Good thing we looked again: apparently this file wasn't
184             # here yet when we checked earlier.
185             logger.error('Local file %s already exists' % outfilename)
186             sys.exit(1)
187         if args.r:
188             arvados.util.mkdir_dash_p(os.path.dirname(outfilename))
189         try:
190             outfile = open(outfilename, 'wb')
191         except Exception as e:
192             logger.error('Open(%s) failed: %s' % (outfilename, e))
193             sys.exit(1)
194     if args.hash:
195         digestor = hashlib.new(args.hash)
196     try:
197         for data in f.readall():
198             if outfile:
199                 outfile.write(data)
200             if digestor:
201                 digestor.update(data)
202             out_bytes += len(data)
203             if args.progress:
204                 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %
205                                  (out_bytes >> 20,
206                                   todo_bytes >> 20,
207                                   (100
208                                    if todo_bytes==0
209                                    else 100.0*out_bytes/todo_bytes)))
210             elif args.batch_progress:
211                 sys.stderr.write('%s %d read %d total\n' %
212                                  (sys.argv[0], os.getpid(),
213                                   out_bytes, todo_bytes))
214         if digestor:
215             sys.stderr.write("%s  %s/%s\n"
216                              % (digestor.hexdigest(), s.name(), f.name()))
217     except KeyboardInterrupt:
218         if outfile:
219             os.unlink(outfilename)
220         break
221
222 if args.progress:
223     sys.stderr.write('\n')