Skip to content

Commit

Permalink
Add extra options to storage constructors
Browse files Browse the repository at this point in the history
- Add library, client_getter overrides in MemcachedStorage constructors
  to optionally use a custom memcache library.
  • Loading branch information
alisaifee committed Mar 20, 2015
1 parent d5b26f0 commit 0857552
Showing 1 changed file with 17 additions and 9 deletions.
26 changes: 17 additions & 9 deletions limits/storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ class Storage(object):
Base class to extend when implementing a storage backend.
"""

def __init__(self, uri=None):
def __init__(self, uri=None, **options):
self.lock = threading.RLock()

@abstractmethod
Expand Down Expand Up @@ -99,7 +99,7 @@ class MemoryStorage(Storage):
"""
STORAGE_SCHEME = "memory"

def __init__(self, uri=None):
def __init__(self, uri=None, **_):
self.storage = Counter()
self.expirations = {}
self.events = {}
Expand Down Expand Up @@ -227,7 +227,7 @@ class RedisStorage(Storage):
return {oldest, a}
"""

def __init__(self, uri):
def __init__(self, uri, **_):
"""
:param str redis_url: url of the form 'redis://host:port'
:raise ConfigurationError: when the redis library is not available
Expand Down Expand Up @@ -314,7 +314,7 @@ class MemcachedStorage(Storage):
MAX_CAS_RETRIES = 10
STORAGE_SCHEME = "memcached"

def __init__(self, uri):
def __init__(self, uri, **options):
"""
:param str host: memcached host
:param int port: memcached port
Expand All @@ -325,22 +325,30 @@ def __init__(self, uri):
for loc in parsed.netloc.split(","):
host, port = loc.split(":")
self.cluster.append((host, int(port)))
self.library = options.get('library', 'pymemcache.client')
self.client_getter = options.get('client_getter', self.get_client)

if not get_dependency("pymemcache"):
if not get_dependency(self.library):
raise ConfigurationError("memcached prerequisite not available."
" please install pymemcache") # pragma: no cover
" please install %s" % self.library) # pragma: no cover
self.local_storage = threading.local()
self.local_storage.storage = None

def get_client(self, hosts):
"""
returns a memcached client.
:param hosts:
:return:
"""
return get_dependency(self.library).Client(*hosts)

@property
def storage(self):
"""
lazily creates a memcached client instance using a thread local
"""
if not (hasattr(self.local_storage, "storage") and self.local_storage.storage):
self.local_storage.storage = get_dependency(
"pymemcache.client"
).Client(*self.cluster)
self.local_storage.storage = self.client_getter(self.cluster)
return self.local_storage.storage

def get(self, key):
Expand Down

0 comments on commit 0857552

Please sign in to comment.