router.py

import threading
from functools import wraps
from uuid import uuid4

from django.db import connections
from django.utils import six


THREAD_LOCAL = threading.local()


class DynamicDbRouter(object):
    """A router that decides what db to read from based on a variable
    local to the current thread.
    """

    def db_for_read(self, model, **hints):
        return getattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE', ['default'])[-1]

    def db_for_write(self, model, **hints):
        return getattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE', ['default'])[-1]

    def allow_relation(self, *args, **kwargs):
        return True

    def allow_syncdb(self, *args, **kwargs):
        return None

    def allow_migrate(self, *args, **kwargs):
        return None


class in_database(object):
    """A decorator and context manager to do queries on a given database.

    :type database: str or dict
    :param database: The database to run queries on. A string
        will route through the matching database in
        ``django.conf.settings.DATABASES``. A dictionary will set up a
        connection with the given configuration and route queries to it.

    :type read: bool, optional
    :param read: Controls whether database reads will route through
        the provided database. If ``False``, reads will route through
        the ``'default'`` database. Defaults to ``True``.

    :type write: bool, optional
    :param write: Controls whether database writes will route to
        the provided database. If ``False``, writes will route to
        the ``'default'`` database. Defaults to ``False``.

    When used as eithe a decorator or a context manager, `in_database`
    requires a single argument, which is the name of the database to
    route queries to, or a configuration dictionary for a database to
    route to.

    Usage as a context manager:

    .. code-block:: python

        from my_django_app.utils import tricky_query

        with in_database('Database_A'):
            results = tricky_query()

    Usage as a decorator:

    .. code-block:: python

        from my_django_app.models import Account

        @in_database('Database_B')
        def lowest_id_account():
            Account.objects.order_by('-id')[0]

    Used with a configuration dictionary:

    .. code-block:: python

        db_config = {'ENGINE': 'django.db.backends.sqlite3',
                     'NAME': 'path/to/mydatabase.db'}
        with in_database(db_config):
            # Run queries
    """
    def __init__(self, database, read=True, write=False):
        self.read = read
        self.write = write
        self.created_db_config = False
        if isinstance(database, six.string_types):
            self.database = database
        elif isinstance(database, dict):
            # Note: this invalidates the docs above. Update them
            # eventually.
            self.created_db_config = True
            self.unique_db_id = str(uuid4())
            connections.databases[self.unique_db_id] = database
            self.database = self.unique_db_id
        else:
            msg = ("database must be an identifier for an existing db, "
                   "or a complete configuration.")
            raise ValueError(msg)

    def __enter__(self):
        if not hasattr(THREAD_LOCAL, 'DB_FOR_READ_OVERRIDE'):
            THREAD_LOCAL.DB_FOR_READ_OVERRIDE = ['default']
        if not hasattr(THREAD_LOCAL, 'DB_FOR_WRITE_OVERRIDE'):
            THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE = ['default']
        read_db = (self.database if self.read
                   else THREAD_LOCAL.DB_FOR_READ_OVERRIDE[-1])
        write_db = (self.database if self.write
                    else THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE[-1])
        THREAD_LOCAL.DB_FOR_READ_OVERRIDE.append(read_db)
        THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE.append(write_db)
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        THREAD_LOCAL.DB_FOR_READ_OVERRIDE.pop()
        THREAD_LOCAL.DB_FOR_WRITE_OVERRIDE.pop()
        if self.created_db_config:
            connections[self.unique_db_id].close()
            del connections.databases[self.unique_db_id]

    def __call__(self, querying_func):
        @wraps(querying_func)
        def inner(*args, **kwargs):
            # Call the function in our context manager
            with self:
                return querying_func(*args, **kwargs)
        return inner