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