- added extra argument con_proxy to ConnectionListener interface checkout/checkin methods

- changed testing connection closer to work on _ConnectionFairy instances, resulting in
pool checkins, not actual closes
- disabled session two phase test for now, needs work
- added some two-phase support to TLEngine, not tested
- TLTransaction is now a wrapper

4 files changed, 61 insertions, 22 deletions

diff --git a/lib/sqlalchemy/engine/threadlocal.py b/lib/sqlalchemy/engine/threadlocal.py
index 4b251de13..164c50f51 100644
--- a/lib/sqlalchemy/engine/threadlocal.py
+++ b/lib/sqlalchemy/engine/threadlocal.py
@@ -30,6 +30,20 @@ class TLSession(object):
 
     def in_transaction(self):
         return self.__tcount > 0
+    
+    def prepare(self):
+        if self.__tcount == 1:
+            try:
+                self.__trans._trans.prepare()
+            finally:
+                self.reset()
+
+    def begin_twophase(self, xid=None):
+        if self.__tcount == 0:
+            self.__transaction = self.get_connection()
+            self.__trans = self.__transaction._begin_twophase(xid=xid)
+        self.__tcount += 1
+        return self.__trans
 
     def begin(self, **kwargs):
         if self.__tcount == 0:
@@ -41,14 +55,14 @@ class TLSession(object):
     def rollback(self):
         if self.__tcount > 0:
             try:
-                self.__trans._rollback_impl()
+                self.__trans._trans.rollback()
             finally:
                 self.reset()
 
     def commit(self):
         if self.__tcount == 1:
             try:
-                self.__trans._commit_impl()
+                self.__trans._trans.commit()
             finally:
                 self.reset()
         elif self.__tcount > 1:
@@ -69,15 +83,21 @@ class TLConnection(base.Connection):
         self.__opencount += 1
         return self
 
-    def _begin(self):
-        return TLTransaction(self)
-
+    def _begin(self, **kwargs):
+        return TLTransaction(super(TLConnection, self).begin(**kwargs), self.__session)
+    
+    def _begin_twophase(self, xid=None):
+        return TLTransaction(super(TLConnection, self).begin_twophase(xid=xid), self.__session)
+        
     def in_transaction(self):
         return self.session.in_transaction()
 
     def begin(self, **kwargs):
         return self.session.begin(**kwargs)
 
+    def begin_twophase(self, xid=None):
+        return self.session.begin_twophase(xid=xid)
+
     def close(self):
         if self.__opencount == 1:
             base.Connection.close(self)
@@ -87,18 +107,29 @@ class TLConnection(base.Connection):
         self.__opencount = 0
         base.Connection.close(self)
 
-class TLTransaction(base.RootTransaction):
-    def _commit_impl(self):
-        base.Transaction.commit(self)
+class TLTransaction(base.Transaction):
+    def __init__(self, trans, session):
+        self._trans = trans
+        self._session = session
 
-    def _rollback_impl(self):
-        base.Transaction.rollback(self)
+    connection = property(lambda s:s._trans.connection)
+    is_active = property(lambda s:s._trans.is_active)
+
+    def rollback(self):
+        self._session.rollback()
 
+    def prepare(self):
+        self._session.prepare()
+        
     def commit(self):
-        self.connection.session.commit()
+        self._session.commit()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self._trans.__exit__(type, value, traceback)
 
-    def rollback(self):
-        self.connection.session.rollback()
 
 class TLEngine(base.Engine):
     """An Engine that includes support for thread-local managed transactions.
diff --git a/lib/sqlalchemy/interfaces.py b/lib/sqlalchemy/interfaces.py
index 227e1b01a..05a8a4a34 100644
--- a/lib/sqlalchemy/interfaces.py
+++ b/lib/sqlalchemy/interfaces.py
@@ -50,17 +50,22 @@ class PoolListener(object):
           ``Connection`` wrapper).
 
         con_record
-          The ``_ConnectionRecord`` that currently owns the connection
+          The ``_ConnectionRecord`` that persistently manages the connection
+          
         """
 
-    def checkout(dbapi_con, con_record):
+    def checkout(dbapi_con, con_record, con_proxy):
         """Called when a connection is retrieved from the Pool.
 
         dbapi_con
           A raw DB-API connection
 
         con_record
-          The ``_ConnectionRecord`` that currently owns the connection
+          The ``_ConnectionRecord`` that persistently manages the connection
+
+        con_proxy
+          The ``_ConnectionFairy`` which manages the connection for the span of
+          the current checkout.
 
         If you raise an ``exceptions.DisconnectionError``, the current
         connection will be disposed and a fresh connection retrieved.
@@ -68,7 +73,7 @@ class PoolListener(object):
         using the new connection.
         """
 
-    def checkin(dbapi_con, con_record):
+    def checkin(dbapi_con, con_record, con_proxy):
         """Called when a connection returns to the pool.
 
         Note that the connection may be closed, and may be None if the
@@ -79,5 +84,10 @@ class PoolListener(object):
           A raw DB-API connection
 
         con_record
-          The _ConnectionRecord that currently owns the connection
+          The ``_ConnectionRecord`` that persistently manages the connection
+
+        con_proxy
+          The ``_ConnectionFairy`` which manages the connection for the span of
+          the current checkout.
+
         """
diff --git a/lib/sqlalchemy/orm/session.py b/lib/sqlalchemy/orm/session.py
index 109c468fc..6263a2e52 100644
--- a/lib/sqlalchemy/orm/session.py
+++ b/lib/sqlalchemy/orm/session.py
@@ -241,9 +241,7 @@ class SessionTransaction(object):
             return
         for t in util.Set(self.__connections.values()):
             if t[2]:
-                # fixme: wrong-
                 # closing the connection will also issue a rollback()
-                t[1].rollback()
                 t[0].close()
         self.session.transaction = None
 
diff --git a/lib/sqlalchemy/pool.py b/lib/sqlalchemy/pool.py
index 5dbee8993..b3fe2c09b 100644
--- a/lib/sqlalchemy/pool.py
+++ b/lib/sqlalchemy/pool.py
@@ -318,7 +318,7 @@ class _ConnectionFairy(object):
         while attempts > 0:
             try:
                 for l in self._pool._on_checkout:
-                    l.checkout(self.connection, self._connection_record)
+                    l.checkout(self.connection, self._connection_record, self)
                 return self
             except exceptions.DisconnectionError, e:
                 self._pool.log(
@@ -372,7 +372,7 @@ class _ConnectionFairy(object):
                 self._pool.log("Connection %s being returned to pool" % repr(self.connection))
             if self._pool._on_checkin:
                 for l in self._pool._on_checkin:
-                    l.checkin(self.connection, self._connection_record)
+                    l.checkin(self.connection, self._connection_record, self)
             self._pool.return_conn(self)
         self.connection = None
         self._connection_record = None