Coverage for python/lsst/daf/persistence/safeFileIo.py: 27%

117 statements  

« prev     ^ index     » next       coverage.py v6.4.4, created at 2022-09-07 09:51 +0000

1# 

2# LSST Data Management System 

3# 

4# Copyright 2008-2015 AURA/LSST. 

5# 

6# This product includes software developed by the 

7# LSST Project (http://www.lsst.org/). 

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 LSST License Statement and 

20# the GNU General Public License along with this program. If not, 

21# see <https://www.lsstcorp.org/LegalNotices/>. 

22# 

23""" 

24Utilities for safe file IO 

25""" 

26from contextlib import contextmanager 

27import errno 

28import fcntl 

29import filecmp 

30import os 

31import tempfile 

32from lsst.log import Log 

33 

34 

35class DoNotWrite(RuntimeError): 

36 pass 

37 

38 

39def safeMakeDir(directory): 

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

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

42 try: 

43 os.makedirs(directory) 

44 except OSError as e: 

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

46 if e.errno != errno.EEXIST: 

47 raise e 

48 

49 

50def setFileMode(filename): 

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

52 # Get the current umask, which we can only do by setting it and then reverting to the original. 

53 umask = os.umask(0o077) 

54 os.umask(umask) 

55 # chmod the new file to match what it would have been if it hadn't started life as a temporary 

56 # file (which have more restricted permissions). 

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

58 

59 

60class FileForWriteOnceCompareSameFailure(RuntimeError): 

61 pass 

62 

63 

64@contextmanager 

65def FileForWriteOnceCompareSame(name): 

66 """Context manager to get a file that can be written only once and all other writes will succeed only if 

67 they match the initial write. 

68 

69 The context manager provides a temporary file object. After the user is done, the temporary file becomes 

70 the permanent file if the file at name does not already exist. If the file at name does exist the 

71 temporary file is compared to the file at name. If they are the same then this is good and the temp file 

72 is silently thrown away. If they are not the same then a runtime error is raised. 

73 """ 

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

75 safeMakeDir(outDir) 

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

77 try: 

78 yield temp 

79 finally: 

80 try: 

81 temp.close() 

82 # If the symlink cannot be created then it will raise. If it can't be created because a file at 

83 # 'name' already exists then we'll do a compare-same check. 

84 os.symlink(temp.name, name) 

85 # If the symlink was created then this is the process that created the first instance of the 

86 # file, and we know its contents match. Move the temp file over the symlink. 

87 os.rename(temp.name, name) 

88 # At this point, we know the file has just been created. Set permissions according to the 

89 # current umask. 

90 setFileMode(name) 

91 except OSError as e: 

92 if e.errno != errno.EEXIST: 

93 raise e 

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

95 os.remove(temp.name) 

96 if filesMatch: 

97 # if the files match then the compare-same check succeeded and we can silently return. 

98 return 

99 else: 

100 # if the files do not match then the calling code was trying to write a non-matching file over 

101 # the previous file, maybe it's a race condition? In any event, raise a runtime error. 

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

103 

104 

105@contextmanager 

106def SafeFile(name): 

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

108 

109 The context manager provides a temporary file object. After the user is done, 

110 we move that file into the desired place and close the fd to avoid resource 

111 leakage. 

112 """ 

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

114 safeMakeDir(outDir) 

115 doWrite = True 

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

117 try: 

118 yield temp 

119 except DoNotWrite: 

120 doWrite = False 

121 finally: 

122 if doWrite: 

123 os.rename(temp.name, name) 

124 setFileMode(name) 

125 

126 

127@contextmanager 

128def SafeFilename(name): 

129 """Context manager for creating a file in a manner avoiding race conditions 

130 

131 The context manager provides a temporary filename with no open file descriptors 

132 (as this can cause trouble on some systems). After the user is done, we move the 

133 file into the desired place. 

134 """ 

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

136 safeMakeDir(outDir) 

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

138 tempName = temp.name 

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

140 try: 

141 yield tempName 

142 finally: 

143 os.rename(tempName, name) 

144 setFileMode(name) 

145 

146 

147@contextmanager 

148def SafeLockedFileForRead(name): 

149 """Context manager for reading a file that may be locked with an exclusive lock via 

150 SafeLockedFileForWrite. This will first acquire a shared lock before returning the file. When the file is 

151 closed the shared lock will be unlocked. 

152 

153 Parameters 

154 ---------- 

155 name : string 

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

157 

158 Yields 

159 ------ 

160 file object 

161 The file to be read from. 

162 """ 

163 log = Log.getLogger("daf.persistence.butler") 

164 try: 

165 with open(name, 'r') as f: 

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

167 fcntl.flock(f, fcntl.LOCK_SH) 

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

169 yield f 

170 finally: 

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

172 

173 

174class SafeLockedFileForWrite: 

175 """File-like object that is used to create a file if needed, lock it with an exclusive lock, and contain 

176 file descriptors to readable and writable versions of the file. 

177 

178 This will only open a file descriptor in 'write' mode if a write operation is performed. If no write 

179 operation is performed, the existing file (if there is one) will not be overwritten. 

180 

181 Contains __enter__ and __exit__ functions so this can be used by a context manager. 

182 """ 

183 def __init__(self, name): 

184 self.log = Log.getLogger("daf.persistence.butler") 

185 self.name = name 

186 self._readable = None 

187 self._writeable = None 

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

189 

190 def __enter__(self): 

191 self.open() 

192 return self 

193 

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

195 self.close() 

196 

197 def open(self): 

198 self._fileHandle = open(self.name, 'a') 

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

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

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

202 

203 def close(self): 

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

205 if self._writeable is not None: 

206 self._writeable.close() 

207 if self._readable is not None: 

208 self._readable.close() 

209 self._fileHandle.close() 

210 

211 @property 

212 def readable(self): 

213 if self._readable is None: 

214 self._readable = open(self.name, 'r') 

215 return self._readable 

216 

217 @property 

218 def writeable(self): 

219 if self._writeable is None: 

220 self._writeable = open(self.name, 'w') 

221 return self._writeable 

222 

223 def read(self, size=None): 

224 if size is not None: 

225 return self.readable.read(size) 

226 return self.readable.read() 

227 

228 def write(self, str): 

229 self.writeable.write(str)