Completed
Push — master ( 522b5a...246873 )
by Bertrand
01:10
created

filter_cachable()   A

Complexity

Conditions 2

Size

Total Lines 6

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 2
c 1
b 0
f 0
dl 0
loc 6
rs 9.4285
1
# coding: utf-8
2
3
from __future__ import unicode_literals
4
import datetime
5
from decimal import Decimal
6
from hashlib import sha1
7
from time import time
8
from uuid import UUID
9
10
from django import VERSION as django_version
11
from django.db import connections
12
from django.db.models.sql import Query
13
from django.db.models.sql.where import ExtraWhere, SubqueryConstraint
14
from django.utils.module_loading import import_string
15
from django.utils.six import text_type, binary_type
16
17
from .settings import cachalot_settings
18
from .signals import post_invalidation
19
from .transaction import AtomicCache
20
21
22
DJANGO_GTE_1_9 = django_version[:2] >= (1, 9)
23
24
25
class UncachableQuery(Exception):
26
    pass
27
28
29
TUPLE_OR_LIST = {tuple, list}
30
31
CACHABLE_PARAM_TYPES = {
32
    bool, int, float, Decimal, binary_type, text_type, type(None),
33
    datetime.date, datetime.time, datetime.datetime, datetime.timedelta, UUID,
34
}
35
36
UNCACHABLE_FUNCS = set()
37
if DJANGO_GTE_1_9:
38
    from django.db.models.functions import Now
39
    from django.contrib.postgres.functions import TransactionNow
40
    UNCACHABLE_FUNCS.update((Now, TransactionNow))
41
42
try:
43
    from psycopg2.extras import (
44
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet, Json)
45
except ImportError:
46
    pass
47
else:
48
    CACHABLE_PARAM_TYPES.update((
49
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet, Json))
50
51
52
def check_parameter_types(params):
53
    for p in params:
54
        cl = p.__class__
55
        if cl not in CACHABLE_PARAM_TYPES:
56
            if cl in TUPLE_OR_LIST:
57
                check_parameter_types(p)
58
            elif cl is dict:
59
                check_parameter_types(p.items())
60
            else:
61
                raise UncachableQuery
62
63
64
def get_query_cache_key(compiler):
65
    """
66
    Generates a cache key from a SQLCompiler.
67
68
    This cache key is specific to the SQL query and its context
69
    (which database is used).  The same query in the same context
70
    (= the same database) must generate the same cache key.
71
72
    :arg compiler: A SQLCompiler that will generate the SQL query
73
    :type compiler: django.db.models.sql.compiler.SQLCompiler
74
    :return: A cache key
75
    :rtype: int
76
    """
77
    sql, params = compiler.as_sql()
78
    check_parameter_types(params)
79
    cache_key = '%s:%s:%s' % (compiler.using, sql, params)
80
    return sha1(cache_key.encode('utf-8')).hexdigest()
81
82
83
def get_table_cache_key(db_alias, table):
84
    """
85
    Generates a cache key from a SQL table.
86
87
    :arg db_alias: Alias of the used database
88
    :type db_alias: str or unicode
89
    :arg table: Name of the SQL table
90
    :type table: str or unicode
91
    :return: A cache key
92
    :rtype: int
93
    """
94
    cache_key = '%s:%s' % (db_alias, table)
95
    return sha1(cache_key.encode('utf-8')).hexdigest()
96
97
98
def _get_query_cache_key(compiler):
99
    return import_string(cachalot_settings.CACHALOT_QUERY_KEYGEN)(compiler)
100
101
102
def _get_table_cache_key(db_alias, table):
103
    return import_string(cachalot_settings.CACHALOT_TABLE_KEYGEN)(db_alias, table)
104
105
106
def _get_tables_from_sql(connection, lowercased_sql):
107
    return [t for t in connection.introspection.django_table_names()
108
            if t in lowercased_sql]
109
110
111
def _find_subqueries(children):
112
    for child in children:
113
        if child.__class__ is SubqueryConstraint:
114
            if child.query_object.__class__ is Query:
115
                yield child.query_object
116
            else:
117
                yield child.query_object.query
118
        else:
119
            rhs = None
120
            if hasattr(child, 'rhs'):
121
                rhs = child.rhs
122
            rhs_class = rhs.__class__
123
            if rhs_class is Query:
124
                yield rhs
125
            elif hasattr(rhs, 'query'):
126
                yield rhs.query
127
            elif rhs_class in UNCACHABLE_FUNCS:
128
                raise UncachableQuery
129
        if hasattr(child, 'children'):
130
            for grand_child in _find_subqueries(child.children):
131
                yield grand_child
132
133
134
def is_cachable(table):
135
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
136
    if whitelist and table not in whitelist:
137
        return False
138
    return table not in cachalot_settings.CACHALOT_UNCACHABLE_TABLES
139
140
141
def are_all_cachable(tables):
142
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
143
    if whitelist and not tables.issubset(whitelist):
144
        return False
145
    return tables.isdisjoint(cachalot_settings.CACHALOT_UNCACHABLE_TABLES)
146
147
148
def filter_cachable(tables):
149
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
150
    tables = tables.difference(cachalot_settings.CACHALOT_UNCACHABLE_TABLES)
151
    if whitelist:
152
        return tables.intersection(whitelist)
153
    return tables
154
155
156
def _get_tables(query, db_alias):
157
    if '?' in query.order_by and not cachalot_settings.CACHALOT_CACHE_RANDOM:
158
        raise UncachableQuery
159
160
    tables = set(query.table_map)
161
    tables.add(query.get_meta().db_table)
162
    subquery_constraints = _find_subqueries(query.where.children)
163
    for subquery in subquery_constraints:
164
        tables.update(_get_tables(subquery, db_alias))
165
    if query.extra_select or hasattr(query, 'subquery') \
166
            or any(c.__class__ is ExtraWhere for c in query.where.children):
167
        sql = query.get_compiler(db_alias).as_sql()[0].lower()
168
        additional_tables = _get_tables_from_sql(connections[db_alias], sql)
169
        tables.update(additional_tables)
170
171
    if not are_all_cachable(tables):
172
        raise UncachableQuery
173
    return tables
174
175
176
def _get_table_cache_keys(compiler):
177
    db_alias = compiler.using
178
    return [_get_table_cache_key(db_alias, t)
179
            for t in _get_tables(compiler.query, db_alias)]
180
181
182
def _invalidate_tables(cache, db_alias, tables):
183
    tables = filter_cachable(set(tables))
184
    if not tables:
185
        return
186
    now = time()
187
    cache.set_many(
188
        {_get_table_cache_key(db_alias, t): now for t in tables},
189
        cachalot_settings.CACHALOT_TIMEOUT)
190
191
    if isinstance(cache, AtomicCache):
192
        cache.to_be_invalidated.update(tables)
193
194
195
def _invalidate_table(cache, db_alias, table):
196
    if not is_cachable(table):
197
        return
198
199
    cache.set(_get_table_cache_key(db_alias, table), time(),
200
              cachalot_settings.CACHALOT_TIMEOUT)
201
202
    if isinstance(cache, AtomicCache):
203
        cache.to_be_invalidated.add(table)
204
    else:
205
        post_invalidation.send(table, db_alias=db_alias)
206