|
- # -*- coding: utf-8 -*-
- from __future__ import (absolute_import, division, print_function,
- unicode_literals)
-
- from contextlib import contextmanager
-
- from redis import Redis
-
- from .local import LocalStack, release_local
-
-
- class NoRedisConnectionException(Exception):
- pass
-
-
- @contextmanager
- def Connection(connection=None): # noqa
- if connection is None:
- connection = Redis()
- push_connection(connection)
- try:
- yield
- finally:
- popped = pop_connection()
- assert popped == connection, \
- 'Unexpected Redis connection was popped off the stack. ' \
- 'Check your Redis connection setup.'
-
-
- def push_connection(redis):
- """Pushes the given connection on the stack."""
- _connection_stack.push(redis)
-
-
- def pop_connection():
- """Pops the topmost connection from the stack."""
- return _connection_stack.pop()
-
-
- def use_connection(redis=None):
- """Clears the stack and uses the given connection. Protects against mixed
- use of use_connection() and stacked connection contexts.
- """
- assert len(_connection_stack) <= 1, \
- 'You should not mix Connection contexts with use_connection()'
- release_local(_connection_stack)
-
- if redis is None:
- redis = Redis()
- push_connection(redis)
-
-
- def get_current_connection():
- """Returns the current Redis connection (i.e. the topmost on the
- connection stack).
- """
- return _connection_stack.top
-
-
- def resolve_connection(connection=None):
- """Convenience function to resolve the given or the current connection.
- Raises an exception if it cannot resolve a connection now.
- """
- if connection is not None:
- return connection
-
- connection = get_current_connection()
- if connection is None:
- raise NoRedisConnectionException('Could not resolve a Redis connection')
- return connection
-
-
- _connection_stack = LocalStack()
-
- __all__ = ['Connection', 'get_current_connection', 'push_connection',
- 'pop_connection', 'use_connection']
|