Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22""" 

23Utilities for safe file IO 

24""" 

25 

26__all__ = ("DoNotWrite", 

27 "safeMakeDir", 

28 "setFileMode", 

29 "FileForWriteOnceCompareSameFailure", 

30 "FileForWriteOnceCompareSame", 

31 "SafeFile", 

32 "SafeFilename", 

33 "SafeLockedFileForRead", 

34 "SafeLockedFileForWrite") 

35 

36from contextlib import contextmanager 

37import errno 

38import fcntl 

39import filecmp 

40import os 

41import tempfile 

42import logging 

43 

44 

45class DoNotWrite(RuntimeError): 

46 pass 

47 

48 

49def safeMakeDir(directory): 

50 """Make a directory in a manner avoiding race conditions""" 

51 if directory != "" and not os.path.exists(directory): 

52 try: 

53 os.makedirs(directory) 

54 except OSError as e: 

55 # Don't fail if directory exists due to race 

56 if e.errno != errno.EEXIST: 

57 raise e 

58 

59 

60def setFileMode(filename): 

61 """Set a file mode according to the user's umask""" 

62 # Get the current umask, which we can only do by setting it and then 

63 # reverting to the original. 

64 umask = os.umask(0o077) 

65 os.umask(umask) 

66 # chmod the new file to match what it would have been if it hadn't started 

67 # life as a temporary file (which have more restricted permissions). 

68 os.chmod(filename, (~umask & 0o666)) 

69 

70 

71class FileForWriteOnceCompareSameFailure(RuntimeError): 

72 pass 

73 

74 

75@contextmanager 

76def FileForWriteOnceCompareSame(name): 

77 """Context manager to get a file that can be written only once and all 

78 other writes will succeed only if they match the initial write. 

79 

80 The context manager provides a temporary file object. After the user is 

81 done, the temporary file becomes the permanent file if the file at name 

82 does not already exist. If the file at name does exist the temporary file 

83 is compared to the file at name. If they are the same then this is good 

84 and the temp file is silently thrown away. If they are not the same then 

85 a runtime error is raised. 

86 """ 

87 outDir, outName = os.path.split(name) 

88 safeMakeDir(outDir) 

89 temp = tempfile.NamedTemporaryFile(mode="w", dir=outDir, prefix=outName, delete=False) 

90 try: 

91 yield temp 

92 finally: 

93 try: 

94 temp.close() 

95 # If the symlink cannot be created then it will raise. If it can't 

96 # be created because a file at "name" already exists then we"ll 

97 # do a compare-same check. 

98 os.symlink(temp.name, name) 

99 # If the symlink was created then this is the process that created 

100 # the first instance of the file, and we know its contents match. 

101 # Move the temp file over the symlink. 

102 os.rename(temp.name, name) 

103 # At this point, we know the file has just been created. Set 

104 # permissions according to the current umask. 

105 setFileMode(name) 

106 except OSError as e: 

107 if e.errno != errno.EEXIST: 

108 raise e 

109 filesMatch = filecmp.cmp(temp.name, name, shallow=False) 

110 os.remove(temp.name) 

111 if filesMatch: 

112 # if the files match then the compare-same check succeeded and 

113 # we can silently return. 

114 return 

115 else: 

116 # if the files do not match then the calling code was trying 

117 # to write a non-matching file over the previous file, maybe 

118 # it's a race condition? In any event, raise a runtime error. 

119 raise FileForWriteOnceCompareSameFailure("Written file does not match existing file.") 

120 

121 

122@contextmanager 

123def SafeFile(name): 

124 """Context manager to create a file in a manner avoiding race conditions 

125 

126 The context manager provides a temporary file object. After the user is 

127 done, we move that file into the desired place and close the fd to avoid 

128 resource leakage. 

129 """ 

130 outDir, outName = os.path.split(name) 

131 safeMakeDir(outDir) 

132 doWrite = True 

