Completed
Push — master ( eb64d1...f6cfeb )
by Bertrand
55s
created

_invalidate_table()   A

Complexity

Conditions 3

Size

Total Lines 11

Duplication

Lines 0
Ratio 0 %

Importance

Changes 2
Bugs 0 Features 0
Metric Value
cc 3
c 2
b 0
f 0
dl 0
loc 11
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 .transaction import AtomicCache
19
20
21
DJANGO_GTE_1_9 = django_version[:2] >= (1, 9)
22
23
24
class UncachableQuery(Exception):
25
    pass
26
27
28
TUPLE_OR_LIST = {tuple, list}
29
30
CACHABLE_PARAM_TYPES = {
31
    bool, int, float, Decimal, binary_type, text_type, type(None),
32
    datetime.date, datetime.time, datetime.datetime, datetime.timedelta, UUID,
33
}
34
35
UNCACHABLE_FUNCS = set()
36
if DJANGO_GTE_1_9:
37
    from django.db.models.functions import Now
38
    from django.contrib.postgres.functions import TransactionNow
39
    UNCACHABLE_FUNCS.update((Now, TransactionNow))
40
41
try:
42
    from psycopg2.extras import (
43
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet, Json)
44
except ImportError:
45
    pass
46
else:
47
    CACHABLE_PARAM_TYPES.update((
48
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet, Json))
49
50
51
def check_parameter_types(params):
52
    for p in params:
53
        cl = p.__class__
54
        if cl not in CACHABLE_PARAM_TYPES:
55
            if cl in TUPLE_OR_LIST:
56
                check_parameter_types(p)
57
            elif cl is dict:
58
                check_parameter_types(p.items())
59
            else:
60
                raise UncachableQuery
61
62
63
def get_query_cache_key(compiler):
64
    """
65
    Generates a cache key from a SQLCompiler.
66
67
    This cache key is specific to the SQL query and its context
68
    (which database is used).  The same query in the same context
69
    (= the same database) must generate the same cache key.
70
71
    :arg compiler: A SQLCompiler that will generate the SQL query
72
    :type compiler: django.db.models.sql.compiler.SQLCompiler
73
    :return: A cache key
74
    :rtype: int
75
    """
76
    sql, params = compiler.as_sql()
77
    check_parameter_types(params)
78
    cache_key = '%s:%s:%s' % (compiler.using, sql, params)
79
    return sha1(cache_key.encode('utf-8')).hexdigest()
80
81
82
def get_table_cache_key(db_alias, table):
83
    """
84
    Generates a cache key from a SQL table.
85
86
    :arg db_alias: Alias of the used database
87
    :type db_alias: str or unicode
88
    :arg table: Name of the SQL table
89
    :type table: str or unicode
90
    :return: A cache key
91
    :rtype: int
92
    """
93
    cache_key = '%s:%s' % (db_alias, table)
94
    return sha1(cache_key.encode('utf-8')).hexdigest()
95
96
97
def _get_query_cache_key(compiler):
98
    return import_string(cachalot_settings.CACHALOT_QUERY_KEYGEN)(compiler)
99
100
101
def _get_table_cache_key(db_alias, table):
102
    return import_string(cachalot_settings.CACHALOT_TABLE_KEYGEN)(db_alias, table)
103
104
105
def _get_tables_from_sql(connection, lowercased_sql):
106
    return [t for t in connection.introspection.django_table_names()
107
            if t in lowercased_sql]
108
109
110
def _find_subqueries(children):
111
    for child in children:
112
        if child.__class__ is SubqueryConstraint:
113
            if child.query_object.__class__ is Query:
114
                yield child.query_object
115
            else:
116
                yield child.query_object.query
117
        else:
118
            rhs = None
119
            if hasattr(child, 'rhs'):
120
                rhs = child.rhs
121
            rhs_class = rhs.__class__
122
            if rhs_class is Query:
123
                yield rhs
124
            elif hasattr(rhs, 'query'):
125
                yield rhs.query
126
            elif rhs_class in UNCACHABLE_FUNCS:
127
                raise UncachableQuery
128
        if hasattr(child, 'children'):
129
            for grand_child in _find_subqueries(child.children):
130
                yield grand_child
131
132
133
def is_cachable(table):
134
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
135
    if whitelist and table not in whitelist:
136
        return False
137
    return table not in cachalot_settings.CACHALOT_UNCACHABLE_TABLES
138
139
140
def are_all_cachable(tables):
141
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
142
    if whitelist and not tables.issubset(whitelist):
143
        return False
144
    return tables.isdisjoint(cachalot_settings.CACHALOT_UNCACHABLE_TABLES)
145
146
147
def filter_cachable(tables):
148
    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
149
    tables = tables.difference(cachalot_settings.CACHALOT_UNCACHABLE_TABLES)
150
    if whitelist:
151
        return tables.intersection(whitelist)
152
    return tables
153
154
155
def _get_tables(query, db_alias):
156
    if ('?' in query.order_by and not cachalot_settings.CACHALOT_CACHE_RANDOM) \
157
            or query.select_for_update:
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