934b2bc0e9
Part of blueprint taskflow-redis-jobs Change-Id: I7c94e2201c5d933c8a1ec73fc0cf705962e5eef6
134 lines
4.3 KiB
Python
134 lines
4.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright (C) 2015 Yahoo! Inc. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"); you may
|
|
# not use this file except in compliance with the License. You may obtain
|
|
# a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
|
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
|
# License for the specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
import enum
|
|
|
|
import redis
|
|
from redis import exceptions as redis_exceptions
|
|
import six
|
|
|
|
|
|
def _raise_on_closed(meth):
|
|
|
|
@six.wraps(meth)
|
|
def wrapper(self, *args, **kwargs):
|
|
if self.closed:
|
|
raise redis_exceptions.ConnectionError("Connection has been"
|
|
" closed")
|
|
return meth(self, *args, **kwargs)
|
|
|
|
return wrapper
|
|
|
|
|
|
class RedisClient(redis.StrictRedis):
|
|
"""A redis client that can be closed (and raises on-usage after closed).
|
|
|
|
TODO(harlowja): if https://github.com/andymccurdy/redis-py/issues/613 ever
|
|
gets resolved or merged or other then we can likely remove this.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super(RedisClient, self).__init__(*args, **kwargs)
|
|
self.closed = False
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
self.connection_pool.disconnect()
|
|
|
|
execute_command = _raise_on_closed(redis.StrictRedis.execute_command)
|
|
transaction = _raise_on_closed(redis.StrictRedis.transaction)
|
|
pubsub = _raise_on_closed(redis.StrictRedis.pubsub)
|
|
|
|
|
|
class UnknownExpire(enum.IntEnum):
|
|
"""Non-expiry (not ttls) results return from :func:`.get_expiry`.
|
|
|
|
See: http://redis.io/commands/ttl or http://redis.io/commands/pttl
|
|
"""
|
|
|
|
DOES_NOT_EXPIRE = -1
|
|
"""
|
|
The command returns ``-1`` if the key exists but has no associated expire.
|
|
"""
|
|
|
|
#: The command returns ``-2`` if the key does not exist.
|
|
KEY_NOT_FOUND = -2
|
|
|
|
|
|
DOES_NOT_EXPIRE = UnknownExpire.DOES_NOT_EXPIRE
|
|
KEY_NOT_FOUND = UnknownExpire.KEY_NOT_FOUND
|
|
|
|
_UNKNOWN_EXPIRE_MAPPING = dict((e.value, e) for e in list(UnknownExpire))
|
|
|
|
|
|
def get_expiry(client, key, prior_version=None):
|
|
"""Gets an expiry for a key (using **best** determined ttl method)."""
|
|
is_new_enough, _prior_version = is_server_new_enough(
|
|
client, (2, 6), prior_version=prior_version)
|
|
if is_new_enough:
|
|
result = client.pttl(key)
|
|
try:
|
|
return _UNKNOWN_EXPIRE_MAPPING[result]
|
|
except KeyError:
|
|
return result / 1000.0
|
|
else:
|
|
result = client.ttl(key)
|
|
try:
|
|
return _UNKNOWN_EXPIRE_MAPPING[result]
|
|
except KeyError:
|
|
return float(result)
|
|
|
|
|
|
def apply_expiry(client, key, expiry, prior_version=None):
|
|
"""Applies an expiry to a key (using **best** determined expiry method)."""
|
|
is_new_enough, _prior_version = is_server_new_enough(
|
|
client, (2, 6), prior_version=prior_version)
|
|
if is_new_enough:
|
|
# Use milliseconds (as that is what pexpire uses/expects...)
|
|
ms_expiry = expiry * 1000.0
|
|
ms_expiry = max(0, int(ms_expiry))
|
|
result = client.pexpire(key, ms_expiry)
|
|
else:
|
|
# Only supports seconds (not subseconds...)
|
|
sec_expiry = int(expiry)
|
|
sec_expiry = max(0, sec_expiry)
|
|
result = client.expire(key, sec_expiry)
|
|
return bool(result)
|
|
|
|
|
|
def is_server_new_enough(client, min_version,
|
|
default=False, prior_version=None):
|
|
"""Checks if a client is attached to a new enough redis server."""
|
|
if not prior_version:
|
|
try:
|
|
server_info = client.info()
|
|
except redis_exceptions.ResponseError:
|
|
server_info = {}
|
|
version_text = server_info.get('redis_version', '')
|
|
else:
|
|
version_text = prior_version
|
|
version_pieces = []
|
|
for p in version_text.split("."):
|
|
try:
|
|
version_pieces.append(int(p))
|
|
except ValueError:
|
|
break
|
|
if not version_pieces:
|
|
return (default, version_text)
|
|
else:
|
|
version_pieces = tuple(version_pieces)
|
|
return (version_pieces >= min_version, version_text)
|