- 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 86a1e98bade52ae9d6501e31ddae6f26c0d01b54..f53aec09f87509e2f95bf47b13459f54f4ed328c 100644 (file)
--- 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 5694ccd4b5d11922054847ecde3059ab23406120..8b8732b24f96208df06e5974f2bdc5d5a538673b 100644 (file)
--- 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 b2ca0ebfc2d5eeaecd42e681f5edd8ebd50d985d..0aaf036ede38959916b13222b29819625a3031da 100644 (file)
--- 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 9977afd8ebbac07d2ad86912d62888a9747ccd41..e2f31d0d11b54d4e2725f0d9148df75a83e6c51f 100644 (file)
--- 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'