summaryrefslogtreecommitdiff
path: root/lib/sqlalchemy/engine/util.py
blob: 4cc7df79074a02b92d8d7712f2c1a76c0c545214 (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
# engine/util.py
# Copyright (C) 2005-2021 the SQLAlchemy authors and contributors
# <see AUTHORS file>
#
# This module is part of SQLAlchemy and is released under
# the MIT License: https://www.opensource.org/licenses/mit-license.php

from .. import exc
from .. import util

try:
    from sqlalchemy.cyextension.util import _distill_params_20  # noqa
    from sqlalchemy.cyextension.util import _distill_raw_params  # noqa
except ImportError:
    from ._py_util import _distill_params_20  # noqa
    from ._py_util import _distill_raw_params  # noqa


def connection_memoize(key):
    """Decorator, memoize a function in a connection.info stash.

    Only applicable to functions which take no arguments other than a
    connection.  The memo will be stored in ``connection.info[key]``.
    """

    @util.decorator
    def decorated(fn, self, connection):
        connection = connection.connect()
        try:
            return connection.info[key]
        except KeyError:
            connection.info[key] = val = fn(self, connection)
            return val

    return decorated


class TransactionalContext:
    """Apply Python context manager behavior to transaction objects.

    Performs validation to ensure the subject of the transaction is not
    used if the transaction were ended prematurely.

    """

    _trans_subject = None

    def _transaction_is_active(self):
        raise NotImplementedError()

    def _transaction_is_closed(self):
        raise NotImplementedError()

    def _rollback_can_be_called(self):
        """indicates the object is in a state that is known to be acceptable
        for rollback() to be called.

        This does not necessarily mean rollback() will succeed or not raise
        an error, just that there is currently no state detected that indicates
        rollback() would fail or emit warnings.

        It also does not mean that there's a transaction in progress, as
        it is usually safe to call rollback() even if no transaction is
        present.

        .. versionadded:: 1.4.28

        """
        raise NotImplementedError()

    def _get_subject(self):
        raise NotImplementedError()

    @classmethod
    def _trans_ctx_check(cls, subject):
        trans_context = subject._trans_context_manager
        if trans_context:
            if not trans_context._transaction_is_active():
                raise exc.InvalidRequestError(
                    "Can't operate on closed transaction inside context "
                    "manager.  Please complete the context manager "
                    "before emitting further commands."
                )

    def __enter__(self):
        subject = self._get_subject()

        # none for outer transaction, may be non-None for nested
        # savepoint, legacy nesting cases
        trans_context = subject._trans_context_manager
        self._outer_trans_ctx = trans_context

        self._trans_subject = subject
        subject._trans_context_manager = self
        return self

    def __exit__(self, type_, value, traceback):
        subject = self._trans_subject

        # simplistically we could assume that
        # "subject._trans_context_manager is self".  However, any calling
        # code that is manipulating __exit__ directly would break this
        # assumption.  alembic context manager
        # is an example of partial use that just calls __exit__ and
        # not __enter__ at the moment.  it's safe to assume this is being done
        # in the wild also
        out_of_band_exit = (
            subject is None or subject._trans_context_manager is not self
        )

        if type_ is None and self._transaction_is_active():
            try:
                self.commit()
            except:
                with util.safe_reraise():
                    if self._rollback_can_be_called():
                        self.rollback()
            finally:
                if not out_of_band_exit:
                    subject._trans_context_manager = self._outer_trans_ctx
                self._trans_subject = self._outer_trans_ctx = None
        else:
            try:
                if not self._transaction_is_active():
                    if not self._transaction_is_closed():
                        self.close()
                else:
                    if self._rollback_can_be_called():
                        self.rollback()
            finally:
                if not out_of_band_exit:
                    subject._trans_context_manager = self._outer_trans_ctx
                self._trans_subject = self._outer_trans_ctx = None