Completed
Push — master ( 0e0c0d...203901 )
by Bertrand
54s
created

cachalot.inner()   B

Complexity

Conditions 6

Size

Total Lines 5

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 6
dl 0
loc 5
rs 8
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
from django.db.transaction import Atomic, get_connection
14
15
from .api import invalidate
16
from .cache import cachalot_caches
17
from .settings import cachalot_settings
18
from .utils import (
19
    _get_query_cache_key, _invalidate_table,
20
    _get_table_cache_keys, _get_tables_from_sql, UncachableQuery)
21
22
23
WRITE_COMPILERS = (SQLInsertCompiler, SQLUpdateCompiler, SQLDeleteCompiler)
24
25
26
def _unset_raw_connection(original):
27
    def inner(compiler, *args, **kwargs):
28
        compiler.connection.raw = False
29
        out = original(compiler, *args, **kwargs)
30
        compiler.connection.raw = True
31
        return out
32
    return inner
33
34
35
TUPLE_OR_LIST = (tuple, list)
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 new_table_cache_keys:
46
        now = time()
47
        d = {}
48
        for k in new_table_cache_keys:
49
            d[k] = now
50
        cache.set_many(d, None)
51
    elif cache_key in data:
52
        timestamp, result = data.pop(cache_key)
53
        table_times = data.values()
54
        if table_times and timestamp > max(table_times):
55
            return result
56
57
    result = execute_query_func()
58
    if isinstance(result, Iterable) and result.__class__ not in TUPLE_OR_LIST:
59
        result = list(result)
60
61
    cache.set(cache_key, (time(), result), None)
62
63
    return result
64
65
66
def _patch_compiler(original):
67
    @wraps(original)
68
    @_unset_raw_connection
69
    def inner(compiler, *args, **kwargs):
70
        execute_query_func = lambda: original(compiler, *args, **kwargs)
71
        if not cachalot_settings.CACHALOT_ENABLED \
72
                or isinstance(compiler, WRITE_COMPILERS):
73
            return execute_query_func()
74
75
        try:
76
            cache_key = _get_query_cache_key(compiler)
77
            table_cache_keys = _get_table_cache_keys(compiler)
78
        except (EmptyResultSet, UncachableQuery):
79
            return execute_query_func()
80
81
        return _get_result_or_execute_query(
82
            execute_query_func,
83
            cachalot_caches.get_cache(db_alias=compiler.using),
84
            cache_key, table_cache_keys)
85
86
    return inner
87
88
89
def _patch_write_compiler(original):
90
    @wraps(original)
91
    @_unset_raw_connection
92
    def inner(write_compiler, *args, **kwargs):
93
        db_alias = write_compiler.using
94
        table = write_compiler.query.get_meta().db_table
95
        _invalidate_table(cachalot_caches.get_cache(db_alias=db_alias),
96
                          db_alias, table)
97
        return original(write_compiler, *args, **kwargs)
98
99
    return inner
100
101
102
def _patch_orm():
103
    SQLCompiler.execute_sql = _patch_compiler(SQLCompiler.execute_sql)
104
    for compiler in WRITE_COMPILERS:
105
        compiler.execute_sql = _patch_write_compiler(compiler.execute_sql)
106
107
108
def _patch_cursor():
109
    def _patch_cursor_execute(original):
110
        @wraps(original)
111
        def inner(cursor, sql, *args, **kwargs):
112
            out = original(cursor, sql, *args, **kwargs)
113
            if getattr(cursor.db, 'raw', True) \
114
                    and cachalot_settings.CACHALOT_INVALIDATE_RAW:
115
                sql = sql.lower()
116
                if 'update' in sql or 'insert' in sql or 'delete' in sql:
117
                    tables = _get_tables_from_sql(cursor.db, sql)
118
                    invalidate(*tables, db_alias=cursor.db.alias)
119
            return out
120
121
        return inner
122
123
    CursorWrapper.execute = _patch_cursor_execute(CursorWrapper.execute)
124
    CursorWrapper.executemany = _patch_cursor_execute(CursorWrapper.executemany)
125
126
127
def _patch_atomic():
128
    def patch_enter(original):
129
        @wraps(original)
130
        def inner(self):
131
            cachalot_caches.enter_atomic(self.using)
132
            original(self)
133
134
        return inner
135
136
    def patch_exit(original):
137
        @wraps(original)
138
        def inner(self, exc_type, exc_value, traceback):
139
            needs_rollback = get_connection(self.using).needs_rollback
140
            original(self, exc_type, exc_value, traceback)
141
            cachalot_caches.exit_atomic(
142
                self.using, exc_type is None and not needs_rollback)
143
144
        return inner
145
146
    Atomic.__enter__ = patch_enter(Atomic.__enter__)
147
    Atomic.__exit__ = patch_exit(Atomic.__exit__)
148
149
150
def _invalidate_on_migration(sender, **kwargs):
151
    invalidate(*sender.get_models(), db_alias=kwargs['using'])
152
153
154
def patch():
155
    post_migrate.connect(_invalidate_on_migration)
156
157
    _patch_cursor()
158
    _patch_atomic()
159
    _patch_orm()
160