From 4c34fc71dcfc20862723de8b62573ec2f8fbaf32 Mon Sep 17 00:00:00 2001
From: Greg Brockman <gdb@mit.edu>
Date: Tue, 6 Jul 2010 00:28:05 -0400
Subject: [PATCH] Cleaned up atomic file-writing code in invirt-build-conf

svn path=/trunk/packages/invirt-dev/; revision=3039
---
 invirt-build-conf |   39 +++++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/invirt-build-conf b/invirt-build-conf
index d43bc7c..2966427 100755
--- a/invirt-build-conf
+++ b/invirt-build-conf
@@ -6,8 +6,9 @@ This script generates the remctl ACL and configuration for each build
 pocket defined in the configuration. It also updates the .k5login for
 the git user that developers can push through.
 """
+from __future__ import with_statement
 
-
+import contextlib
 import os
 import tempfile
 
@@ -26,30 +27,31 @@ def userToPrinc(user):
 
     return princ.replace('.', '/') + '@' + realm
 
+def acl_path(pocket):
+    return '/etc/remctl/acl/build-%s' % pocket
+
+@contextlib.contextmanager
+def atomic_write(file):
+    tmp_fd, tmp_name = tempfile.mkstemp()
+    tmp = os.fdopen(tmp_fd, 'r+')
+    yield tmp
+    tmp.close()
+    os.rename(tmp_name, file)
 
 def main():
     all_devs = set()
-
-    # Python could really use a file-like object that gets written to
-    # a temporary path and moved to its final resting place on
-    # .close(). Oh well.
-    conf_fd, conf_name = tempfile.mkstemp()
-    conf = os.fdopen(conf_fd, 'r+')
     build_handler = '/usr/bin/invirt-submit-build'
 
     for pocket in config.build.pockets:
         acl = authz.expandAdmin(getattr(config.build.pockets, pocket).acl, None)
+        with atomic_write(acl_path(pocket)) as f:
+            princs = [userToPrinc(a) for a in acl]
+            print >>f, '\n'.join(princs)
+            all_devs.update(set(princs))
 
-        acl_fd, acl_name = tempfile.mkstemp()
-        acl_fd = os.fdopen(acl_fd, 'r+')
-        print >>acl_fd, '\n'.join(userToPrinc(a) for a in acl)
-
-        all_devs.update(set(userToPrinc(a) for a in acl))
-
-        acl_path = os.path.join('/etc/remctl/acl/build-%s' % pocket)
-
-        os.rename(acl_name, acl_path)
-        print >>conf, 'build %s %s %s' % (pocket, build_handler, acl_path)
+    with atomic_write('/etc/remctl/conf.d/build') as f:
+        for pocket in config.build.pockets:
+            print >>f, 'build %s %s %s' % (pocket, build_handler, acl_path(pocket))
 
     os.rename(conf_name, '/etc/remctl/conf.d/build')
 
@@ -57,7 +59,8 @@ def main():
     k5login = os.fdopen(k5login_fd, 'r+')
     print >>k5login, '\n'.join(all_devs)
 
-    os.rename(k5login_name, os.path.join(builder._REPO_DIR, '.k5login'))
+    with atomic_write(os.path.join(builder._REPO_DIR, '.k5login')) as f:
+        print >>f, '\n'.join(all_devs)
 
 
 if __name__ == '__main__':
-- 
1.7.9.5