Completed
Push — master ( 78c434...bc49fd )
by Bertrand
01:05
created

_patch_compiler()   B

Complexity

Conditions 6

Size

Total Lines 22

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 6
dl 0
loc 22
rs 7.7857
c 1
b 0
f 0
1
# coding: utf-8
2
3
from __future__ import unicode_literals
4
from collections import Iterable
5
from functools import wraps
6
from time import time
7
8
from django.db.backends.utils import CursorWrapper
9
from django.db.models.query import EmptyResultSet
10
from django.db.models.signals import post_migrate
11
from django.db.models.sql.compiler import (
12
    SQLCompiler, SQLInsertCompiler, SQLUpdateCompiler, SQLDeleteCompiler,
13
)
14
from django.db.transaction import Atomic, get_connection
15
from django.utils.six import binary_type
16
17
from .api import invalidate
18
from .cache import cachalot_caches
19
from .settings import cachalot_settings, ITERABLES
20
from .utils import (
21
    _get_table_cache_keys, _get_tables_from_sql,
22
    UncachableQuery, is_cachable, filter_cachable,
23
)
24
25
26
WRITE_COMPILERS = (SQLInsertCompiler, SQLUpdateCompiler, SQLDeleteCompiler)
27
28
29
def _unset_raw_connection(original):
30
    def inner(compiler, *args, **kwargs):
31
        compiler.connection.raw = False
32
        out = original(compiler, *args, **kwargs)
33
        compiler.connection.raw = True
34
        return out
35
    return inner
36
37
38
def _get_result_or_execute_query(execute_query_func, cache,
39
                                 cache_key, table_cache_keys):
40
    data = cache.get_many(table_cache_keys + [cache_key])
41
42
    new_table_cache_keys = set(table_cache_keys)
43
    new_table_cache_keys.difference_update(data)
44
45
    if not new_table_cache_keys and cache_key in data:
46
        timestamp, result = data.pop(cache_key)
47
        if timestamp >= max(data.values()):
48
            return result
49
50
    result = execute_query_func()
51
    if result.__class__ not in ITERABLES and isinstance(result, Iterable):
52
        result = list(result)
53
54
    now = time()
55
    to_be_set = {k: now for k in new_table_cache_keys}
56
    to_be_set[cache_key] = (now, result)
57
    cache.set_many(to_be_set, cachalot_settings.CACHALOT_TIMEOUT)
58
59
    return result
60
61
62
def _patch_compiler(original):
63
    @wraps(original)
64
    @_unset_raw_connection
65
    def inner(compiler, *args, **kwargs):
66
        execute_query_func = lambda: original(compiler, *args, **kwargs)
67
        db_alias = compiler.using
68
        if db_alias not in cachalot_settings.CACHALOT_DATABASES \
69
                or isinstance(compiler, WRITE_COMPILERS):
70
            return execute_query_func()
71
72
        try:
73
            cache_key = cachalot_settings.CACHALOT_QUERY_KEYGEN(compiler)
74
            table_cache_keys = _get_table_cache_keys(compiler)
75
        except (EmptyResultSet, UncachableQuery):
76
            return execute_query_func()
77
78
        return _get_result_or_execute_query(
79
            execute_query_func,
80
            cachalot_caches.get_cache(db_alias=db_alias),
81
            cache_key, table_cache_keys)
82
83
    return inner
84
85
86
def _patch_write_compiler(original):
87
    @wraps(original)
88
    @_unset_raw_connection
89
    def inner(write_compiler, *args, **kwargs):
90
        db_alias = write_compiler.using
91
        table = write_compiler.query.get_meta().db_table
92
        if is_cachable(table):
93
            invalidate(table, db_alias=db_alias,
94
                       cache_alias=cachalot_settings.CACHALOT_CACHE)
95
        return original(write_compiler, *args, **kwargs)
96
97
    return inner
98
99
100
def _patch_orm():
101
    if cachalot_settings.CACHALOT_ENABLED:
102
        SQLCompiler.execute_sql = _patch_compiler(SQLCompiler.execute_sql)
103
    for compiler in WRITE_COMPILERS:
104
        compiler.execute_sql = _patch_write_compiler(compiler.execute_sql)
105
106
107
def _unpatch_orm():
108
    if hasattr(SQLCompiler.execute_sql, '__wrapped__'):
109
        SQLCompiler.execute_sql = SQLCompiler.execute_sql.__wrapped__
110
    for compiler in WRITE_COMPILERS:
111
        compiler.execute_sql = compiler.execute_sql.__wrapped__
112
113
114
def _patch_cursor():
115
    def _patch_cursor_execute(original):
116
        @wraps(original)
117
        def inner(cursor, sql, *args, **kwargs):
118
            out = original(cursor, sql, *args, **kwargs)
119
            connection = cursor.db
120
            if getattr(connection, 'raw', True):
121
                if isinstance(sql, binary_type):
122
                    sql = sql.decode('utf-8')
123
                sql = sql.lower()
124
                if 'update' in sql or 'insert' in sql or 'delete' in sql \
125
                        or 'alter' in sql or 'create' in sql or 'drop' in sql:
126
                    tables = filter_cachable(
127
                        _get_tables_from_sql(connection, sql))
128
                    if tables:
129
                        invalidate(*tables, db_alias=connection.alias,
130
                                   cache_alias=cachalot_settings.CACHALOT_CACHE)
131
            return out
132
133
        return inner
134
135
    if cachalot_settings.CACHALOT_INVALIDATE_RAW:
136
        CursorWrapper.execute = _patch_cursor_execute(CursorWrapper.execute)
137
        CursorWrapper.executemany = \
138
            _patch_cursor_execute(CursorWrapper.executemany)
139
140
141
def _unpatch_cursor():
142
    if hasattr(CursorWrapper.execute, '__wrapped__'):
143
        CursorWrapper.execute = CursorWrapper.execute.__wrapped__
144
        CursorWrapper.executemany = CursorWrapper.executemany.__wrapped__
145
146
147
def _patch_atomic():
148
    def patch_enter(original):
149
        @wraps(original)
150
        def inner(self):
151
            cachalot_caches.enter_atomic(self.using)
152
            original(self)
153
154
        return inner
155
156
    def patch_exit(original):
157
        @wraps(original)
158
        def inner(self, exc_type, exc_value, traceback):
159
            needs_rollback = get_connection(self.using).needs_rollback
160
            original(self, exc_type, exc_value, traceback)
161
            cachalot_caches.exit_atomic(
162
                self.using, exc_type is None and not needs_rollback)
163
164
        return inner
165
166
    Atomic.__enter__ = patch_enter(Atomic.__enter__)
167
    Atomic.__exit__ = patch_exit(Atomic.__exit__)
168
169
170
def _unpatch_atomic():
171
    Atomic.__enter__ = Atomic.__enter__.__wrapped__
172
    Atomic.__exit__ = Atomic.__exit__.__wrapped__
173
174
175
def _invalidate_on_migration(sender, **kwargs):
176
    invalidate(*sender.get_models(), db_alias=kwargs['using'],
177
               cache_alias=cachalot_settings.CACHALOT_CACHE)
178
179
180
def patch():
181
    post_migrate.connect(_invalidate_on_migration)
182
183
    _patch_cursor()
184
    _patch_atomic()
185
    _patch_orm()
186
187
188
def unpatch():
189
    post_migrate.disconnect(_invalidate_on_migration)
190
191
    _unpatch_cursor()
192
    _unpatch_atomic()
193
    _unpatch_orm()
194