core/cache/cache.py

331 lines
12 KiB
Python
Raw Normal View History

2024-06-11 14:51:34 +00:00
import asyncio
2024-02-27 12:40:53 +00:00
import json
2024-05-30 16:16:50 +00:00
from typing import List
2024-08-12 08:00:01 +00:00
from sqlalchemy import and_, join, select
2024-08-06 16:59:27 +00:00
from orm.author import Author, AuthorFollower
2024-08-06 17:20:20 +00:00
from orm.shout import Shout, ShoutAuthor, ShoutTopic
2024-08-12 08:00:01 +00:00
from orm.topic import Topic, TopicFollower
2024-05-06 21:06:31 +00:00
from services.db import local_session
2024-08-07 06:51:09 +00:00
from services.redis import redis
2024-08-12 08:00:01 +00:00
from utils.encoders import CustomJSONEncoder
2024-08-07 05:57:56 +00:00
from utils.logger import root_logger as logger
2024-02-25 13:43:04 +00:00
2024-02-27 12:40:53 +00:00
DEFAULT_FOLLOWS = {
2024-04-17 15:32:23 +00:00
"topics": [],
"authors": [],
2024-08-07 06:51:09 +00:00
"shouts": [],
2024-04-17 15:32:23 +00:00
"communities": [{"id": 1, "name": "Дискурс", "slug": "discours", "pic": ""}],
2024-02-27 12:40:53 +00:00
}
2024-02-25 13:43:04 +00:00
2024-08-07 06:51:09 +00:00
# Cache topic data
2024-05-20 22:40:57 +00:00
async def cache_topic(topic: dict):
2024-05-30 16:15:11 +00:00
payload = json.dumps(topic, cls=CustomJSONEncoder)
2024-08-07 06:51:09 +00:00
# Cache by id and slug for quick access
2024-08-06 16:55:27 +00:00
await asyncio.gather(
2024-08-06 17:55:19 +00:00
redis.execute("SET", f"topic:id:{topic['id']}", payload),
redis.execute("SET", f"topic:slug:{topic['slug']}", payload),
2024-08-06 16:55:27 +00:00
)
2024-05-20 22:40:57 +00:00
2024-08-07 06:51:09 +00:00
# Cache author data
2024-04-09 08:17:32 +00:00
async def cache_author(author: dict):
2024-03-06 19:00:37 +00:00
payload = json.dumps(author, cls=CustomJSONEncoder)
2024-08-07 06:51:09 +00:00
# Cache author data by user and id
2024-08-06 16:55:27 +00:00
await asyncio.gather(
2024-08-06 17:23:23 +00:00
redis.execute("SET", f"author:user:{author['user'].strip()}", str(author["id"])),
redis.execute("SET", f"author:id:{author['id']}", payload),
2024-08-06 16:55:27 +00:00
)
2024-03-12 12:50:57 +00:00
2024-02-27 12:40:53 +00:00
2024-08-07 06:51:09 +00:00
# Cache follows data
2024-05-30 09:49:46 +00:00
async def cache_follows(follower_id: int, entity_type: str, entity_id: int, is_insert=True):
2024-08-06 16:55:27 +00:00
key = f"author:follows-{entity_type}s:{follower_id}"
2024-08-08 14:54:15 +00:00
follows_str = await redis.execute("get", key)
2024-08-07 06:51:09 +00:00
follows = json.loads(follows_str) if follows_str else DEFAULT_FOLLOWS[entity_type]
2024-05-20 22:40:57 +00:00
if is_insert:
2024-08-06 16:55:27 +00:00
if entity_id not in follows:
follows.append(entity_id)
2024-05-20 22:40:57 +00:00
else:
follows = [eid for eid in follows if eid != entity_id]
2024-08-07 04:18:49 +00:00
await redis.execute("set", key, json.dumps(follows, cls=CustomJSONEncoder))
2024-08-07 06:51:09 +00:00
await update_follower_stat(follower_id, entity_type, len(follows))
2024-05-20 22:40:57 +00:00
2024-08-06 16:55:27 +00:00
2024-08-07 06:51:09 +00:00
# Update follower statistics
2024-08-06 16:55:27 +00:00
async def update_follower_stat(follower_id, entity_type, count):
follower_key = f"author:id:{follower_id}"
2024-08-08 14:54:15 +00:00
follower_str = await redis.execute("get", follower_key)
2024-08-06 16:55:27 +00:00
follower = json.loads(follower_str) if follower_str else None
2024-08-06 16:45:42 +00:00
if follower:
2024-08-06 16:55:27 +00:00
follower["stat"] = {f"{entity_type}s": count}
2024-05-20 22:40:57 +00:00
await cache_author(follower)
2024-02-29 21:51:49 +00:00
2024-08-07 06:51:09 +00:00
# Get author from cache
2024-08-09 06:37:06 +00:00
async def get_cached_author(author_id: int, get_with_stat):
2024-08-06 16:55:27 +00:00
author_key = f"author:id:{author_id}"
2024-08-08 14:54:15 +00:00
result = await redis.execute("get", author_key)
2024-08-06 16:55:27 +00:00
if result:
return json.loads(result)
2024-08-07 06:51:09 +00:00
# Load from database if not found in cache
2024-08-09 06:37:06 +00:00
q = select(Author).where(Author.id == author_id)
author = get_with_stat(q)
if author:
await cache_author(author.dict())
return author.dict()
2024-08-06 16:55:27 +00:00
return None
2024-08-07 06:51:09 +00:00
# Function to get cached topic
async def get_cached_topic(topic_id: int):
"""
Fetch topic data from cache or database by id.
Args:
topic_id (int): The identifier for the topic.
Returns:
dict: Topic data or None if not found.
"""
topic_key = f"topic:id:{topic_id}"
2024-08-08 14:54:15 +00:00
cached_topic = await redis.execute("get", topic_key)
2024-08-07 06:51:09 +00:00
if cached_topic:
return json.loads(cached_topic)
# If not in cache, fetch from the database
with local_session() as session:
topic = session.execute(select(Topic).where(Topic.id == topic_id)).scalar_one_or_none()
if topic:
topic_dict = topic.dict()
2024-08-08 14:54:15 +00:00
await redis.execute("set", topic_key, json.dumps(topic_dict, cls=CustomJSONEncoder))
2024-08-07 06:51:09 +00:00
return topic_dict
return None
# Get topic by slug from cache
2024-08-08 14:46:25 +00:00
async def get_cached_topic_by_slug(slug: str, get_with_stat):
2024-08-06 16:55:27 +00:00
topic_key = f"topic:slug:{slug}"
2024-08-08 14:54:15 +00:00
result = await redis.execute("get", topic_key)
2024-08-06 16:55:27 +00:00
if result:
return json.loads(result)
2024-08-07 06:51:09 +00:00
# Load from database if not found in cache
2024-08-08 14:46:25 +00:00
topic_query = select(Topic).where(Topic.slug == slug)
topic = get_with_stat(topic_query)
if topic:
topic_dict = topic.dict()
await cache_topic(topic_dict)
return topic_dict
2024-08-06 16:55:27 +00:00
return None
2024-08-06 16:45:42 +00:00
2024-08-07 06:51:09 +00:00
# Get list of authors by ID from cache
2024-08-06 16:45:42 +00:00
async def get_cached_authors_by_ids(author_ids: List[int]) -> List[dict]:
2024-08-07 06:51:09 +00:00
# Fetch all author data concurrently
2024-08-06 16:55:27 +00:00
keys = [f"author:id:{author_id}" for author_id in author_ids]
2024-08-08 14:54:15 +00:00
results = await asyncio.gather(*(redis.execute("get", key) for key in keys))
2024-08-06 16:55:27 +00:00
authors = [json.loads(result) if result else None for result in results]
2024-08-07 06:51:09 +00:00
# Load missing authors from database and cache
2024-08-06 16:55:27 +00:00
missing_indices = [index for index, author in enumerate(authors) if author is None]
if missing_indices:
missing_ids = [author_ids[index] for index in missing_indices]
2024-05-30 17:23:32 +00:00
with local_session() as session:
2024-08-06 16:45:42 +00:00
query = select(Author).where(Author.id.in_(missing_ids))
2024-08-06 16:55:27 +00:00
missing_authors = session.execute(query).scalars().all()
await asyncio.gather(*(cache_author(author.dict()) for author in missing_authors))
2024-08-07 06:51:09 +00:00
for index, author in zip(missing_indices, missing_authors):
authors[index] = author.dict()
2024-08-06 16:45:42 +00:00
return authors
2024-08-06 16:59:27 +00:00
2024-08-06 17:05:24 +00:00
2024-08-06 16:59:27 +00:00
async def get_cached_topic_followers(topic_id: int):
2024-08-08 15:00:36 +00:00
"""
Получает подписчиков темы по ID, используя кеш Redis.
Если данные отсутствуют в кеше, извлекает из базы данных и кеширует их.
2024-08-06 16:59:27 +00:00
2024-08-08 15:00:36 +00:00
:param topic_id: Идентификатор темы, подписчиков которой необходимо получить.
:return: Список подписчиков темы, каждый элемент представляет собой словарь с ID и именем автора.
"""
try:
# Попытка получить данные из кеша
cached = await redis.get(f"topic:followers:{topic_id}")
if cached:
2024-08-09 04:26:04 +00:00
followers_ids = json.loads(cached)
logger.debug(f"Cached {len(followers_ids)} followers for topic #{topic_id}")
followers = await get_cached_authors_by_ids(followers_ids)
2024-08-08 15:00:36 +00:00
return followers
# Если данные не найдены в кеше, загрузка из базы данных
async with local_session() as session:
result = await session.execute(
session.query(Author.id)
.join(TopicFollower, TopicFollower.follower == Author.id)
.filter(TopicFollower.topic == topic_id)
)
followers_ids = [f[0] for f in result.scalars().all()]
# Кеширование результатов
await redis.set(f"topic:followers:{topic_id}", json.dumps(followers_ids))
# Получение подробной информации о подписчиках по их ID
followers = await get_cached_authors_by_ids(followers_ids)
2024-08-09 04:14:33 +00:00
logger.debug(followers)
2024-08-08 15:00:36 +00:00
return followers
except Exception as e:
2024-08-09 04:26:04 +00:00
logger.error(f"Ошибка при получении подписчиков для темы #{topic_id}: {str(e)}")
2024-08-08 15:00:36 +00:00
return []
2024-08-06 16:59:27 +00:00
2024-08-06 17:05:24 +00:00
2024-08-07 06:51:09 +00:00
# Get cached author followers
2024-08-06 16:59:27 +00:00
async def get_cached_author_followers(author_id: int):
2024-08-07 06:51:09 +00:00
# Check cache for data
2024-08-08 14:54:15 +00:00
cached = await redis.execute("get", f"author:followers:{author_id}")
2024-08-06 16:59:27 +00:00
if cached:
followers_ids = json.loads(cached)
followers = await get_cached_authors_by_ids(followers_ids)
2024-08-09 04:22:55 +00:00
logger.debug(f"Cached followers for author #{author_id}: {len(followers)}")
2024-08-06 16:59:27 +00:00
return followers
2024-08-07 06:51:09 +00:00
# Query database if cache is empty
2024-08-06 16:59:27 +00:00
with local_session() as session:
2024-08-06 17:05:24 +00:00
followers_ids = [
f[0]
for f in session.query(Author.id)
.join(AuthorFollower, AuthorFollower.follower == Author.id)
.filter(AuthorFollower.author == author_id, Author.id != author_id)
.all()
]
2024-08-06 17:23:23 +00:00
await redis.execute("SET", f"author:followers:{author_id}", json.dumps(followers_ids))
2024-08-06 16:59:27 +00:00
followers = await get_cached_authors_by_ids(followers_ids)
return followers
2024-08-06 17:05:24 +00:00
2024-08-07 06:51:09 +00:00
# Get cached follower authors
2024-08-06 16:59:27 +00:00
async def get_cached_follower_authors(author_id: int):
2024-08-07 06:51:09 +00:00
# Attempt to retrieve authors from cache
2024-08-08 14:54:15 +00:00
cached = await redis.execute("get", f"author:follows-authors:{author_id}")
2024-08-06 16:59:27 +00:00
if cached:
authors_ids = json.loads(cached)
else:
2024-08-07 06:51:09 +00:00
# Query authors from database
2024-08-06 16:59:27 +00:00
with local_session() as session:
2024-08-06 17:05:24 +00:00
authors_ids = [
a[0]
for a in session.execute(
select(Author.id)
.select_from(join(Author, AuthorFollower, Author.id == AuthorFollower.author))
.where(AuthorFollower.follower == author_id)
).all()
]
2024-08-06 17:23:23 +00:00
await redis.execute("SET", f"author:follows-authors:{author_id}", json.dumps(authors_ids))
2024-08-06 16:59:27 +00:00
authors = await get_cached_authors_by_ids(authors_ids)
return authors
2024-08-06 17:05:24 +00:00
2024-08-07 06:51:09 +00:00
# Get cached follower topics
2024-08-06 16:59:27 +00:00
async def get_cached_follower_topics(author_id: int):
2024-08-07 06:51:09 +00:00
# Attempt to retrieve topics from cache
2024-08-08 14:54:15 +00:00
cached = await redis.execute("get", f"author:follows-topics:{author_id}")
2024-08-06 16:59:27 +00:00
if cached:
topics_ids = json.loads(cached)
else:
2024-08-07 06:51:09 +00:00
# Load topics from database and cache them
2024-08-06 16:59:27 +00:00
with local_session() as session:
2024-08-06 17:05:24 +00:00
topics_ids = [
t[0]
for t in session.query(Topic.id)
.join(TopicFollower, TopicFollower.topic == Topic.id)
.where(TopicFollower.follower == author_id)
.all()
]
2024-08-06 17:23:23 +00:00
await redis.execute("SET", f"author:follows-topics:{author_id}", json.dumps(topics_ids))
2024-08-06 16:59:27 +00:00
topics = []
for topic_id in topics_ids:
2024-08-08 14:54:15 +00:00
topic_str = await redis.execute("get", f"topic:id:{topic_id}")
2024-08-06 16:59:27 +00:00
if topic_str:
topic = json.loads(topic_str)
if topic and topic not in topics:
topics.append(topic)
logger.debug(f"Cached topics for author#{author_id}: {len(topics)}")
return topics
2024-08-06 17:05:24 +00:00
2024-08-07 06:51:09 +00:00
# Get author by user ID from cache
2024-08-06 17:05:24 +00:00
async def get_cached_author_by_user_id(user_id: str):
"""
2024-08-07 06:51:09 +00:00
Retrieve author information by user_id, checking the cache first, then the database.
2024-08-06 17:05:24 +00:00
Args:
2024-08-07 06:51:09 +00:00
user_id (str): The user identifier for which to retrieve the author.
2024-08-06 17:05:24 +00:00
Returns:
2024-08-07 06:51:09 +00:00
dict: Dictionary with author data or None if not found.
2024-08-06 17:05:24 +00:00
"""
2024-08-07 06:51:09 +00:00
# Attempt to find author ID by user_id in Redis cache
2024-08-08 14:54:15 +00:00
author_id = await redis.execute("get", f"author:user:{user_id.strip()}")
2024-08-06 17:05:24 +00:00
if author_id:
2024-08-07 06:51:09 +00:00
# If ID is found, get full author data by ID
2024-08-08 14:54:15 +00:00
author_data = await redis.execute("get", f"author:id:{author_id}")
2024-08-06 17:05:24 +00:00
if author_data:
return json.loads(author_data)
2024-08-07 06:51:09 +00:00
# If data is not found in cache, query the database
2024-08-06 17:05:24 +00:00
with local_session() as session:
author = session.execute(select(Author).where(Author.user == user_id)).scalar_one_or_none()
if author:
2024-08-07 06:51:09 +00:00
# Cache the retrieved author data
2024-08-06 17:05:24 +00:00
author_dict = author.dict()
await asyncio.gather(
2024-08-06 17:23:23 +00:00
redis.execute("SET", f"author:user:{user_id.strip()}", str(author.id)),
redis.execute("SET", f"author:id:{author.id}", json.dumps(author_dict)),
2024-08-06 17:05:24 +00:00
)
return author_dict
2024-08-07 06:51:09 +00:00
# Return None if author is not found
2024-08-06 17:05:24 +00:00
return None
2024-08-06 17:20:20 +00:00
2024-08-06 17:55:19 +00:00
2024-08-07 06:51:09 +00:00
# Get cached topic authors
2024-08-06 17:20:20 +00:00
async def get_cached_topic_authors(topic_id: int):
"""
2024-08-07 06:51:09 +00:00
Retrieve a list of authors for a given topic, using cache or database.
2024-08-06 17:20:20 +00:00
Args:
2024-08-07 06:51:09 +00:00
topic_id (int): The identifier of the topic for which to retrieve authors.
2024-08-06 17:20:20 +00:00
Returns:
2024-08-07 06:51:09 +00:00
List[dict]: A list of dictionaries containing author data.
2024-08-06 17:20:20 +00:00
"""
2024-08-07 06:51:09 +00:00
# Attempt to get a list of author IDs from cache
2024-08-06 17:20:20 +00:00
rkey = f"topic:authors:{topic_id}"
2024-08-08 14:54:15 +00:00
cached_authors_ids = await redis.execute("get", rkey)
2024-08-06 17:20:20 +00:00
if cached_authors_ids:
authors_ids = json.loads(cached_authors_ids)
else:
2024-08-07 06:51:09 +00:00
# If cache is empty, get data from the database
2024-08-06 17:20:20 +00:00
with local_session() as session:
query = (
select(ShoutAuthor.author)
.select_from(join(ShoutTopic, Shout, ShoutTopic.shout == Shout.id))
.join(ShoutAuthor, ShoutAuthor.shout == Shout.id)
.where(and_(ShoutTopic.topic == topic_id, Shout.published_at.is_not(None), Shout.deleted_at.is_(None)))
)
2024-08-06 17:55:19 +00:00
authors_ids = [author_id for (author_id,) in session.execute(query).all()]
2024-08-07 06:51:09 +00:00
# Cache the retrieved author IDs
2024-08-07 04:18:49 +00:00
await redis.execute("set", rkey, json.dumps(authors_ids))
2024-08-06 17:20:20 +00:00
2024-08-07 06:51:09 +00:00
# Retrieve full author details from cached IDs
2024-08-06 17:20:20 +00:00
if authors_ids:
authors = await get_cached_authors_by_ids(authors_ids)
logger.debug(f"Topic#{topic_id} authors fetched and cached: {len(authors)} authors found.")
return authors
return []