http://scons.tigris.org/issues/show_bug.cgi?id=2329
[scons.git] / src / engine / SCons / dblite.py
index c9845a751d545fce2a788d73ff42dae464f19af7..2383f1df0452349c573dcb21fea1bdf42d96ea47 100644 (file)
@@ -5,57 +5,101 @@ import cPickle
 import time
 import shutil
 import os
-import types
 import __builtin__
 
-_open = __builtin__.open # avoid name clash
-
 keep_all_files = 00000
 ignore_corrupt_dbfiles = 0
 
-if hasattr(types, 'UnicodeType'):
+def corruption_warning(filename):
+    print "Warning: Discarding corrupt database:", filename
+
+try: unicode
+except NameError:
     def is_string(s):
-        t = type(s)
-        return t is types.StringType or t is types.UnicodeType
+        return isinstance(s, str)
 else:
     def is_string(s):
-        return type(s) is types.StringType
+        return type(s) in (str, unicode)
 
 try:
     unicode('a')
 except NameError:
     def unicode(s): return s
 
+dblite_suffix = '.dblite'
+tmp_suffix = '.tmp'
+
 class dblite:
 
+  # Squirrel away references to the functions in various modules
+  # that we'll use when our __del__() method calls our sync() method
+  # during shutdown.  We might get destroyed when Python is in the midst
+  # of tearing down the different modules we import in an essentially
+  # arbitrary order, and some of the various modules's global attributes
+  # may already be wiped out from under us.
+  #
+  # See the discussion at:
+  #   http://mail.python.org/pipermail/python-bugs-list/2003-March/016877.html
+
+  _open = __builtin__.open
+  _cPickle_dump = cPickle.dump
+  _os_chmod = os.chmod
+  try:
+      _os_chown = os.chown
+  except AttributeError:
+      _os_chown = None
+  _os_rename = os.rename
+  _os_unlink = os.unlink
+  _shutil_copyfile = shutil.copyfile
+  _time_time = time.time
+
   def __init__(self, file_base_name, flag, mode):
     assert flag in (None, "r", "w", "c", "n")
     if (flag is None): flag = "r"
-    if file_base_name[-7:] != '.dblite':
-        file_base_name = file_base_name + '.dblite'
-    self._file_name = file_base_name
+    base, ext = os.path.splitext(file_base_name)
+    if ext == dblite_suffix:
+      # There's already a suffix on the file name, don't add one.
+      self._file_name = file_base_name
+      self._tmp_name = base + tmp_suffix
+    else:
+      self._file_name = file_base_name + dblite_suffix
+      self._tmp_name = file_base_name + tmp_suffix
     self._flag = flag
     self._mode = mode
     self._dict = {}
     self._needs_sync = 00000
+    if self._os_chown is not None and (os.geteuid()==0 or os.getuid()==0):
+      # running as root; chown back to current owner/group when done
+      try:
+        statinfo = os.stat(self._file_name)
+        self._chown_to = statinfo.st_uid
+        self._chgrp_to = statinfo.st_gid
+      except OSError, e:
+        # db file doesn't exist yet.
+        # Check os.environ for SUDO_UID, use if set
+        self._chown_to = int(os.environ.get('SUDO_UID', -1))
+        self._chgrp_to = int(os.environ.get('SUDO_GID', -1))
+    else:
+      self._chown_to = -1        # don't chown
+      self._chgrp_to = -1        # don't chgrp
     if (self._flag == "n"):
-      _open(self._file_name, "wb", self._mode)
+      self._open(self._file_name, "wb", self._mode)
     else:
       try:
-        f = _open(self._file_name, "rb")
+        f = self._open(self._file_name, "rb")
       except IOError, e:
         if (self._flag != "c"):
           raise e
-        _open(self._file_name, "wb", self._mode)
+        self._open(self._file_name, "wb", self._mode)
       else:
         p = f.read()
         if (len(p) > 0):
           try:
             self._dict = cPickle.loads(p)
-          except:
+          except (cPickle.UnpicklingError, EOFError):
             if (ignore_corrupt_dbfiles == 0): raise
             if (ignore_corrupt_dbfiles == 1):
-              print "Warning: Discarding corrupt database:", self._file_name
+              corruption_warning(self._file_name)
 
   def __del__(self):
     if (self._needs_sync):
@@ -63,14 +107,29 @@ class dblite:
 
   def sync(self):
     self._check_writable()
-    f = _open(self._file_name, "wb", self._mode)
-    cPickle.dump(self._dict, f, 1)
+    f = self._open(self._tmp_name, "wb", self._mode)
+    self._cPickle_dump(self._dict, f, 1)
     f.close()
+    # Windows doesn't allow renaming if the file exists, so unlink
+    # it first, chmod'ing it to make sure we can do so.  On UNIX, we
+    # may not be able to chmod the file if it's owned by someone else
+    # (e.g. from a previous run as root).  We should still be able to
+    # unlink() the file if the directory's writable, though, so ignore
+    # any OSError exception  thrown by the chmod() call.
+    try: self._os_chmod(self._file_name, 0777)
+    except OSError: pass
+    self._os_unlink(self._file_name)
+    self._os_rename(self._tmp_name, self._file_name)
+    if self._os_chown is not None and self._chown_to > 0: # don't chown to root or -1
+      try:
+        self._os_chown(self._file_name, self._chown_to, self._chgrp_to)
+      except OSError:
+        pass
     self._needs_sync = 00000
     if (keep_all_files):
-      shutil.copyfile(
+      self._shutil_copyfile(
         self._file_name,
-        self._file_name + "_" + str(int(time.time())))
+        self._file_name + "_" + str(int(self._time_time())))
 
   def _check_writable(self):
     if (self._flag == "r"):
@@ -180,3 +239,9 @@ def _exercise():
 
 if (__name__ == "__main__"):
   _exercise()
+
+# Local Variables:
+# tab-width:4
+# indent-tabs-mode:nil
+# End:
+# vim: set expandtab tabstop=4 shiftwidth=4: