You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
rq/rq/connections.py

78 lines
2.1 KiB
Python

# -*- coding: utf-8 -*-
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from contextlib import contextmanager
from redis import StrictRedis
from .compat.connections import patch_connection
from .local import LocalStack, release_local
class NoRedisConnectionException(Exception):
pass
@contextmanager
def Connection(connection=None):
if connection is None:
connection = StrictRedis()
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(patch_connection(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 = StrictRedis()
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 patch_connection(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']