django-cachalot/cachalot/utils.py

# coding: utf-8

from __future__ import unicode_literals
import datetime
from decimal import Decimal
from hashlib import sha1
from time import time
from uuid import UUID

from django.db import connections
from django.db.models.sql import Query
from django.db.models.sql.where import ExtraWhere, SubqueryConstraint
from django.utils.module_loading import import_string
from django.utils.six import text_type, binary_type

from .settings import cachalot_settings
from .signals import post_invalidation
from .transaction import AtomicCache


class UncachableQuery(Exception):
    pass


CACHABLE_PARAM_TYPES = {
    bool, int, float, Decimal, binary_type, text_type, type(None),
    datetime.date, datetime.time, datetime.datetime, datetime.timedelta, UUID,
}

try:
    from psycopg2.extras import (
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet)
except ImportError:
    pass
else:
    CACHABLE_PARAM_TYPES.update((
        NumericRange, DateRange, DateTimeRange, DateTimeTZRange, Inet))


def check_parameter_types(params):
    for p in params:
        cl = p.__class__
        if cl not in CACHABLE_PARAM_TYPES:
            if cl is list or cl is tuple:
                check_parameter_types(p)
            elif cl is dict:
                check_parameter_types(p.items())
            else:
                raise UncachableQuery


def get_query_cache_key(compiler):
    """
    Generates a cache key from a SQLCompiler.

    This cache key is specific to the SQL query and its context
    (which database is used).  The same query in the same context
    (= the same database) must generate the same cache key.

    :arg compiler: A SQLCompiler that will generate the SQL query
    :type compiler: django.db.models.sql.compiler.SQLCompiler
    :return: A cache key
    :rtype: str
    """
    sql, params = compiler.as_sql()
    check_parameter_types(params)
    cache_key = '%s:%s:%s' % (compiler.using, sql, params)
    return sha1(cache_key.encode('utf-8')).hexdigest()


def get_table_cache_key(db_alias, table):
    """
    Generates a cache key from a SQL table.

    :arg db_alias: Alias of the used database
    :type db_alias: str or unicode
    :arg table: Name of the SQL table
    :type table: str or unicode
    :return: A cache key
    :rtype: str
    """
    cache_key = '%s:%s' % (db_alias, table)
    return sha1(cache_key.encode('utf-8')).hexdigest()


def _get_query_cache_key(compiler):
    return import_string(cachalot_settings.CACHALOT_QUERY_KEYGEN)(compiler)


def _get_table_cache_key(db_alias, table):
    return import_string(cachalot_settings.CACHALOT_TABLE_KEYGEN)(db_alias, table)


def _get_tables_from_sql(connection, lowercased_sql):
    return [t for t in connection.introspection.django_table_names()
            if t in lowercased_sql]


def _find_subqueries(children):
    for child in children:
        if child.__class__ is SubqueryConstraint:
            if child.query_object.__class__ is Query:
                yield child.query_object
            else:
                yield child.query_object.query
        else:
            rhs = None
            if hasattr(child, 'rhs'):
                rhs = child.rhs
            elif child.__class__ is tuple:
                rhs = child[-1]
            if rhs.__class__ is Query:
                yield rhs
            elif hasattr(rhs, 'query'):
                yield rhs.query
        if hasattr(child, 'children'):
            for grand_child in _find_subqueries(child.children):
                yield grand_child


def _get_tables(query, db_alias):
    if '?' in query.order_by and not cachalot_settings.CACHALOT_CACHE_RANDOM:
        raise UncachableQuery

    tables = set(query.table_map)
    tables.add(query.get_meta().db_table)
    subquery_constraints = _find_subqueries(query.where.children
                                            + query.having.children)
    for subquery in subquery_constraints:
        tables.update(_get_tables(subquery, db_alias))
    if query.extra_select or hasattr(query, 'subquery') \
            or any(c.__class__ is ExtraWhere for c in query.where.children):
        sql = query.get_compiler(db_alias).as_sql()[0].lower()
        additional_tables = _get_tables_from_sql(connections[db_alias], sql)
        tables.update(additional_tables)

    whitelist = cachalot_settings.CACHALOT_ONLY_CACHABLE_TABLES
    blacklist = cachalot_settings.CACHALOT_UNCACHABLE_TABLES
    if (whitelist and not tables.issubset(whitelist)) \
            or not tables.isdisjoint(blacklist):
        raise UncachableQuery
    return tables


def _get_table_cache_keys(compiler):
    db_alias = compiler.using
    tables = _get_tables(compiler.query, db_alias)
    return [_get_table_cache_key(db_alias, t) for t in tables]


def _invalidate_tables(cache, db_alias, tables):
    now = time()
    d = {}
    for table in tables:
        d[_get_table_cache_key(db_alias, table)] = now
    cache.set_many(d, None)

    if isinstance(cache, AtomicCache):
        cache.to_be_invalidated.update(tables)


def _invalidate_table(cache, db_alias, table):
    cache.set(_get_table_cache_key(db_alias, table), time(), None)

    if isinstance(cache, AtomicCache):
        cache.to_be_invalidated.add(table)
    else:
        post_invalidation.send(table, db_alias=db_alias)