Coverage for python/lsst/daf/butler/core/safeFileIo.py : 22%

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/>.
22"""
23Utilities for safe file IO
24"""
26__all__ = ("DoNotWrite",
27 "safeMakeDir",
28 "setFileMode",
29 "FileForWriteOnceCompareSameFailure",
30 "FileForWriteOnceCompareSame",
31 "SafeFile",
32 "SafeFilename",
33 "SafeLockedFileForRead",
34 "SafeLockedFileForWrite")
36from contextlib import contextmanager
37import errno
38import fcntl
39import filecmp
40import os
41import tempfile
42import logging
45class DoNotWrite(RuntimeError):
46 pass
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
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))
71class FileForWriteOnceCompareSameFailure(RuntimeError):
72 pass
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.
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.")
122@contextmanager
123def SafeFile(name):
124 """Context manager to create a file in a manner avoiding race conditions
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)
144@contextmanager
145def SafeFilename(name):
146 """Context manager for creating a file in a manner avoiding race
147 conditions.
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)
165@contextmanager
166def SafeLockedFileForRead(name):
167 """Context manager for reading a file that may be locked with an exclusive
168 lock via SafeLockedFileForWrite.
170 This will first acquire a shared lock before returning the file. When
171 the file is closed the shared lock will be unlocked.
173 Parameters
174 ----------
175 name : string
176 The file name to be opened, may include path.
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))
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.
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.
203 Contains __enter__ and __exit__ functions so this can be used by a
204 context manager.
205 """
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])
214 def __enter__(self):
215 self.open()
216 return self
218 def __exit__(self, type, value, traceback):
219 self.close()
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))
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()
235 @property
236 def readable(self):
237 if self._readable is None:
238 self._readable = open(self.name, "r")
239 return self._readable
241 @property
242 def writeable(self):
243 if self._writeable is None:
244 self._writeable = open(self.name, "w")
245 return self._writeable
247 def read(self, size=None):
248 if size is not None:
249 return self.readable.read(size)
250 return self.readable.read()
252 def write(self, str):
253 self.writeable.write(str)