summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy/engine/threadlocal.py
blob: 610eedeaaa0fb4cab64aa3fc22d0d460e89a23cd (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
from sqlalchemy import schema, exceptions, util, sql, types
import StringIO, sys, re
import base, default

"""provides a thread-local transactional wrapper around the basic ComposedSQLEngine.  multiple calls to engine.connect()
will return the same connection for the same thread. also provides begin/commit methods on the engine itself
which correspond to a thread-local transaction."""

class TLSession(object):
    def __init__(self, engine):
        self.engine = engine
        self.__tcount = 0
    def get_connection(self, close_with_result=False):
        try:
            return self.__transaction
        except AttributeError:
            return base.Connection(self.engine, close_with_result=close_with_result)
    def begin(self):
        if self.__tcount == 0:
            self.__transaction = self.get_connection()
            self.__trans = self.__transaction.begin()
        self.__tcount += 1
    def rollback(self):
        if self.__tcount > 0:
            try:
                self.__trans.rollback()
            finally:
                del self.__transaction
                del self.__trans
                self.__tcount = 0
    def commit(self):
        if self.__tcount == 1:
            try:
                self._trans.commit()
            finally:
                del self.__transaction
                del self._trans
                self.__tcount = 0
        elif self.__tcount > 1:
            self.__tcount -= 1
    def is_begun(self):
        return self.__tcount > 0

    
class TLEngine(base.ComposedSQLEngine):
    """a ComposedSQLEngine that includes support for thread-local managed transactions.  This engine
    is better suited to be used with threadlocal Pool object."""
    def __init__(self, *args, **kwargs):
        """the TLEngine relies upon the ConnectionProvider having "threadlocal" behavior,
        so that once a connection is checked out for the current thread, you get that same connection
        repeatedly."""
        base.ComposedSQLEngine.__init__(self, *args, **kwargs)
        self.context = util.ThreadLocal()
    def raw_connection(self):
        """returns a DBAPI connection."""
        return self.connection_provider.get_connection()
    def connect(self, **kwargs):
        """returns a Connection that is not thread-locally scoped.  this is the equilvalent to calling
        "connect()" on a ComposedSQLEngine."""
        return base.Connection(self, self.connection_provider.unique_connection())

    def _session(self):
        if not hasattr(self.context, 'session'):
            self.context.session = TLSession(self)
        return self.context.session
    session = property(_session, doc="returns the current thread's TLSession")

    def contextual_connect(self, **kwargs):
        """returns a TLConnection which is thread-locally scoped."""
        return self.session.get_connection(**kwargs)
        
    def begin(self):
        return self.session.begin()
    def commit(self):
        self.session.commit()
    def rollback(self):
        self.session.rollback()

class TLocalConnectionProvider(default.PoolConnectionProvider):
    def unique_connection(self):
        return self._pool.unique_connection()