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 import arvados.commands._util as arv_cmd

13

14 from arvados._version import __version__

15

16 logger = logging.getLogger('arvados.arv-get')

17

18 def abort(msg, code=1):

19 print >>sys.stderr, "arv-get:", msg

20 exit(code)

21

22 parser = argparse.ArgumentParser(

23 description='Copy data from Keep to a local file or pipe.',

24 parents=[arv_cmd.retry_opt])

25 parser.add_argument('--version', action='version',

26 version="%s %s" % (sys.argv[0], __version__),

27 help='Print version and exit.')

28 parser.add_argument('locator', type=str,

29 help="""

30 Collection locator, optionally with a file path or prefix.

31 """)

32 parser.add_argument('destination', type=str, nargs='?', default='-',

33 help="""

34 Local file or directory where the data is to be written. Default: stdout.

35 """)

36 group = parser.add_mutually_exclusive_group()

37 group.add_argument('--progress', action='store_true',

38 help="""

39 Display human-readable progress on stderr (bytes and, if possible,

40 percentage of total data size). This is the default behavior when it

41 is not expected to interfere with the output: specifically, stderr is

42 a tty _and_ either stdout is not a tty, or output is being written to

43 named files rather than stdout.

44 """)

45 group.add_argument('--no-progress', action='store_true',

46 help="""

47 Do not display human-readable progress on stderr.

48 """)

49 group.add_argument('--batch-progress', action='store_true',

50 help="""

51 Display machine-readable progress on stderr (bytes and, if known,

52 total data size).

53 """)

54 group = parser.add_mutually_exclusive_group()

55 group.add_argument('--hash',

56 help="""

57 Display the hash of each file as it is read from Keep, using the given

58 hash algorithm. Supported algorithms include md5, sha1, sha224,

59 sha256, sha384, and sha512.

60 """)

61 group.add_argument('--md5sum', action='store_const',

62 dest='hash', const='md5',

63 help="""

64 Display the MD5 hash of each file as it is read from Keep.

65 """)

66 parser.add_argument('-n', action='store_true',

67 help="""

68 Do not write any data -- just read from Keep, and report md5sums if

69 requested.

70 """)

71 parser.add_argument('-r', action='store_true',

72 help="""

73 Retrieve all files in the specified collection/prefix. This is the

74 default behavior if the "locator" argument ends with a forward slash.

75 """)

76 group = parser.add_mutually_exclusive_group()

77 group.add_argument('-f', action='store_true',

78 help="""

79 Overwrite existing files while writing. The default behavior is to

80 refuse to write *anything* if any of the output files already

81 exist. As a special case, -f is not needed to write to stdout.

82 """)

83 group.add_argument('--skip-existing', action='store_true',

84 help="""

85 Skip files that already exist. The default behavior is to refuse to

86 write *anything* if any files exist that would have to be

87 overwritten. This option causes even devices, sockets, and fifos to be

88 skipped.

89 """)

90

91 args = parser.parse_args()

92

93 if args.locator[-1] == os.sep:

94 args.r = True

95 if (args.r and

96 not args.n and

97 not (args.destination and

98 os.path.isdir(args.destination))):

99 parser.error('Destination is not a directory.')

100 if not args.r and (os.path.isdir(args.destination) or

101 args.destination[-1] == os.path.sep):

102 args.destination = os.path.join(args.destination,

103 os.path.basename(args.locator))

104 logger.debug("Appended source file name to destination directory: %s",

105 args.destination)

106

107 if args.destination == '/dev/stdout':

108 args.destination = "-"

109

110 if args.destination == '-':

111 # Normally you have to use -f to write to a file (or device) that

112 # already exists, but "-" and "/dev/stdout" are common enough to

113 # merit a special exception.

114 args.f = True

115 else:

116 args.destination = args.destination.rstrip(os.sep)

117

118 # Turn on --progress by default if stderr is a tty and output is

119 # either going to a named file, or going (via stdout) to something

120 # that isn't a tty.

121 if (not (args.batch_progress or args.no_progress)

122 and sys.stderr.isatty()

123 and (args.destination != '-'

124 or not sys.stdout.isatty())):

125 args.progress = True

126

127

