git.arvados.org - arvados.git/blob - 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')