From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 8 Jan 2011 20:33:38 +0000 (-0500)
Subject: - Threadlocal engine returns itself upon begin(),
X-Git-Tag: rel_0_6_6~3
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=f51fa433e93a6aa8df798ec83f81f7f6535e79f9;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

- Threadlocal engine returns itself upon begin(),
begin_nested(); engine then implements contextmanager
methods to allow the "with" statement. [ticket:2004]
---

diff --git a/CHANGES b/CHANGES
index 86a1e98bad..f53aec09f8 100644
--- a/CHANGES
+++ b/CHANGES
@@ -122,6 +122,10 @@ CHANGES
     prepare() won't raise if no transaction is in progress;
     this was a regression introduced in 0.6.  [ticket:1998]
 
+  - Threadlocal engine returns itself upon begin(),
+    begin_nested(); engine then implements contextmanager
+    methods to allow the "with" statement. [ticket:2004]
+
 - postgresql
   - Single element tuple expressions inside an IN clause
     parenthesize correctly, also from [ticket:1984]
diff --git a/lib/sqlalchemy/engine/threadlocal.py b/lib/sqlalchemy/engine/threadlocal.py
index 5694ccd4b5..8b8732b24f 100644
--- a/lib/sqlalchemy/engine/threadlocal.py
+++ b/lib/sqlalchemy/engine/threadlocal.py
@@ -33,7 +33,6 @@ class TLConnection(base.Connection):
         self.__opencount = 0
         base.Connection.close(self)
 
-
 class TLEngine(base.Engine):
     """An Engine that includes support for thread-local managed transactions."""
 
@@ -66,16 +65,28 @@ class TLEngine(base.Engine):
         if not hasattr(self._connections, 'trans'):
             self._connections.trans = []
         self._connections.trans.append(self.contextual_connect().begin_twophase(xid=xid))
+        return self
 
     def begin_nested(self):
         if not hasattr(self._connections, 'trans'):
             self._connections.trans = []
         self._connections.trans.append(self.contextual_connect().begin_nested())
+        return self
 
     def begin(self):
         if not hasattr(self._connections, 'trans'):
             self._connections.trans = []
         self._connections.trans.append(self.contextual_connect().begin())
+        return self
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        if type is None:
+            self.commit()
+        else:
+            self.rollback()
 
     def prepare(self):
         if not hasattr(self._connections, 'trans') or \
diff --git a/test/engine/test_transaction.py b/test/engine/test_transaction.py
index b2ca0ebfc2..0aaf036ede 100644
--- a/test/engine/test_transaction.py
+++ b/test/engine/test_transaction.py
@@ -638,6 +638,27 @@ class TLTransactionTest(TestBase):
         finally:
             external_connection.close()
 
+    def test_with_interface(self):
+        trans = tlengine.begin()
+        tlengine.execute(users.insert(), user_id=1, user_name='user1')
+        tlengine.execute(users.insert(), user_id=2, user_name='user2')
+        trans.commit()
+
+        trans = tlengine.begin()
+        tlengine.execute(users.insert(), user_id=3, user_name='user3')
+        trans.__exit__(Exception, "fake", None)
+        trans = tlengine.begin()
+        tlengine.execute(users.insert(), user_id=4, user_name='user4')
+        trans.__exit__(None, None, None)
+        eq_(
+            tlengine.execute(users.select().order_by(users.c.user_id)).fetchall(),
+            [
+                (1, 'user1'),
+                (2, 'user2'),
+                (4, 'user4'),
+            ]
+        )
+
     def test_commits(self):
         connection = tlengine.connect()
         assert connection.execute('select count(*) from query_users'
diff --git a/test/orm/test_transaction.py b/test/orm/test_transaction.py
index 9977afd8eb..e2f31d0d11 100644
--- a/test/orm/test_transaction.py
+++ b/test/orm/test_transaction.py
@@ -23,8 +23,6 @@ class TransactionTest(FixtureTest):
             })
         mapper(Address, addresses)
 
-
-
 class FixtureDataTest(TransactionTest):
     run_inserts = 'each'