128 r = re.search(r'^(.*?)(/.*)?$', args.locator)

129 collection = r.group(1)

130 get_prefix = r.group(2)

131 if args.r and not get_prefix:

132 get_prefix = os.sep

133 api_client = arvados.api('v1')

134 reader = arvados.CollectionReader(collection, num_retries=args.retries)

135

136 if not get_prefix:

137 if not args.n:

138 open_flags = os.O_CREAT | os.O_WRONLY

139 if not args.f:

140 open_flags |= os.O_EXCL

141 try:

142 if args.destination == "-":

143 sys.stdout.write(reader.manifest_text())

144 else:

145 out_fd = os.open(args.destination, open_flags)

146 with os.fdopen(out_fd, 'wb') as out_file:

147 out_file.write(reader.manifest_text())

148 except (IOError, OSError) as error:

149 abort("can't write to '{}': {}".format(args.destination, error))

150 except (arvados.errors.ApiError, arvados.errors.KeepReadError) as error:

151 abort("failed to download '{}': {}".format(collection, error))

152 sys.exit(0)

153

154 reader.normalize()

155

156 # Scan the collection. Make an array of (stream, file, local

157 # destination filename) tuples, and add up total size to extract.

158 todo = []

159 todo_bytes = 0

160 try:

161 for s in reader.all_streams():

162 for f in s.all_files():

163 if get_prefix and get_prefix[-1] == os.sep:

164 if 0 != string.find(os.path.join(s.name(), f.name()),

165 '.' + get_prefix):

166 continue

167 if args.destination == "-":

168 dest_path = "-"

169 else:

170 dest_path = os.path.join(

171 args.destination,

172 os.path.join(s.name(), f.name())[len(get_prefix)+1:])

173 if (not (args.n or args.f or args.skip_existing) and

174 os.path.exists(dest_path)):

175 abort('Local file %s already exists.' % (dest_path,))

176 else:

177 if os.path.join(s.name(), f.name()) != '.' + get_prefix:

178 continue

179 dest_path = args.destination

180 todo += [(s, f, dest_path)]

181 todo_bytes += f.size()

182 except arvados.errors.NotFoundError as e:

183 abort(e)

184

185 # Read data, and (if not -n) write to local file(s) or pipe.

186

187 out_bytes = 0

188 for s,f,outfilename in todo:

189 outfile = None

190 digestor = None

191 if not args.n:

192 if outfilename == "-":

193 outfile = sys.stdout

194 else:

195 if args.skip_existing and os.path.exists(outfilename):

196 logger.debug('Local file %s exists. Skipping.', outfilename)

197 continue

198 elif not args.f and (os.path.isfile(outfilename) or

199 os.path.isdir(outfilename)):

200 # Good thing we looked again: apparently this file wasn't

201 # here yet when we checked earlier.

202 abort('Local file %s already exists.' % (outfilename,))

203 if args.r:

204 arvados.util.mkdir_dash_p(os.path.dirname(outfilename))

205 try:

206 outfile = open(outfilename, 'wb')

207 except Exception as error:

208 abort('Open(%s) failed: %s' % (outfilename, error))

209 if args.hash:

210 digestor = hashlib.new(args.hash)

211 try:

212 for data in f.readall():

213 if outfile:

214 outfile.write(data)

215 if digestor:

216 digestor.update(data)

217 out_bytes += len(data)

218 if args.progress:

219 sys.stderr.write('\r%d MiB / %d MiB %.1f%%' %

220 (out_bytes >> 20,

221 todo_bytes >> 20,

222 (100

223 if todo_bytes==0

224 else 100.0*out_bytes/todo_bytes)))

225 elif args.batch_progress:

226 sys.stderr.write('%s %d read %d total\n' %

227 (sys.argv[0], os.getpid(),

228 out_bytes, todo_bytes))

229 if digestor:

230 sys.stderr.write("%s %s/%s\n"

231 % (digestor.hexdigest(), s.name(), f.name()))

232 except KeyboardInterrupt:

233 if outfile and (outfile.fileno() > 2) and not outfile.closed:

234 os.unlink(outfile.name)

235 break

236

237 if args.progress:

238 sys.stderr.write('\n')