SCons/dblite.py

   1 # MIT License
   2 #
   3 # Copyright The SCons Foundation
   4 #
   5 # Permission is hereby granted, free of charge, to any person obtaining
   6 # a copy of this software and associated documentation files (the
   7 # "Software"), to deal in the Software without restriction, including
   8 # without limitation the rights to use, copy, modify, merge, publish,
   9 # distribute, sublicense, and/or sell copies of the Software, and to
  10 # permit persons to whom the Software is furnished to do so, subject to
  11 # the following conditions:
  12 #
  13 # The above copyright notice and this permission notice shall be included
  14 # in all copies or substantial portions of the Software.
  15 #
  16 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
  17 # KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
  18 # WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  19 # NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
  20 # LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
  21 # OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
  22 # WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  23
  24 """
  25 dblite.py module contributed by Ralf W. Grosse-Kunstleve.
  26 Extended for Unicode by Steven Knight.
  27
  28 This is a very simple-minded "database" used for saved signature
  29 information, with an interface modeled on the Python dbm database
  30 interface module.
  31 """
  32
  33 import io
  34 import os
  35 import pickle
  36 import shutil
  37 import time
  38
  39 from SCons.compat import PICKLE_PROTOCOL
  40
  41 KEEP_ALL_FILES = False
  42 IGNORE_CORRUPT_DBFILES = False
  43
  44
  45 def corruption_warning(filename) -> None:
  46     """Local warning for corrupt db.
  47
  48     Used for self-tests. SCons overwrites this with a
  49     different warning function in SConsign.py.
  50     """
  51     print("Warning: Discarding corrupt database:", filename)
  52
  53
  54 DBLITE_SUFFIX = ".dblite"
  55 TMP_SUFFIX = ".tmp"
  56
  57
  58 class _Dblite:
  59     """Lightweight signature database class.
  60
  61     Behaves like a dict when in memory, loads from a pickled disk
  62     file on open and writes back out to it on close.
  63
  64     Open the database file using a path derived from *file_base_name*.
  65     The optional *flag* argument can be:
  66
  67     +---------+---------------------------------------------------+
  68     | Value   | Meaning                                           |
  69     +=========+===================================================+
  70     | ``'r'`` | Open existing database for reading only (default) |
  71     +---------+---------------------------------------------------+
  72     | ``'w'`` | Open existing database for reading and  writing   |
  73     +---------+---------------------------------------------------+
  74     | ``'c'`` | Open database for reading and writing, creating   |
  75     |         | it if it doesn't exist                            |
  76     +---------+---------------------------------------------------+
  77     | ``'n'`` | Always create a new, empty database, open for     |
  78     |         | reading and writing                               |
  79     +---------+---------------------------------------------------+
  80
  81     The optional *mode* argument is the POSIX mode of the file, used only
  82     when the database has to be created.  It defaults to octal ``0o666``.
  83     """
  84
  85     # Because open() is defined at module level, overwriting builtin open
  86     # in the scope of this module, we use io.open to avoid ambiguity.
  87     _open = staticmethod(io.open)
  88
  89     # we need to squirrel away references to functions from various modules
  90     # that we'll use when sync() is called: this may happen at Python
  91     # teardown time (we call it from our __del__), and the global module
  92     # references themselves may already have been rebound to None.
  93     _pickle_dump = staticmethod(pickle.dump)
  94     _pickle_protocol = PICKLE_PROTOCOL
  95     try:
  96         _os_chown = staticmethod(os.chown)
  97     except AttributeError:
  98         _os_chown = None
  99     _os_replace = staticmethod(os.replace)
 100     _os_chmod = staticmethod(os.chmod)
 101     _shutil_copyfile = staticmethod(shutil.copyfile)
 102     _time_time = staticmethod(time.time)
 103
 104     def __init__(self, file_base_name, flag='r', mode=0o666) -> None:
 105         assert flag in ("r", "w", "c", "n")
 106
 107         base, ext = os.path.splitext(file_base_name)
 108         if ext == DBLITE_SUFFIX:
 109             # There's already a suffix on the file name, don't add one.
 110             self._file_name = file_base_name
 111             self._tmp_name = base + TMP_SUFFIX
 112         else:
 113             self._file_name = file_base_name + DBLITE_SUFFIX
 114             self._tmp_name = file_base_name + TMP_SUFFIX
 115
 116         self._flag = flag
 117         self._mode = mode
 118         self._dict = {}
 119         self._needs_sync = False
 120
 121         if self._os_chown is not None and 0 in (os.geteuid(), os.getegid()):
 122             # running as root; chown back to current owner/group when done
 123             try:
 124                 statinfo = os.stat(self._file_name)
 125                 self._chown_to = statinfo.st_uid
 126                 self._chgrp_to = statinfo.st_gid
 127             except OSError:
 128                 # db file doesn't exist yet.
 129                 # Check os.environ for SUDO_UID, use if set
 130                 self._chown_to = int(os.environ.get('SUDO_UID', -1))
 131                 self._chgrp_to = int(os.environ.get('SUDO_GID', -1))
 132         else:
 133             self._chown_to = -1  # don't chown
 134             self._chgrp_to = -1  # don't chgrp
 135
 136         if self._flag == "n":
 137             with io.open(self._file_name, "wb", opener=self.opener):
 138                 return  # just make sure it exists
 139         else:
 140             # We only need the disk file to slurp in the data.  Updates are
 141             # handled on close, db is mainained only in memory until then.
 142             try:
 143                 with io.open(self._file_name, "rb") as f:
 144                     p = f.read()
 145             except OSError as e:
 146                 # an error for file not to exist, unless flag is create
 147                 if self._flag != "c":
 148                     raise e
 149                 with io.open(self._file_name, "wb", opener=self.opener):
 150                     return  # just make sure it exists
 151             if len(p) > 0:
 152                 try:
 153                     self._dict = pickle.loads(p, encoding='bytes')
 154                 except (
 155                     pickle.UnpicklingError,
 156                     # Python3 docs:
 157                     # Note that other exceptions may also be raised during
 158                     # unpickling, including (but not necessarily limited to)
 159                     # AttributeError, EOFError, ImportError, and IndexError.
 160                     AttributeError,
 161                     EOFError,
 162                     ImportError,
 163                     IndexError,
 164                 ):
 165                     if IGNORE_CORRUPT_DBFILES:
 166                         corruption_warning(self._file_name)
 167                     else:
 168                         raise
 169
 170     def opener(self, path, flags):
 171         """Database open helper when creation may be needed.
 172
 173         The high-level Python open() function cannot specify a file mode
 174         for creation. Using this as the opener with the saved mode lets
 175         us do that.
 176         """
 177         return os.open(path, flags, mode=self._mode)
 178
 179     def close(self) -> None:
 180         if self._needs_sync:
 181             self.sync()
 182
 183     def __del__(self) -> None:
 184         self.close()
 185
 186     def sync(self) -> None:
 187         """Flush the database to disk.
 188
 189         This routine *must* succeed, since the in-memory and on-disk
 190         copies are out of sync as soon as we do anything that changes
 191         the in-memory version. Thus, to be cautious, flush to a
 192         temporary file and then move it over with some error handling.
 193         """
 194         self._check_writable()
 195         with self._open(self._tmp_name, "wb", opener=self.opener) as f:
 196             self._pickle_dump(self._dict, f, self._pickle_protocol)
 197
 198         try:
 199             self._os_replace(self._tmp_name, self._file_name)
 200         except PermissionError:
 201             # If we couldn't replace due to perms, try to change and retry.
 202             # This is mainly for Windows - on POSIX the file permissions
 203             # don't matter, the os.replace would have worked anyway.
 204             # We're giving up if the retry fails, just let the Python
 205             # exception abort us.
 206             try:
 207                 self._os_chmod(self._file_name, 0o777)
 208             except PermissionError:
 209                 pass
 210             self._os_replace(self._tmp_name, self._file_name)
 211
 212         if (
 213             self._os_chown is not None and self._chown_to > 0
 214         ):  # don't chown to root or -1
 215             try:
 216                 self._os_chown(self._file_name, self._chown_to, self._chgrp_to)
 217             except OSError:
 218                 pass
 219
 220         self._needs_sync = False
 221         if KEEP_ALL_FILES:
 222             self._shutil_copyfile(
 223                 self._file_name, f"{self._file_name}_{int(self._time_time())}"
 224             )
 225
 226     def _check_writable(self):
 227         if self._flag == "r":
 228             raise OSError(f"Read-only database: {self._file_name}")
 229
 230     def __getitem__(self, key):
 231         return self._dict[key]
 232
 233     def __setitem__(self, key, value):
 234         self._check_writable()
 235
 236         if not isinstance(key, str):
 237             raise TypeError(f"key `{key}' must be a string but is {type(key)}")
 238
 239         if not isinstance(value, bytes):
 240             raise TypeError(f"value `{value}' must be bytes but is {type(value)}")
 241
 242         self._dict[key] = value
 243         self._needs_sync = True
 244
 245     def __delitem__(self, key):
 246         del self._dict[key]
 247
 248     def keys(self):
 249         return self._dict.keys()
 250
 251     def items(self):
 252         return self._dict.items()
 253
 254     def values(self):
 255         return self._dict.values()
 256
 257     __iter__ = keys
 258
 259     def __contains__(self, key) -> bool:
 260         return key in self._dict
 261
 262     def __len__(self) -> int:
 263         return len(self._dict)
 264
 265
 266 def open(file, flag="r", mode: int = 0o666):  # pylint: disable=redefined-builtin
 267     return _Dblite(file, flag, mode)
 268
 269
 270 def _exercise():
 271     db = open("tmp", "n")
 272     assert len(db) == 0
 273     db["foo"] = b"bar"
 274     assert db["foo"] == b"bar"
 275     db.sync()
 276
 277     db = open("tmp", "c")
 278     assert len(db) == 1, len(db)
 279     assert db["foo"] == b"bar"
 280     db["bar"] = b"foo"
 281     assert db["bar"] == b"foo"
 282     db.sync()
 283
 284     db = open("tmp")
 285     assert len(db) == 2, len(db)
 286     assert db["foo"] == b"bar"
 287     assert db["bar"] == b"foo"
 288     try:
 289         db.sync()
 290     except OSError as e:
 291         assert str(e) == "Read-only database: tmp.dblite"
 292     else:
 293         raise RuntimeError("IOError expected.")
 294     db = open("tmp", "w")
 295     assert len(db) == 2, len(db)
 296     db["ping"] = b"pong"
 297     db.sync()
 298
 299     try:
 300         db[(1, 2)] = "tuple"
 301     except TypeError as e:
 302         assert str(e) == "key `(1, 2)' must be a string but is <class 'tuple'>", str(e)
 303     else:
 304         raise RuntimeError("TypeError exception expected")
 305
 306     try:
 307         db["list"] = [1, 2]
 308     except TypeError as e:
 309         assert str(e) == "value `[1, 2]' must be bytes but is <class 'list'>", str(e)
 310     else:
 311         raise RuntimeError("TypeError exception expected")
 312
 313     db = open("tmp")
 314     assert len(db) == 3, len(db)
 315
 316     db = open("tmp", "n")
 317     assert len(db) == 0, len(db)
 318     _Dblite._open("tmp.dblite", "w")
 319
 320     db = open("tmp")
 321     _Dblite._open("tmp.dblite", "w").write("x")
 322     try:
 323         db = open("tmp")
 324     except pickle.UnpicklingError:
 325         pass
 326     else:
 327         raise RuntimeError("pickle exception expected.")
 328
 329     global IGNORE_CORRUPT_DBFILES
 330     IGNORE_CORRUPT_DBFILES = True
 331     db = open("tmp")
 332     assert len(db) == 0, len(db)
 333     os.unlink("tmp.dblite")
 334     try:
 335         db = open("tmp", "w")
 336     except OSError as e:
 337         assert str(e) == "[Errno 2] No such file or directory: 'tmp.dblite'", str(e)
 338     else:
 339         raise RuntimeError("IOError expected.")
 340
 341     print("Completed _exercise()")
 342
 343
 344 if __name__ == "__main__":
 345     _exercise()
 346
 347 # Local Variables:
 348 # tab-width:4
 349 # indent-tabs-mode:nil
 350 # End:
 351 # vim: set expandtab tabstop=4 shiftwidth=4: