Add `client.get_fulluser` with cache

pull/1/head
hikariatama 2022-08-09 22:05:35 +00:00
parent ed3b8529d3
commit 7f53220524
3 changed files with 101 additions and 4 deletions

View File

@ -16,6 +16,7 @@
- Add possible cause of error in logs (module and method)
- Add `client.get_perms_cached` to cache native `client.get_permissions`
- Add `client.get_fullchannel` with cache
- Add `client.get_fulluser` with cache
- Add `exp` and `force` params to `client.get_perms_cached` and `client.get_entity`
- Add `exp` cached values check in `client.get_perms_cached` and `client.get_entity`
- Change errors format in web to more human-readable

View File

@ -56,7 +56,12 @@ from . import database, loader, utils, heroku
from .dispatcher import CommandDispatcher
from .translations import Translator
from .version import __version__
from .tl_cache import install_entity_caching, install_perms_caching, install_fullchannel_caching
from .tl_cache import (
install_entity_caching,
install_perms_caching,
install_fullchannel_caching,
install_fulluser_caching,
)
try:
from .web import core
@ -526,6 +531,7 @@ class Hikka:
install_entity_caching(client)
install_perms_caching(client)
install_fullchannel_caching(client)
install_fulluser_caching(client)
self.clients += [client]
except sqlite3.OperationalError:

View File

@ -11,11 +11,11 @@ import time
import asyncio
import logging
from typing import Optional, Union
from xml.dom.minidom import Entity
from telethon.hints import EntityLike
from telethon import TelegramClient
from telethon.tl.functions.channels import GetFullChannelRequest
from telethon.tl.types import ChannelFull
from telethon.tl.functions.users import GetFullUserRequest
from telethon.tl.types import ChannelFull, UserFull
logger = logging.getLogger(__name__)
@ -116,6 +116,33 @@ class CacheRecordFullChannel:
)
class CacheRecordFullUser:
def __init__(self, user_id: int, full_user: UserFull, exp: int):
self.user_id = user_id
self.full_user = full_user
self._exp = round(time.time() + exp)
self.ts = time.time()
def expired(self):
return self._exp < time.time()
def __eq__(self, record: "CacheRecordFullUser"):
return hash(record) == hash(self)
def __hash__(self):
return hash((self._hashable_entity, self._hashable_user))
def __str__(self):
return f"CacheRecordFullUser of {self.user_id}"
def __repr__(self):
return (
f"CacheRecordFullUser(channel_id={self.user_id}(...),"
f" exp={self._exp})"
)
def install_entity_caching(client: TelegramClient):
client._hikka_entity_cache = {}
@ -337,7 +364,7 @@ def install_fullchannel_caching(client: TelegramClient):
:param entity: Channel to fetch ChannelFull of
:param exp: Expiration time of the cache record and maximum time of already cached record
:param force: Whether to force refresh the cache (make API request)
:return: :obj:`FullChannel`
:return: :obj:`ChannelFull`
"""
if not hashable(entity):
try:
@ -385,3 +412,66 @@ def install_fullchannel_caching(client: TelegramClient):
client.get_fullchannel = get_fullchannel
asyncio.ensure_future(cleaner(client))
logger.debug("Monkeypatched client with fullchannel cacher")
def install_fulluser_caching(client: TelegramClient):
client._hikka_fulluser_cache = {}
async def get_fulluser(
entity: EntityLike,
exp: Optional[int] = 300,
force: Optional[bool] = False,
) -> ChannelFull:
"""
Gets the FullUserRequest and cache it
:param entity: User to fetch UserFull of
:param exp: Expiration time of the cache record and maximum time of already cached record
:param force: Whether to force refresh the cache (make API request)
:return: :obj:`UserFull`
"""
if not hashable(entity):
try:
hashable_entity = next(
getattr(entity, attr)
for attr in {"user_id", "chat_id", "id"}
if getattr(entity, attr, None)
)
except StopIteration:
logger.debug(
f"Can't parse hashable from {entity=}, using legacy fulluser request"
)
return await client(GetFullUserRequest(entity))
else:
hashable_entity = entity
if str(hashable_entity).isdigit() and int(hashable_entity) < 0:
hashable_entity = int(str(hashable_entity)[4:])
if (
not force
and client._hikka_fulluser_cache.get(hashable_entity)
and not client._hikka_fulluser_cache[hashable_entity].expired()
and client._hikka_fulluser_cache[hashable_entity].ts + exp > time.time()
):
return client._hikka_fulluser_cache[hashable_entity].full_channel
result = await client(GetFullUserRequest(entity))
client._hikka_fulluser_cache[hashable_entity] = CacheRecordFullUser(
hashable_entity,
result,
exp,
)
return result
async def cleaner(client: TelegramClient):
while True:
for user_id, record in client._hikka_fulluser_cache.copy().items():
if record.expired():
del client._hikka_fulluser_cache[user_id]
logger.debug(f"Cleaned outdated fulluser cache {user_id=}")
await asyncio.sleep(3)
client.get_fulluser = get_fulluser
asyncio.ensure_future(cleaner(client))
logger.debug("Monkeypatched client with fulluser cacher")