pushlog: use a context manager for sqlite connection draft
authorGregory Szorc <gps@mozilla.com>
Tue, 09 Sep 2014 13:21:44 -0700
changeset 1688 525d372b77e1a5fe8004ed04ba1ae84f267a1e09
parent 1687 ccb181db0e093818f5d845ac24e2220e59458670
child 1689 20cbe3e2e9256ecfed56e3d045083dcf668b4521
push idunknown
push userunknown
push dateunknown
pushlog: use a context manager for sqlite connection Previously we were performing manual connection management on the connection to the SQLite database. This method is prone to bugs due to uncaught exceptions, etc. We switch to using a context manager, which is much safer and easier to reason about.
hgext/pushlog/__init__.py
--- a/hgext/pushlog/__init__.py
+++ b/hgext/pushlog/__init__.py
@@ -1,13 +1,14 @@
 # This software may be used and distributed according to the terms of the
 # GNU General Public License version 2 or any later version.
 
 '''Record pushes to Mercurial repositories.'''
 
+import contextlib
 import os
 import sqlite3
 import stat
 import time
 import weakref
 
 testedwith = '3.0 3.1 3.2'
 
@@ -25,67 +26,60 @@ class pushlog(object):
     '''An interface to pushlog data.'''
 
     def __init__(self, repo):
         '''Create an instance bound to a sqlite database in a path.'''
         # Use a weak ref to avoid cycle. This object's lifetime should be no
         # greater than the repo's.
         self.repo = weakref.proxy(repo)
 
-    @property
+    @contextlib.contextmanager
     def conn(self):
-        if not hasattr(self, '_conn'):
-            path = self.repo.vfs.join('pushlog2.db')
-            create = False
-            if not os.path.exists(path):
+        path = self.repo.vfs.join('pushlog2.db')
+        create = False
+        if not os.path.exists(path):
+            create = True
+
+        conn = sqlite3.connect(path)
+        if not create:
+            res = conn.execute(
+                "SELECT COUNT(*) FROM SQLITE_MASTER WHERE name='pushlog'")
+            if res.fetchone()[0] != 1:
                 create = True
 
-            conn = sqlite3.connect(path)
-            if not create:
-                res = conn.execute(
-                    "SELECT COUNT(*) FROM SQLITE_MASTER WHERE name='pushlog'")
-                if res.fetchone()[0] != 1:
-                    create = True
+        if create:
+            for sql in SCHEMA:
+                conn.execute(sql)
+            conn.commit()
+            st = os.stat(path)
+            os.chmod(path, st.st_mode | stat.S_IWGRP)
 
-            if create:
-                for sql in SCHEMA:
-                    conn.execute(sql)
-                conn.commit()
-                st = os.stat(path)
-                os.chmod(path, st.st_mode | stat.S_IWGRP)
-
-            self._conn = conn
-
-        return self._conn
+        try:
+            yield conn
+        finally:
+            conn.close()
 
     def recordpush(self, nodes, user, when):
         '''Record a push into the pushlog.
 
         A push consists of a list of nodes, a username, and a time of the
         push.
         '''
-        c = self.conn
-        try:
+        with self.conn() as c:
             res = c.execute('INSERT INTO pushlog (user, date) VALUES (?, ?)', (user, when))
             pushid = res.lastrowid
             for e in nodes:
                 ctx = self.repo[e]
                 rev = ctx.rev()
                 node = ctx.hex()
 
                 c.execute('INSERT INTO changesets (pushid, rev, node) '
                         'VALUES (?, ?, ?)', (pushid, rev, node))
 
             c.commit()
-        finally:
-            # We mimic the behavior of the old pushlog for now and destroy the
-            # connection after it is used. This should eventually be changed to use
-            # more sane resource management, such as a context manager.
-            c.close()
-            self._conn = None
 
 def pretxnchangegrouphook(ui, repo, node=None, source=None, **kwargs):
     # This hook is executed whenever changesets are introduced. We ignore
     # new changesets unless they come from a push. ``source`` can be
     # ``push`` for ssh or ``serve`` for HTTP pushes.
     #
     # This is arguably the wrong thing to do: designing a system to record
     # all changes to the store is the right thing to do. However, things are