mirror of
https://github.com/GSA/notifications-admin.git
synced 2026-08-14 23:08:07 -04:00
Localize notification_utils to the admin
This changeset pulls in all of the notification_utils code directly into the admin and removes it as an external dependency. We are doing this to cut down on operational maintenance of the project and will begin removing parts of it no longer needed for the admin. Signed-off-by: Carlo Costino <carlo.costino@gsa.gov>
This commit is contained in:
Carlo Costino
13
notifications_utils/clients/redis/__init__.py
Normal file
13
notifications_utils/clients/redis/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
from datetime import datetime
|
||||
|
||||
from .request_cache import RequestCache # noqa: F401 (unused import)
|
||||
|
||||
|
||||
def total_limit_cache_key(service_id):
|
||||
return "{}-{}-{}".format(
|
||||
str(service_id), datetime.utcnow().strftime("%Y-%m-%d"), "total-count"
|
||||
)
|
||||
|
||||
|
||||
def rate_limit_cache_key(service_id, api_key_type):
|
||||
return "{}-{}".format(str(service_id), api_key_type)
|
||||
184
notifications_utils/clients/redis/redis_client.py
Normal file
184
notifications_utils/clients/redis/redis_client.py
Normal file
@@ -0,0 +1,184 @@
|
||||
import numbers
|
||||
import uuid
|
||||
from time import time
|
||||
|
||||
from flask import current_app
|
||||
from flask_redis import FlaskRedis
|
||||
|
||||
|
||||
def prepare_value(val):
|
||||
"""
|
||||
Only bytes, strings and numbers (ints, longs and floats) are acceptable
|
||||
for keys and values. Previously redis-py attempted to cast other types
|
||||
to str() and store the result. This caused must confusion and frustration
|
||||
when passing boolean values (cast to 'True' and 'False') or None values
|
||||
(cast to 'None'). It is now the user's responsibility to cast all
|
||||
key names and values to bytes, strings or numbers before passing the
|
||||
value to redis-py.
|
||||
"""
|
||||
# things redis-py natively supports
|
||||
if isinstance(
|
||||
val,
|
||||
(
|
||||
bytes,
|
||||
str,
|
||||
numbers.Number,
|
||||
),
|
||||
):
|
||||
return val
|
||||
# things we know we can safely cast to string
|
||||
elif isinstance(val, (uuid.UUID,)):
|
||||
return str(val)
|
||||
else:
|
||||
raise ValueError("cannot cast {} to a string".format(type(val)))
|
||||
|
||||
|
||||
class RedisClient:
|
||||
redis_store = FlaskRedis()
|
||||
active = False
|
||||
scripts = {}
|
||||
|
||||
def init_app(self, app):
|
||||
self.active = app.config.get("REDIS_ENABLED")
|
||||
if self.active:
|
||||
self.redis_store.init_app(app)
|
||||
|
||||
self.register_scripts()
|
||||
|
||||
def register_scripts(self):
|
||||
# delete keys matching a pattern supplied as a parameter. Does so in batches of 5000 to prevent unpack from
|
||||
# exceeding lua's stack limit, and also to prevent errors if no keys match the pattern.
|
||||
# Inspired by https://gist.github.com/ddre54/0a4751676272e0da8186
|
||||
self.scripts["delete-keys-by-pattern"] = self.redis_store.register_script(
|
||||
"""
|
||||
local keys = redis.call('keys', ARGV[1])
|
||||
local deleted = 0
|
||||
for i=1, #keys, 5000 do
|
||||
deleted = deleted + redis.call('del', unpack(keys, i, math.min(i + 4999, #keys)))
|
||||
end
|
||||
return deleted
|
||||
"""
|
||||
)
|
||||
|
||||
def delete_by_pattern(self, pattern, raise_exception=False):
|
||||
r"""
|
||||
Deletes all keys matching a given pattern, and returns how many keys were deleted.
|
||||
Pattern is defined as in the KEYS command: https://redis.io/commands/keys
|
||||
|
||||
* h?llo matches hello, hallo and hxllo
|
||||
* h*llo matches hllo and heeeello
|
||||
* h[ae]llo matches hello and hallo, but not hillo
|
||||
* h[^e]llo matches hallo, hbllo, ... but not hello
|
||||
* h[a-b]llo matches hallo and hbllo
|
||||
|
||||
Use \ to escape special characters if you want to match them verbatim
|
||||
"""
|
||||
if self.active:
|
||||
try:
|
||||
return self.scripts["delete-keys-by-pattern"](args=[pattern])
|
||||
except Exception as e:
|
||||
self.__handle_exception(
|
||||
e, raise_exception, "delete-by-pattern", pattern
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
def exceeded_rate_limit(self, cache_key, limit, interval, raise_exception=False):
|
||||
"""
|
||||
Rate limiting.
|
||||
- Uses Redis sorted sets
|
||||
- Also uses redis "multi" which is abstracted into pipeline() by FlaskRedis/PyRedis
|
||||
- Sends all commands to redis as a group to be executed atomically
|
||||
|
||||
Method:
|
||||
(1) Add event, scored by timestamp (zadd). The score determines order in set.
|
||||
(2) Use zremrangebyscore to delete all set members with a score between
|
||||
- Earliest entry (lowest score == earliest timestamp) - represented as '-inf'
|
||||
and
|
||||
- Current timestamp minus the interval
|
||||
- Leaves only relevant entries in the set (those between now and now - interval)
|
||||
(3) Count the set
|
||||
(4) If count > limit fail request
|
||||
(5) Ensure we expire the set key to preserve space
|
||||
|
||||
Notes:
|
||||
- Failed requests count. If over the limit and keep making requests you'll stay over the limit.
|
||||
- The actual value in the set is just the timestamp, the same as the score. We don't store any requets details.
|
||||
- return value of pipe.execute() is an array containing the outcome of each call.
|
||||
- result[2] == outcome of pipe.zcard()
|
||||
- If redis is inactive, or we get an exception, allow the request
|
||||
|
||||
:param cache_key:
|
||||
:param limit: Number of requests permitted within interval
|
||||
:param interval: Interval we measure requests in
|
||||
:param raise_exception: Should throw exception
|
||||
:return:
|
||||
"""
|
||||
cache_key = prepare_value(cache_key)
|
||||
if self.active:
|
||||
try:
|
||||
pipe = self.redis_store.pipeline()
|
||||
when = time()
|
||||
pipe.zadd(cache_key, {when: when})
|
||||
pipe.zremrangebyscore(cache_key, "-inf", when - interval)
|
||||
pipe.zcard(cache_key)
|
||||
pipe.expire(cache_key, interval)
|
||||
result = pipe.execute()
|
||||
return result[2] > limit
|
||||
except Exception as e:
|
||||
self.__handle_exception(
|
||||
e, raise_exception, "rate-limit-pipeline", cache_key
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return False
|
||||
|
||||
def raw_set(self, key, value, ex=None, px=None, nx=False, xx=False):
|
||||
self.redis_store.set(key, value, ex, px, nx, xx)
|
||||
|
||||
def set(
|
||||
self, key, value, ex=None, px=None, nx=False, xx=False, raise_exception=False
|
||||
):
|
||||
key = prepare_value(key)
|
||||
value = prepare_value(value)
|
||||
if self.active:
|
||||
try:
|
||||
self.redis_store.set(key, value, ex, px, nx, xx)
|
||||
except Exception as e:
|
||||
self.__handle_exception(e, raise_exception, "set", key)
|
||||
|
||||
def incr(self, key, raise_exception=False):
|
||||
key = prepare_value(key)
|
||||
if self.active:
|
||||
try:
|
||||
return self.redis_store.incr(key)
|
||||
except Exception as e:
|
||||
self.__handle_exception(e, raise_exception, "incr", key)
|
||||
|
||||
def raw_get(self, key):
|
||||
return self.redis_store.get(key)
|
||||
|
||||
def get(self, key, raise_exception=False):
|
||||
key = prepare_value(key)
|
||||
if self.active:
|
||||
try:
|
||||
return self.redis_store.get(key)
|
||||
except Exception as e:
|
||||
self.__handle_exception(e, raise_exception, "get", key)
|
||||
|
||||
return None
|
||||
|
||||
def delete(self, *keys, raise_exception=False):
|
||||
keys = [prepare_value(k) for k in keys]
|
||||
if self.active:
|
||||
try:
|
||||
self.redis_store.delete(*keys)
|
||||
except Exception as e:
|
||||
self.__handle_exception(e, raise_exception, "delete", ", ".join(keys))
|
||||
|
||||
def __handle_exception(self, e, raise_exception, operation, key_name):
|
||||
current_app.logger.exception(
|
||||
"Redis error performing {} on {}".format(operation, key_name)
|
||||
)
|
||||
if raise_exception:
|
||||
raise e
|
||||
95
notifications_utils/clients/redis/request_cache.py
Normal file
95
notifications_utils/clients/redis/request_cache.py
Normal file
@@ -0,0 +1,95 @@
|
||||
import json
|
||||
from contextlib import suppress
|
||||
from datetime import timedelta
|
||||
from functools import wraps
|
||||
from inspect import signature
|
||||
|
||||
|
||||
class RequestCache:
|
||||
DEFAULT_TTL = int(timedelta(days=7).total_seconds())
|
||||
|
||||
def __init__(self, redis_client):
|
||||
self.redis_client = redis_client
|
||||
|
||||
@staticmethod
|
||||
def _get_argument(argument_name, client_method, args, kwargs):
|
||||
with suppress(KeyError):
|
||||
return kwargs[argument_name]
|
||||
|
||||
with suppress(ValueError, IndexError):
|
||||
argument_index = list(signature(client_method).parameters).index(
|
||||
argument_name
|
||||
)
|
||||
return args[argument_index]
|
||||
|
||||
with suppress(KeyError):
|
||||
return signature(client_method).parameters[argument_name].default
|
||||
|
||||
raise TypeError(
|
||||
"{}() takes no argument called '{}'".format(
|
||||
client_method.__name__, argument_name
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _make_key(key_format, client_method, args, kwargs):
|
||||
return key_format.format(
|
||||
**{
|
||||
argument_name: RequestCache._get_argument(
|
||||
argument_name, client_method, args, kwargs
|
||||
)
|
||||
for argument_name in list(signature(client_method).parameters)
|
||||
}
|
||||
)
|
||||
|
||||
def set(self, key_format, *, ttl_in_seconds=DEFAULT_TTL):
|
||||
def _set(client_method):
|
||||
@wraps(client_method)
|
||||
def new_client_method(*args, **kwargs):
|
||||
redis_key = RequestCache._make_key(
|
||||
key_format, client_method, args, kwargs
|
||||
)
|
||||
cached = self.redis_client.get(redis_key)
|
||||
if cached:
|
||||
return json.loads(cached.decode("utf-8"))
|
||||
api_response = client_method(*args, **kwargs)
|
||||
self.redis_client.set(
|
||||
redis_key,
|
||||
json.dumps(api_response),
|
||||
ex=int(ttl_in_seconds),
|
||||
)
|
||||
return api_response
|
||||
|
||||
return new_client_method
|
||||
|
||||
return _set
|
||||
|
||||
def delete(self, key_format):
|
||||
def _delete(client_method):
|
||||
@wraps(client_method)
|
||||
def new_client_method(*args, **kwargs):
|
||||
try:
|
||||
api_response = client_method(*args, **kwargs)
|
||||
finally:
|
||||
redis_key = self._make_key(key_format, client_method, args, kwargs)
|
||||
self.redis_client.delete(redis_key)
|
||||
return api_response
|
||||
|
||||
return new_client_method
|
||||
|
||||
return _delete
|
||||
|
||||
def delete_by_pattern(self, key_format):
|
||||
def _delete(client_method):
|
||||
@wraps(client_method)
|
||||
def new_client_method(*args, **kwargs):
|
||||
try:
|
||||
api_response = client_method(*args, **kwargs)
|
||||
finally:
|
||||
redis_key = self._make_key(key_format, client_method, args, kwargs)
|
||||
self.redis_client.delete_by_pattern(redis_key)
|
||||
return api_response
|
||||
|
||||
return new_client_method
|
||||
|
||||
return _delete
|
||||
Reference in New Issue
Block a user