133 with tempfile.NamedTemporaryFile(mode="w", dir=outDir, prefix=outName, delete=False) as temp: 

134 try: 

135 yield temp 

136 except DoNotWrite: 

137 doWrite = False 

138 finally: 

139 if doWrite: 

140 os.rename(temp.name, name) 

141 setFileMode(name) 

142 

143 

144@contextmanager 

145def SafeFilename(name): 

146 """Context manager for creating a file in a manner avoiding race 

147 conditions. 

148 

149 The context manager provides a temporary filename with no open file 

150 descriptors (as this can cause trouble on some systems). After the user is 

151 done, we move the file into the desired place. 

152 """ 

153 outDir, outName = os.path.split(name) 

154 safeMakeDir(outDir) 

155 temp = tempfile.NamedTemporaryFile(mode="w", dir=outDir, prefix=outName, delete=False) 

156 tempName = temp.name 

157 temp.close() # We don't use the fd, just want a filename 

158 try: 

159 yield tempName 

160 finally: 

161 os.rename(tempName, name) 

162 setFileMode(name) 

163 

164 

165@contextmanager 

166def SafeLockedFileForRead(name): 

167 """Context manager for reading a file that may be locked with an exclusive 

168 lock via SafeLockedFileForWrite. 

169 

170 This will first acquire a shared lock before returning the file. When 

171 the file is closed the shared lock will be unlocked. 

172 

173 Parameters 

174 ---------- 

175 name : string 

176 The file name to be opened, may include path. 

177 

178 Yields 

179 ------ 

180 file object 

181 The file to be read from. 

182 """ 

183 log = logging.getLogger("daf.butler") 

184 try: 

185 with open(name, "r") as f: 

186 log.debug("Acquiring shared lock on {}".format(name)) 

187 fcntl.flock(f, fcntl.LOCK_SH) 

188 log.debug("Acquired shared lock on {}".format(name)) 

189 yield f 

190 finally: 

191 log.debug("Releasing shared lock on {}".format(name)) 

192 

193 

194class SafeLockedFileForWrite: 

195 """File-like object that is used to create a file if needed, lock it with 

196 an exclusive lock, and contain file descriptors to readable and writable 

197 versions of the file. 

198 

199 This will only open a file descriptor in "write" mode if a write operation 

200 is performed. If no write operation is performed, the existing file (if 

201 there is one) will not be overwritten. 

202 

203 Contains __enter__ and __exit__ functions so this can be used by a 

204 context manager. 

205 """ 

206 

207 def __init__(self, name): 

208 self.log = logging.getLogger("daf.butler") 

209 self.name = name 

210 self._readable = None 

211 self._writeable = None 

212 safeMakeDir(os.path.split(name)[0]) 

213 

214 def __enter__(self): 

215 self.open() 

216 return self 

217 

218 def __exit__(self, type, value, traceback): 

219 self.close() 

220 

221 def open(self): 

222 self._fileHandle = open(self.name, "a") 

223 self.log.debug("Acquiring exclusive lock on {}".format(self.name)) 

224 fcntl.flock(self._fileHandle, fcntl.LOCK_EX) 

225 self.log.debug("Acquired exclusive lock on {}".format(self.name)) 

226 

227 def close(self): 

228 self.log.debug("Releasing exclusive lock on {}".format(self.name)) 

229 if self._writeable is not None: 

230 self._writeable.close() 

231 if self._readable is not None: 

232 self._readable.close() 

233 self._fileHandle.close() 

234 

235 @property 

236 def readable(self): 

237 if self._readable is None: 

238 self._readable = open(self.name, "r") 

239 return self._readable 

240 

241 @property 

242 def writeable(self): 

243 if self._writeable is None: 

244 self._writeable = open(self.name, "w") 

245 return self._writeable 

246 

247 def read(self, size=None): 

248 if size is not None: 

249 return self.readable.read(size) 

250 return self.readable.read() 

251 

252 def write(self, str): 

253 self.writeable.write(str)