summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy/orm/scoping.py
blob: 28eb63819ec4ba5a52cb754c869f844383116415 (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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
# scoping.py
# Copyright (C) the SQLAlchemy authors and contributors
#
# This module is part of SQLAlchemy and is released under
# the MIT License: http://www.opensource.org/licenses/mit-license.php

import sqlalchemy.exceptions as sa_exc
from sqlalchemy.util import ScopedRegistry, to_list, get_cls_kwargs, deprecated
from sqlalchemy.orm import (
    EXT_CONTINUE, MapperExtension, class_mapper, object_session
    )
from sqlalchemy.orm import exc as orm_exc
from sqlalchemy.orm.session import Session


__all__ = ['ScopedSession']


class ScopedSession(object):
    """Provides thread-local management of Sessions.

    Usage::

      Session = scoped_session(sessionmaker(autoflush=True))

      ... use session normally.

    """

    def __init__(self, session_factory, scopefunc=None):
        self.session_factory = session_factory
        self.registry = ScopedRegistry(session_factory, scopefunc)
        self.extension = _ScopedExt(self)

    def __call__(self, **kwargs):
        if kwargs:
            scope = kwargs.pop('scope', False)
            if scope is not None:
                if self.registry.has():
                    raise sa_exc.InvalidRequestError("Scoped session is already present; no new arguments may be specified.")
                else:
                    sess = self.session_factory(**kwargs)
                    self.registry.set(sess)
                    return sess
            else:
                return self.session_factory(**kwargs)
        else:
            return self.registry()

    def remove(self):
        if self.registry.has():
            self.registry().close()
        self.registry.clear()

    @deprecated("Session.mapper is deprecated.  "
        "Please see http://www.sqlalchemy.org/trac/wiki/UsageRecipes/SessionAwareMapper "
        "for information on how to replicate its behavior.")
    def mapper(self, *args, **kwargs):
        """return a mapper() function which associates this ScopedSession with the Mapper.
        
        DEPRECATED.
        
        """

        from sqlalchemy.orm import mapper

        extension_args = dict((arg, kwargs.pop(arg))
                              for arg in get_cls_kwargs(_ScopedExt)
                              if arg in kwargs)

        kwargs['extension'] = extension = to_list(kwargs.get('extension', []))
        if extension_args:
            extension.append(self.extension.configure(**extension_args))
        else:
            extension.append(self.extension)
        return mapper(*args, **kwargs)

    def configure(self, **kwargs):
        """reconfigure the sessionmaker used by this ScopedSession."""

        self.session_factory.configure(**kwargs)

    def query_property(self, query_cls=None):
        """return a class property which produces a `Query` object against the
        class when called.

        e.g.::
            Session = scoped_session(sessionmaker())

            class MyClass(object):
                query = Session.query_property()

            # after mappers are defined
            result = MyClass.query.filter(MyClass.name=='foo').all()

        Produces instances of the session's configured query class by
        default.  To override and use a custom implementation, provide
        a ``query_cls`` callable.  The callable will be invoked with
        the class's mapper as a positional argument and a session
        keyword argument.

        There is no limit to the number of query properties placed on
        a class.

        """
        class query(object):
            def __get__(s, instance, owner):
                try:
                    mapper = class_mapper(owner)
                    if mapper:
                        if query_cls:
                            # custom query class
                            return query_cls(mapper, session=self.registry())
                        else:
                            # session's configured query class
                            return self.registry().query(mapper)
                except orm_exc.UnmappedClassError:
                    return None
        return query()

def instrument(name):
    def do(self, *args, **kwargs):
        return getattr(self.registry(), name)(*args, **kwargs)
    return do
for meth in Session.public_methods:
    setattr(ScopedSession, meth, instrument(meth))

def makeprop(name):
    def set(self, attr):
        setattr(self.registry(), name, attr)
    def get(self):
        return getattr(self.registry(), name)
    return property(get, set)
for prop in ('bind', 'dirty', 'deleted', 'new', 'identity_map', 'is_active'):
    setattr(ScopedSession, prop, makeprop(prop))

def clslevel(name):
    def do(cls, *args, **kwargs):
        return getattr(Session, name)(*args, **kwargs)
    return classmethod(do)
for prop in ('close_all', 'object_session', 'identity_key'):
    setattr(ScopedSession, prop, clslevel(prop))

class _ScopedExt(MapperExtension):
    def __init__(self, context, validate=False, save_on_init=True):
        self.context = context
        self.validate = validate
        self.save_on_init = save_on_init
        self.set_kwargs_on_init = True

    def validating(self):
        return _ScopedExt(self.context, validate=True)

    def configure(self, **kwargs):
        return _ScopedExt(self.context, **kwargs)

    def instrument_class(self, mapper, class_):
        class query(object):
            def __getattr__(s, key):
                return getattr(self.context.registry().query(class_), key)
            def __call__(s):
                return self.context.registry().query(class_)
            def __get__(self, instance, cls):
                return self

        if not 'query' in class_.__dict__:
            class_.query = query()

        if self.set_kwargs_on_init and class_.__init__ is object.__init__:
            class_.__init__ = self._default__init__(mapper)

    def _default__init__(ext, mapper):
        def __init__(self, **kwargs):
            for key, value in kwargs.iteritems():
                if ext.validate:
                    if not mapper.get_property(key, resolve_synonyms=False,
                                               raiseerr=False):
                        raise sa_exc.ArgumentError(
                            "Invalid __init__ argument: '%s'" % key)
                setattr(self, key, value)
        return __init__

    def init_instance(self, mapper, class_, oldinit, instance, args, kwargs):
        if self.save_on_init:
            session = kwargs.pop('_sa_session', None)
            if session is None:
                session = self.context.registry()
            session._save_without_cascade(instance)
        return EXT_CONTINUE

    def init_failed(self, mapper, class_, oldinit, instance, args, kwargs):
        sess = object_session(instance)
        if sess:
            sess.expunge(instance)
        return EXT_CONTINUE

    def dispose_class(self, mapper, class_):
        if hasattr(class_, '__init__') and hasattr(class_.__init__, '_oldinit'):
            if class_.__init__._oldinit is not None:
                class_.__init__ = class_.__init__._oldinit
            else:
                delattr(class_, '__init__')
        if hasattr(class_, 'query'):
            delattr(class_, 'query')