def test_cache_set_many(app, client):
cache.set_many({'key_1': 1, 'Key_2': 3})
assert cache.get('key_1') == 1
assert cache.get('key_2') == 3
with app.test_request_context('/test'):
assert 'key_1' in flask.g.cache_keys['set_many']
assert 'key_2' in flask.g.cache_keys['set_many']
def test_subscribe_forum_deletes_cache_keys(app, authed_client):
add_permissions(app, ForumPermissions.MODIFY_SUBSCRIPTIONS)
Forum.from_subscribed_user(1)
ForumSubscription.user_ids_from_forum(5)
response = authed_client.post('/subscriptions/forums/5')
assert response.status_code == 200
assert not cache.get(
ForumSubscription.__cache_key_users__.format(forum_id=5)
)
assert not cache.get(
ForumSubscription.__cache_key_of_user__.format(user_id=1)
)
assert ForumSubscription.user_ids_from_forum(5) == [3, 4, 1]
def latest_revision(cls, article_id: int, language_id: int = 1) -> int:
cache_key = cls.__cache_key_latest_id_of_article__.format(
article_id=article_id, language_id=language_id
)
revision_id = cache.get(cache_key)
if revision_id:
return cls.from_attrs(
revision_id=revision_id,
article_id=article_id,
language_id=language_id,
)
else:
latest_revision = (
cls.query.filter(
and_(
cls.article_id == article_id,
cls.language_id == language_id,
)
)
.order_by(cls.revision_id.desc()) # type: ignore
.limit(1)
.scalar()
)
if not latest_revision:
raise WikiNoRevisions
cache.set(cache_key, latest_revision.revision_id)
return latest_revision
def count(
cls,
*,
key: str,
attribute: InstrumentedAttribute,
filter: BinaryExpression = None,
) -> int:
"""
An abstracted function for counting a number of elements matching a query. If the
passed cache key exists, its value will be returned; otherwise, the passed query
will be ran and the resultant count cached and returned.
:param key: The cache key to check
:param attribute: The attribute to count; a model's column
:param filter: The SQLAlchemy filter expression
:return: The number of rows matching the query element
"""
count = cache.get(key)
if not isinstance(count, int):
query = cls._construct_query(
db.session.query(func.count(attribute)), filter)
count = query.scalar()
cache.set(key, count)
return count
def permissions(self) -> List[str]:
"""
A general function to get the permissions of a user from a permission
model and attributes of their user classes. Locked users are restricted
to the permissions defined for them in the config.
:param key: The cache key to cache the permissions under
:param model: The model to query custom permissions from
:param attr: The attribute of the userclasses that should be queried
"""
from core.permissions.models import SecondaryClass
from core.permissions.models import UserPermission
if self.locked: # Locked accounts have restricted permissions.
return app.config['LOCKED_ACCOUNT_PERMISSIONS']
key = self.__cache_key_permissions__.format(id=self.id)
permissions = cache.get(key)
if not permissions:
permissions = copy(self.user_class_model.permissions)
for class_ in SecondaryClass.from_user(self.id):
permissions += class_.permissions
permissions = set(permissions) # De-dupe
for perm, granted in UserPermission.from_user(self.id).items():
if not granted and perm in permissions:
permissions.remove(perm)
if granted and perm not in permissions:
permissions.add(perm)
cache.set(key, permissions)
return permissions
def test_thread_last_viewed_post_deleted(app, authed_client):
thread = ForumThread.from_pk(5)
last_post = thread.last_viewed_post
assert last_post.id == 3
assert last_post.thread_id == 5
assert 3 == cache.get(
ForumLastViewedPost.__cache_key__.format(thread_id=5, user_id=1))
def test_thread_last_viewed_post_cached(app, authed_client):
cache.set(ForumLastViewedPost.__cache_key__.format(thread_id=1, user_id=1),
2)
thread = ForumThread.from_pk(1)
last_post = thread.last_viewed_post
assert last_post.id == 2
assert last_post.thread_id == 3
assert 2 == cache.get(
ForumLastViewedPost.__cache_key__.format(thread_id=1, user_id=1))
def test_cache_model(app, authed_client):
"""Test that caching a model works."""
user = User.from_pk(1)
cache.cache_model(user, timeout=60)
user_data = cache.get('users_1')
assert user_data['id'] == 1
assert user_data['username'] == 'user_one'
assert user_data['enabled'] is True
assert user_data['inviter_id'] is None
def get_pks_of_many(
cls,
key: str = None,
filter: BinaryExpression = None,
order: BinaryExpression = None,
include_dead: bool = False,
expr_override: BinaryExpression = None,
) -> List[Union[int, str]]:
"""
Get a list of object IDs meeting query criteria. Fetching from the
cache with the provided cache key will be attempted first; if the cache
key does not exist then a query will be ran. Calls with ``include_dead=True`` are
saved under a different cache key. ``include_dead`` does not affect the query results
when passing ``expr_override``.
:param key: The cache key to check (and return if present)
:param filter: A SQLAlchemy filter expression to be applied to the query
:param order: A SQLAlchemy order_by expression to be applied to the query
:param include_dead: Whether or not to include dead results in the IDs list
:param expr_override: If passed, this will override filter and order, and be
called verbatim in a ``db.session.execute`` if the cache
key does not exist
:return: A list of IDs
"""
key = f'{key}_incl_dead' if include_dead and key else key
pks = cache.get(key) if key else None
if not pks or not isinstance(pks, list):
if expr_override is not None:
pks = [
x[0] if len(x) == 1 else dict(x)
for x in db.session.execute(expr_override)
]
else:
primary_key = cls.get_primary_key()
if isinstance(primary_key, list):
query = cls._construct_query(
db.session.query(*(getattr(cls, k)
for k in primary_key)),
filter,
order,
)
else:
query = cls._construct_query(
db.session.query(getattr(cls, primary_key)),
filter,
order,
)
if not include_dead and cls.__deletion_attr__:
query = query.filter(
getattr(cls, cls.__deletion_attr__) == 'f')
pks = [
x[0] if len(x) == 1 else x._asdict() for x in query.all()
]
if key:
cache.set(key, pks)
return pks
def test_route():
cache.set('key3', 1)
cache.inc('key1')
cache.get('key2')
cache.ttl('key2')
cache.ttl('key3')
cache.get('key3')
cache.delete('key3')
cache.delete('key4')
assert cache.has('key1')
assert flask.g.cache_keys['inc'] == {'key1'}
assert flask.g.cache_keys['get'] == {'key3'}
assert flask.g.cache_keys['set'] == {'key3'}
assert flask.g.cache_keys['delete'] == {'key3'}
assert flask.g.cache_keys['ttl'] == {'key3'}
assert flask.g.cache_keys['has'] == {'key1'}
return flask.jsonify('complete')
def get_rules(section: str) -> dict:
if section not in SECTION_NAMES:
raise APIException(f'{section} is not a valid section of the rules.')
rules = cache.get(f'rules_{section}')
if not rules:
filename = os.path.join(os.path.dirname(__file__), f'{section}.json')
with open(filename, 'r') as f:
rules = json.load(f)
cache.set(f'rules_{section}', rules, 0)
return rules
def test_get_rules_cache(app, authed_client, monkeypatch):
add_permissions(app, RulePermissions.VIEW)
authed_client.get('/rules/golden') # cache
monkeypatch.setattr('rules.os', None)
response = authed_client.get('/rules/golden').get_json()
assert isinstance(response['response'], dict)
assert 'id' in response['response']
assert response['response']['rules'][0]['rules'][0]['number'] == '1.1'
assert cache.get('rules_golden')
assert cache.ttl('rules_golden') is None
def test_forum_thread_subscriptions_cache_keys_user_ids(app, authed_client):
user_ids = ForumThreadSubscription.user_ids_from_thread(4) # noqa
cache.set(
ForumThreadSubscription.__cache_key_of_user__.format(user_id=1),
[14, 23],
)
cache.set(
ForumThreadSubscription.__cache_key_users__.format(thread_id=4),
[3, 4, 5],
)
assert 3 == len(
cache.get(
ForumThreadSubscription.__cache_key_users__.format(thread_id=4)))
ForumThreadSubscription.clear_cache_keys(user_ids=[1, 2])
assert 3 == len(
cache.get(
ForumThreadSubscription.__cache_key_users__.format(thread_id=4)))
assert not cache.get(
ForumThreadSubscription.__cache_key_of_user__.format(user_id=1))
def test_unsubscribe_thread_deletes_cache_keys(app, authed_client):
add_permissions(app, ForumPermissions.MODIFY_SUBSCRIPTIONS)
ForumThread.from_subscribed_user(1)
ForumThreadSubscription.user_ids_from_thread(4)
assert cache.get(
ForumThreadSubscription.__cache_key_users__.format(thread_id=4)
)
assert cache.get(
ForumThreadSubscription.__cache_key_of_user__.format(user_id=1)
)
response = authed_client.delete('/subscriptions/threads/4')
assert response.status_code == 200
assert not cache.get(
ForumThreadSubscription.__cache_key_users__.format(thread_id=4)
)
assert not cache.get(
ForumThreadSubscription.__cache_key_of_user__.format(user_id=1)
)
assert ForumThreadSubscription.user_ids_from_thread(4) == [2]
def from_cache(cls, key: str, *, query: BaseQuery = None) -> Optional[PKB]:
data = cache.get(key)
obj = cls._create_obj_from_cache(data)
if obj:
return obj
else:
cache.delete(key)
if query:
obj = query.scalar()
cache.cache_model(obj)
return obj
def wrapper(*args, **kwargs):
arg = '-'.join(map(str, args))
key = kw + '-' + arg
data = cache.get(key)
if not data:
# print key, 'notin'
data = func(*args, **kwargs)
if timeout:
cache.set(key, data, timeout)
else:
cache.set(key, data)
# else:
# print key, 'bingo'
return data
def test_get_user_self_and_caches(app, authed_client):
add_permissions(app, 'users_view')
response = authed_client.get('/users/1')
check_json_response(
response,
{
'id': 1,
'username': '******',
'secondary_classes': ['FLS']
},
)
assert response.status_code == 200
user_data = cache.get('users_1')
assert user_data['user_class_id'] == 1
assert user_data['id'] == 1
assert user_data['username'] == 'user_one'
def update_api_key(api_key: APIKey) -> None:
"""
Update the provided session or api key's last seen times,
user agent, and IP fields.
:param session_key: The session or API key to update
"""
cache_key = f'{api_key.cache_key}_updated'
if (
not cache.get(cache_key)
or api_key.ip != flask.request.remote_addr
or api_key.user_agent != flask.request.headers.get('User-Agent')
):
api_key.last_used = datetime.utcnow().replace(tzinfo=pytz.utc)
api_key.user_agent = flask.request.headers.get('User-Agent')
api_key.ip = flask.request.remote_addr
db.session.commit()
cache.delete(api_key.cache_key)
cache.set(
cache_key, 1, timeout=60 * 2
) # 2 minute wait before next update
def post_from_attrs(cls, thread_id, user_id):
cache_key = cls.__cache_key__.format(
thread_id=thread_id, user_id=user_id
)
post_id = cache.get(cache_key)
if not post_id:
last_viewed = cls.query.filter(
and_((cls.thread_id == thread_id), (cls.user_id == user_id))
).scalar()
post_id = last_viewed.post_id if last_viewed else None
post = ForumPost.from_pk(post_id, include_dead=True)
if not post:
return None
if post.deleted: # Get the last non-deleted and read post.
post = ForumPost.from_query(
filter=and_(
(ForumPost.thread_id == thread_id),
(ForumPost.id < post.id),
(ForumPost.deleted == 'f'),
),
order=ForumPost.id.desc(),
) # type: ignore
cache.set(cache_key, post.id if post else None)
return post
def get_col_from_many(
cls,
*,
column: InstrumentedAttribute,
key: str = None,
filter: BinaryExpression = None,
order: BinaryExpression = None,
) -> List[Any]:
"""
Get the values of a specific column from every row in the database.
:param column: The desired column
:param key: A cache key to save the resultant list in
:param filter: Filters to apply to the database query
:param order: How to order the values from the database query
:return: A list of values from the column
"""
values = cache.get(key) if key else None
if values is None:
query = cls._construct_query(db.session.query(column), filter,
order)
values = [x[0] for x in query.all()]
cache.set(key, values)
return values
def _get_saved_cookie(self):
key = 'yugong.toutiao.cookie.%s' % config.get('app.toutiao.account')
return cache.get(key)
def test_thread_last_viewed_post_none(app, authed_client):
thread = ForumThread.from_pk(1)
assert thread.last_viewed_post is None
assert not cache.get(
ForumLastViewedPost.__cache_key__.format(thread_id=1, user_id=1))
def test_modify_poll_unfeature(app, authed_client):
add_permissions(app, 'forums_view', 'modify_forum_polls')
authed_client.put('/polls/3', data=json.dumps({'featured': False}))
assert ForumPoll.get_featured() is None
assert not cache.get(ForumPoll.__cache_key_featured__)
def test_thread_subscriptions(app, authed_client):
threads = ForumThread.from_subscribed_user(1)
assert all(t.id in {1, 3, 4} for t in threads)
assert {1, 3, 4} == set(
cache.get(
ForumThreadSubscription.__cache_key_of_user__.format(user_id=1)))
def test_from_cache_bad(app, client):
"""Receiving bad model data from the cache; delete it."""
cache.set('users_1', {'id': 2, 'username': '******'}, timeout=60)
assert not User.from_cache('users_1')
assert not cache.get('users_1')
