async def test_depth(
monkeypatch,
input_file: str,
expected_file: str,
get_redis_client: Redis,
) -> None:
await test_init(monkeypatch, get_redis_client)
input_data = FileUtils.load_params_from_json(os.path.join(PATH, input_file))
expected_data = FileUtils.load_params_from_json(os.path.join(PATH, expected_file))
# monkeypatch.chdir(os.path.abspath(os.path.dirname(__file__)))
if len(input_data) == len(expected_data):
for i in range(len(input_data)):
# logger.info(f"input line = {i}, and input_data[i]= {input_data[i]}")
depth = get_results(input_data, i)
# logger.info(f"depth : {depth}")
for key, value in depth.items():
# logger.info(f"key, value - {key} and {value}")
# below: get value for result key in expected results
expected_result = expected_data[i][key]
logger.info(f"expected_result : {expected_result}")
assert value == expected_result
await test_init(monkeypatch, get_redis_client)
get_redis_client.close()
await get_redis_client.wait_closed()
async def _print_stats(loop):
worker_queue = config["WORKER_QUEUE"]
pool = await create_connection_pool(loop)
redis = Redis(pool)
worker_queue_length = await redis.llen(worker_queue)
requests = await get_counter(pool, "requests", 60)
connects = await get_counter(pool, "connects", 60)
errors = await get_counter(pool, "errors", 60)
intercepted = await get_counter(pool, "intercepted", 60)
not_authorized = await get_counter(pool, "not_authorized", 60)
authorized = await get_counter(pool, "authorized", 60)
ratelimited = await get_counter(pool, "ratelimited", 60)
responses = await get_counter(pool, "responses", 60)
session_count = ProxySession.query \
.count()
request_count = Request.query \
.count()
active_user_count = User.query \
.filter_by(active=True) \
.count()
not_active_user_count = User.query \
.filter_by(active=False) \
.count()
#intercept_count = Intercept.query \
# .count()
print("proxy_sessions.value", session_count)
print("requests.value", request_count)
print("active_users.value", active_user_count)
print("not_active_users.value", not_active_user_count)
#print("user_intercepts.value", intercept_count)
def calc_rate(l):
rate = sum(map(int, l[1:6])) / 5
return '{0:.2f}'.format(rate)
print("requests_5_minute.value " + calc_rate(requests))
print("connects_5_minute.value " + calc_rate(connects))
print("errors_5_minutes.value " + calc_rate(errors))
print("intercepted_5_minute.value " + calc_rate(intercepted))
print("not_authorized_5_minute.value " + calc_rate(not_authorized))
print("authorized_5_minute.value " + calc_rate(authorized))
print("ratelimited_5_minute.value " + calc_rate(ratelimited))
print("responses_5_minute.value " + calc_rate(responses))
print("worker_queue.value " + str(worker_queue_length))
redis.close()
pool.close()
await pool.wait_closed()
async def create_redis_room_store(redis: Redis) -> RedisRoomStore:
(append_to_room, lreplace, delete_room) = await asyncio.gather(
redis.script_load(_APPEND_TO_ROOM),
redis.script_load(_LREPLACE),
redis.script_load(_DELETE_ROOM),
)
return RedisRoomStore(redis, append_to_room, lreplace, delete_room)
async def close_redis_connection(connection: aioredis.Redis) -> None:
logger.info("Closing connection to Redis")
# gracefully closing underlying connection
connection.close()
await connection.wait_closed()
logger.info("Connection to Redis closed")
async def get_ride_id(redis_cli: Redis, user_id: str) -> Optional[str]:
async for key in redis_cli.iscan(match=_get_ride_key("*", user_id)):
return str(key.decode().split(":")[1])
async for _ in redis_cli.iscan(match=_get_ride_key(user_id, "*")): # noqa
return user_id
return None
def add_namespace(self, pipe: Redis, keys: NamespaceKeys) -> None:
"""Add the namespace to be cleaned up.
Args:
pipe: Piped redis commands.
keys: The namespace key group.
"""
cleanup_val = keys.root.named['namespace']
pipe.rpush(self.keys.namespaces, cleanup_val)
def add_root(self, pipe: Redis, root: RedisKey) -> None:
"""Add the content to be cleaned up.
Args:
pipe: Piped redis commands.
root: The redis key prefix.
"""
cleanup_val = root.wildcard
pipe.rpush(self.keys.roots, cleanup_val)
async def delete_rule(_id: str, db: AioRedis):
"""
Deletes a rate limiter rule.
@param id: (string) the ID of the rate limiter rule to delete
@param db: (object) db connection
"""
await asyncio.gather(
db.delete(_id),
RateLimiter._clear_indexes(_id, db),
db.srem(rules_set, _id)
)
async def delete(_id: str, db: AioRedis):
"""
deletes a endpoint cache rule.
@param id: (string) id of endpoint cache to delete
@param db: (object) db connection
"""
await asyncio.gather(
db.delete(_id),
EndpointCacher._clear_indexes(_id, db),
db.srem(endpoint_cache_set, _id),
)
def add_mailbox(self, pipe: Redis, keys: MailboxKeys) -> None:
"""Add the mailbox to be cleaned up.
Args:
pipe: Piped redis commands.
keys: The mailbox key group.
"""
namespace = keys.root.named['namespace']
mailbox_id = keys.root.named['mailbox_id']
cleanup_val = b'%b\x00%b' % (namespace, mailbox_id)
pipe.rpush(self.keys.mailboxes, cleanup_val)
def add_content(self, pipe: Redis, keys: ContentKeys) -> None:
"""Add the content to be cleaned up.
Args:
pipe: Piped redis commands.
keys: The content key group.
"""
namespace = keys.root.named['namespace']
email_id = keys.root.named['email_id']
cleanup_val = b'%b\x00%b' % (namespace, email_id)
pipe.rpush(self.keys.contents, cleanup_val)
async def delete_entry(_id: str, db: AioRedis):
"""
Deletes a rate limiter entry.
@param host: (string) the hostname of the rate limiter entry to delete
@param db: (object) db connection
"""
await asyncio.gather(
db.delete(_id),
RateLimiter._clear_indexes(_id, db),
db.srem(entry_set, _id)
)
async def increment_counter(redis_pool, name, count=1, now=None):
try:
now = now or time.time()
redis = Redis(redis_pool)
pipe = redis.pipeline()
for prec in PRECISION:
pnow = int(now / prec) * prec
hash = '%s:%s' % (prec, name)
pipe.zadd('known:', 0, hash)
pipe.hincrby('count:' + hash, pnow, count)
await pipe.execute()
except Exception as e:
logger.info("Error saving stats", e)
async def create_rule(ctx: object, db: AioRedis):
"""
Creates a rate limiter rule.
@param ctx: (object) data to be inserted
@param db: (object) db connection
"""
ctx['_id'] = str(bson.ObjectId())
await asyncio.gather(
RateLimiter._set_indexes(ctx, db),
db.hmset_dict(ctx['_id'], ctx),
db.sadd(rules_set, ctx['_id']),
)
def add_message(self, pipe: Redis, keys: MessageKeys) -> None:
"""Add the message to be cleaned up.
Args:
pipe: Piped redis commands.
keys: The message key group.
"""
namespace = keys.root.named['namespace']
mailbox_id = keys.root.named['mailbox_id']
msg_uid = keys.root.named['uid']
cleanup_val = b'%b\x00%b\x00%b' \
% (namespace, mailbox_id, msg_uid)
pipe.rpush(self.keys.messages, cleanup_val)
async def create_entry(ctx, db: AioRedis):
"""
Creates a rate limiter entry.
@param ctx: (object) data to be inserted
@param db: (object) db connection
"""
ctx['_id'] = str(bson.ObjectId())
await asyncio.gather(
RateLimiter._set_indexes(ctx, db),
db.hmset_dict(ctx['_id'], ctx),
db.sadd(entry_set, ctx['_id']),
db.expire(ctx['_id'], int(ctx['timeout']))
)
async def _cleanup(self, conn: aioredis.Redis, sid: str) -> None:
logger.debug("cleaning up scheduler data: sid=%s", sid)
async for worker_id, job_data in conn.ihscan(self._stash_key, match=f'{sid}#*'):
job_id, run_at = job_data.decode().split(',')
logger.debug("releasing dead job: sid=%s, id=%s, run_at=%s", sid, job_id, run_at)
with await self._conn_pool as conn:
tx = conn.multi_exec()
tx.zadd(self._timeline_key, float(run_at), job_id)
tx.hdel(self._stash_key, worker_id)
await tx.execute()
await conn.hdel(self._heartbeats_key, field=sid)
async def write_combinations(redis: aioredis.Redis):
keys = await redis.mget(*ALL_KEYS)
# do nothing, 60sec is enough to pass all tests
if set(keys) != {None}:
return
pipe = redis.pipeline()
for i, comb in enumerate(ALL_COMBINATIONS):
key_ = mk_key(comb)
if i % 3 == 0:
bundle = {
random.choice(COMB_PARTS): random.choice(COMB_PARTS)
for k in range(len(COMB_PARTS))
}
pipe.hmset_dict(key_, **bundle)
elif i % 5 == 0:
bundle = {
random.choice(COMB_PARTS): random.randint(1, 10)
for k in range(len(COMB_PARTS))
}
pipe.zadd(key_,
*chain(*[reversed(pair) for pair in bundle.items()]))
elif i % 7 == 0:
pipe.sadd(key_, *COMB_PARTS)
elif i % 11 == 0:
pipe.lpush(key_, *COMB_PARTS)
else:
pipe.set(key_, key_)
pipe.expire(key_, 60)
return await pipe.execute()
async def connect(self):
"""Connect to the database.
This method will connect to a Redis database. By default it will
connect to Redis on localhost on port 6379
"""
try:
self.client = Redis(
host=self.host,
port=int(self.port),
db=self.database,
password=self.password,
)
await self.client.ping() # to actually initiate a connection
_LOGGER.info(
_("Connected to Redis database %s from %s on port %s."),
self.database,
self.host,
self.port,
)
except OSError:
_LOGGER.warning(
_("Unable to connect to Redis database on address: %s port: %s."
),
self.host,
self.port,
)
async def init(self):
if self._pool:
return
if type(self.redis_host) is str and self.sentinel:
raise ConfigurationError(
"str provided for `redis_host` but `sentinel` is true, list of sentinels expected"
)
if self.sentinel:
addr = self.redis_host
async def pool_factory(*args: Any,
**kwargs: Any) -> ConnectionsPool:
client = await aioredis.sentinel.create_sentinel_pool(
*args, ssl=self.ssl, **kwargs)
return client.master_for(self.sentinel_master)
else:
pool_factory = functools.partial(aioredis.create_pool,
ssl=self.ssl)
addr = self.redis_host, self.redis_port
self._pool = await pool_factory(addr,
db=self.redis_db,
password=self.redis_password,
encoding="utf8")
self._redis = Redis(self._pool)
await Tortoise.init(config=self.tortoise_config)
await Tortoise.generate_schemas()
async def kline_interval(market: str,
interval: int,
count: int,
redis_client: Redis = Depends(get_redis_database)):
# f"^kline-{market}.*-{interval}.*-latest$"
pattern = f"kline-{market}-{interval}-*"
kline_keys = []
async for key in redis_client.iscan(match=pattern):
kline_keys.append(key)
kline_keys.sort(reverse=True)
# results = []
# for key in kline_keys:
# value = await CrudRedisGeneral.get(redis_client, key)
# results.append({key:value})
results = collections.OrderedDict()
for i in range(len(kline_keys)):
if i >= count:
break
key = kline_keys[i]
value = await CrudRedisGeneral.get(redis_client, key)
results[key] = value
return results
async def dump_bytecode(self, bucket):
key = self.make_key(bucket)
pool = await self.backend
context = await pool
with context as conn:
conn = Redis(conn)
await conn.set(key, bucket.bytecode_to_string())
async def wait_for_result(self, task_id):
key = (self.task_keyprefix + task_id).encode()
# TODO share a single connection for all waiting results
# connections effectively get a task to process messages
# coming in. The built-in Connection then wants to put that
# on a queue. Do we then need another task processing those
# queues? The case of two concurrent waits on the same
# celery task is the tricky one. We can't just have the call
# to wait_for_result own that queue. Maybe the first call
# could own that queue and setup a future for any other calls
# to wait on.
async with self.redis_pool.get() as conn:
conn = Redis(conn)
chan, = await conn.subscribe(key)
try:
while True:
encodedmeta = await Redis(self.redis_pool).get(key)
if not encodedmeta:
meta = {'status': states.PENDING, 'result': None}
else:
meta = self.celery.backend.decode_result(encodedmeta)
if meta['status'] in states.READY_STATES:
if meta['status'] in states.PROPAGATE_STATES:
raise meta['result']
return meta['result']
await chan.get()
finally:
await conn.unsubscribe(key)
async def __call__(
self, config: Config = Depends(config_dependency)) -> Redis:
"""Creates the Redis pool if necessary and returns it."""
if not self.redis:
password = config.redis_password
self.redis = Redis.from_url(config.redis_url, password=password)
assert self.redis
return self.redis
async def reset_redis(redis: Redis,
loop: AbstractEventLoop,
hard: bool = False) -> None:
if hard:
await redis.flushall()
else:
await redis.delete(q_key, info_key, seen_key, scope_key, done_key),
await redis.hset(info_key, "crawl_depth", 2),
for url in default_seed_list:
await asyncio.gather(
redis.rpush(q_key, ujson.dumps({
"url": url,
"depth": 0
})),
redis.sadd(seen_key, url),
loop=loop,
)
async def test_context_manager(self, r: aioredis.Redis):
async with r.pubsub() as pubsub:
await pubsub.subscribe("foo")
assert pubsub.connection is not None
assert pubsub.connection is None
assert pubsub.channels == {}
assert pubsub.patterns == {}
async def _set_indexes(ctx: object, db: AioRedis):
"""
sets secondary indexes
@param ctx: indexess to set
@param db: redis instance
"""
coroutines = []
for index in [('status_code', rule_status_code_index),
('service_id', rule_service_id_index)]:
if index[0] in ctx:
coroutines.append(db.hset(index[1], ctx['_id'], ctx[index[0]]))
for index in [('rule_id', entry_rule_id_index),
('host', entry_host_index)]:
if index[0] in ctx:
coroutines.append(db.hset(index[1], ctx['_id'], ctx[index[0]]))
await Async.all(coroutines)
def _handle_event(self, connection: Redis, event_result):
done, pending = event_result
for task in done:
if task is self.closed_event_wait:
self.closed_event_wait = None
raise StopException("Closed event")
if task is self.add_event_wait:
self.add_callback_event = asyncio.Future()
self.add_event_wait = asyncio.ensure_future(self.add_callback_event)
tup = task.result()
ix = next((i for i, v in enumerate(self.callbacks) if v is tup), None)
if ix is None:
continue
asyncio.ensure_future(
asyncio.gather(*(connection.psubscribe(i) if i.is_pattern else connection.subscribe(i)
for i in map(operator.itemgetter(1), self.callbacks[ix:]))))
def _get_event(self, connection: Redis) -> Awaitable:
if self.closed_event_wait is None:
self.closed_event_wait = asyncio.ensure_future(connection.wait_closed())
if self.add_event_wait is None:
self.add_callback_event = asyncio.Future()
self.add_event_wait = asyncio.ensure_future(self.add_callback_event)
if len(self.callbacks):
self.add_callback_event.set_result(self.callbacks[0])
return asyncio.wait([self.add_event_wait, self.closed_event_wait], return_when=asyncio.FIRST_COMPLETED)
async def bulk_update(new_data: List[CurrencyRecord], invalidate_others: bool,
redis: Redis):
if invalidate_others:
await redis.flushdb()
tr = redis.multi_exec()
for one_line in new_data:
tr.set(str((one_line.source, one_line.dest)), one_line.price)
tr.set(str((one_line.dest, one_line.source)), 1. / one_line.price)
await tr.execute()
async def _get_password(cls, redis: Redis, config: Config,
name: str) -> Tuple[str, bytes]:
key = config.users_key.format(name=name)
multi = redis.multi_exec()
if config.users_hash is None:
multi.get(key)
else:
multi.hget(config.users_hash, key)
new_namespace = uuid.uuid4().hex.encode('ascii')
multi.hsetnx(config.namespace_hash, key, new_namespace)
multi.hget(config.namespace_hash, key)
value, _, namespace = await multi.execute()
if value is None:
raise InvalidAuth()
elif config.users_json:
value_obj = json.loads(value)
try:
return value_obj['password'], namespace
except KeyError as exc:
raise InvalidAuth() from exc
else:
return value.decode('utf-8'), namespace
