mirror of
https://github.com/ada-dmitry/kurator_bot.git
synced 2026-09-25 16:50:15 +00:00
Init
This commit is contained in:
ada-dmitry
@@ -0,0 +1,13 @@
|
||||
from telebot.asyncio_storage.memory_storage import StateMemoryStorage
|
||||
from telebot.asyncio_storage.redis_storage import StateRedisStorage
|
||||
from telebot.asyncio_storage.pickle_storage import StatePickleStorage
|
||||
from telebot.asyncio_storage.base_storage import StateDataContext, StateStorageBase
|
||||
|
||||
|
||||
__all__ = [
|
||||
"StateStorageBase",
|
||||
"StateDataContext",
|
||||
"StateMemoryStorage",
|
||||
"StateRedisStorage",
|
||||
"StatePickleStorage",
|
||||
]
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,123 @@
|
||||
import copy
|
||||
|
||||
|
||||
class StateStorageBase:
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
async def set_data(self, chat_id, user_id, key, value):
|
||||
"""
|
||||
Set data for a user in a particular chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_data(self, chat_id, user_id):
|
||||
"""
|
||||
Get data for a user in a particular chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def set_state(self, chat_id, user_id, state):
|
||||
"""
|
||||
Set state for a particular user.
|
||||
|
||||
! Note that you should create a
|
||||
record if it does not exist, and
|
||||
if a record with state already exists,
|
||||
you need to update a record.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def delete_state(self, chat_id, user_id):
|
||||
"""
|
||||
Delete state for a particular user.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def reset_data(self, chat_id, user_id):
|
||||
"""
|
||||
Reset data for a particular user in a chat.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def get_state(self, chat_id, user_id):
|
||||
raise NotImplementedError
|
||||
|
||||
def get_interactive_data(self, chat_id, user_id):
|
||||
"""
|
||||
Should be sync, but should provide a context manager
|
||||
with __aenter__ and __aexit__ methods.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
async def save(self, chat_id, user_id, data):
|
||||
raise NotImplementedError
|
||||
|
||||
def _get_key(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
prefix: str,
|
||||
separator: str,
|
||||
business_connection_id: str = None,
|
||||
message_thread_id: int = None,
|
||||
bot_id: int = None,
|
||||
) -> str:
|
||||
"""
|
||||
Convert parameters to a key.
|
||||
"""
|
||||
params = [prefix]
|
||||
if bot_id:
|
||||
params.append(str(bot_id))
|
||||
if business_connection_id:
|
||||
params.append(business_connection_id)
|
||||
if message_thread_id:
|
||||
params.append(str(message_thread_id))
|
||||
params.append(str(chat_id))
|
||||
params.append(str(user_id))
|
||||
|
||||
return separator.join(params)
|
||||
|
||||
|
||||
class StateDataContext:
|
||||
"""
|
||||
Class for data.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
obj,
|
||||
chat_id,
|
||||
user_id,
|
||||
business_connection_id=None,
|
||||
message_thread_id=None,
|
||||
bot_id=None,
|
||||
):
|
||||
self.obj = obj
|
||||
self.data = None
|
||||
self.chat_id = chat_id
|
||||
self.user_id = user_id
|
||||
self.bot_id = bot_id
|
||||
self.business_connection_id = business_connection_id
|
||||
self.message_thread_id = message_thread_id
|
||||
|
||||
async def __aenter__(self):
|
||||
data = await self.obj.get_data(
|
||||
chat_id=self.chat_id,
|
||||
user_id=self.user_id,
|
||||
business_connection_id=self.business_connection_id,
|
||||
message_thread_id=self.message_thread_id,
|
||||
bot_id=self.bot_id,
|
||||
)
|
||||
self.data = copy.deepcopy(data)
|
||||
return self.data
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
return await self.obj.save(
|
||||
self.chat_id,
|
||||
self.user_id,
|
||||
self.data,
|
||||
self.business_connection_id,
|
||||
self.message_thread_id,
|
||||
self.bot_id,
|
||||
)
|
||||
@@ -0,0 +1,225 @@
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateDataContext
|
||||
from typing import Optional, Union
|
||||
|
||||
|
||||
class StateMemoryStorage(StateStorageBase):
|
||||
"""
|
||||
Memory storage for states.
|
||||
|
||||
Stores states in memory as a dictionary.
|
||||
|
||||
.. code-block:: python3
|
||||
|
||||
storage = StateMemoryStorage()
|
||||
bot = AsyncTeleBot(token, storage=storage)
|
||||
|
||||
:param separator: Separator for keys, default is ":".
|
||||
:type separator: Optional[str]
|
||||
|
||||
:param prefix: Prefix for keys, default is "telebot".
|
||||
:type prefix: Optional[str]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, separator: Optional[str] = ":", prefix: Optional[str] = "telebot"
|
||||
) -> None:
|
||||
self.separator = separator
|
||||
self.prefix = prefix
|
||||
if not self.prefix:
|
||||
raise ValueError("Prefix cannot be empty")
|
||||
|
||||
self.data = (
|
||||
{}
|
||||
) # key: telebot:bot_id:business_connection_id:message_thread_id:chat_id:user_id
|
||||
|
||||
async def set_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
state: str,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
if hasattr(state, "name"):
|
||||
state = state.name
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
self.data[_key] = {"state": state, "data": {}}
|
||||
else:
|
||||
self.data[_key]["state"] = state
|
||||
|
||||
return True
|
||||
|
||||
async def get_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Union[str, None]:
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
return None
|
||||
|
||||
return self.data[_key]["state"]
|
||||
|
||||
async def delete_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
return False
|
||||
|
||||
del self.data[_key]
|
||||
return True
|
||||
|
||||
async def set_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
key: str,
|
||||
value: Union[str, int, float, dict],
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
raise RuntimeError(f"MemoryStorage: key {_key} does not exist.")
|
||||
self.data[_key]["data"][key] = value
|
||||
return True
|
||||
|
||||
async def get_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> dict:
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
return self.data.get(_key, {}).get("data", {})
|
||||
|
||||
async def reset_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
return False
|
||||
self.data[_key]["data"] = {}
|
||||
return True
|
||||
|
||||
def get_interactive_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Optional[dict]:
|
||||
return StateDataContext(
|
||||
self,
|
||||
chat_id=chat_id,
|
||||
user_id=user_id,
|
||||
business_connection_id=business_connection_id,
|
||||
message_thread_id=message_thread_id,
|
||||
bot_id=bot_id,
|
||||
)
|
||||
|
||||
async def save(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
data: dict,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
|
||||
if self.data.get(_key) is None:
|
||||
return False
|
||||
self.data[_key]["data"] = data
|
||||
return True
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"<StateMemoryStorage: {self.data}>"
|
||||
@@ -0,0 +1,274 @@
|
||||
aiofiles_installed = True
|
||||
try:
|
||||
import aiofiles
|
||||
except ImportError:
|
||||
aiofiles_installed = False
|
||||
|
||||
import os
|
||||
import pickle
|
||||
import asyncio
|
||||
from typing import Optional, Union, Callable, Any
|
||||
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateDataContext
|
||||
|
||||
|
||||
def with_lock(func: Callable) -> Callable:
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
async with self.lock:
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class StatePickleStorage(StateStorageBase):
|
||||
"""
|
||||
State storage based on pickle file.
|
||||
|
||||
.. warning::
|
||||
|
||||
This storage is not recommended for production use.
|
||||
Data may be corrupted. If you face a case where states do not work as expected,
|
||||
try to use another storage.
|
||||
|
||||
.. code-block:: python3
|
||||
|
||||
storage = StatePickleStorage()
|
||||
bot = AsyncTeleBot(token, storage=storage)
|
||||
|
||||
:param file_path: Path to file where states will be stored.
|
||||
:type file_path: str
|
||||
|
||||
:param prefix: Prefix for keys, default is "telebot".
|
||||
:type prefix: Optional[str]
|
||||
|
||||
:param separator: Separator for keys, default is ":".
|
||||
:type separator: Optional[str]
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str = "./.state-save/states.pkl",
|
||||
prefix="telebot",
|
||||
separator: Optional[str] = ":",
|
||||
) -> None:
|
||||
|
||||
if not aiofiles_installed:
|
||||
raise ImportError("Please install aiofiles using `pip install aiofiles`")
|
||||
|
||||
self.file_path = file_path
|
||||
self.prefix = prefix
|
||||
self.separator = separator
|
||||
self.lock = asyncio.Lock()
|
||||
self.create_dir()
|
||||
|
||||
async def _read_from_file(self) -> dict:
|
||||
async with aiofiles.open(self.file_path, "rb") as f:
|
||||
data = await f.read()
|
||||
return pickle.loads(data)
|
||||
|
||||
async def _write_to_file(self, data: dict) -> None:
|
||||
async with aiofiles.open(self.file_path, "wb") as f:
|
||||
await f.write(pickle.dumps(data))
|
||||
|
||||
def create_dir(self):
|
||||
"""
|
||||
Create directory .save-handlers.
|
||||
"""
|
||||
dirs, filename = os.path.split(self.file_path)
|
||||
os.makedirs(dirs, exist_ok=True)
|
||||
if not os.path.isfile(self.file_path):
|
||||
with open(self.file_path, "wb") as file:
|
||||
pickle.dump({}, file)
|
||||
|
||||
@with_lock
|
||||
async def set_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
state: str,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
if _key not in data:
|
||||
data[_key] = {"state": state, "data": {}}
|
||||
else:
|
||||
data[_key]["state"] = state
|
||||
await self._write_to_file(data)
|
||||
return True
|
||||
|
||||
@with_lock
|
||||
async def get_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Union[str, None]:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
return data.get(_key, {}).get("state")
|
||||
|
||||
@with_lock
|
||||
async def delete_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
if _key in data:
|
||||
del data[_key]
|
||||
await self._write_to_file(data)
|
||||
return True
|
||||
return False
|
||||
|
||||
@with_lock
|
||||
async def set_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
key: str,
|
||||
value: Union[str, int, float, dict],
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
state_data = data.get(_key, {})
|
||||
state_data["data"][key] = value
|
||||
if _key not in data:
|
||||
raise RuntimeError(f"StatePickleStorage: key {_key} does not exist.")
|
||||
else:
|
||||
data[_key]["data"][key] = value
|
||||
await self._write_to_file(data)
|
||||
return True
|
||||
|
||||
@with_lock
|
||||
async def get_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> dict:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
return data.get(_key, {}).get("data", {})
|
||||
|
||||
@with_lock
|
||||
async def reset_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
if _key in data:
|
||||
data[_key]["data"] = {}
|
||||
await self._write_to_file(data)
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_interactive_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Optional[dict]:
|
||||
return StateDataContext(
|
||||
self,
|
||||
chat_id=chat_id,
|
||||
user_id=user_id,
|
||||
business_connection_id=business_connection_id,
|
||||
message_thread_id=message_thread_id,
|
||||
bot_id=bot_id,
|
||||
)
|
||||
|
||||
@with_lock
|
||||
async def save(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
data: dict,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self._read_from_file()
|
||||
data[_key]["data"] = data
|
||||
await self._write_to_file(data)
|
||||
return True
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"StatePickleStorage({self.file_path}, {self.prefix})"
|
||||
@@ -0,0 +1,339 @@
|
||||
redis_installed = True
|
||||
try:
|
||||
import redis
|
||||
from redis.asyncio import Redis, ConnectionPool
|
||||
except ImportError:
|
||||
redis_installed = False
|
||||
|
||||
import json
|
||||
from typing import Optional, Union, Callable, Coroutine
|
||||
import asyncio
|
||||
|
||||
from telebot.asyncio_storage.base_storage import StateStorageBase, StateDataContext
|
||||
|
||||
|
||||
def async_with_lock(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]:
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
async with self.lock:
|
||||
return await func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def async_with_pipeline(func: Callable[..., Coroutine]) -> Callable[..., Coroutine]:
|
||||
async def wrapper(self, *args, **kwargs):
|
||||
async with self.redis.pipeline() as pipe:
|
||||
pipe.multi()
|
||||
result = await func(self, pipe, *args, **kwargs)
|
||||
await pipe.execute()
|
||||
return result
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
class StateRedisStorage(StateStorageBase):
|
||||
"""
|
||||
State storage based on Redis.
|
||||
|
||||
.. code-block:: python3
|
||||
|
||||
storage = StateRedisStorage(...)
|
||||
bot = AsyncTeleBot(token, storage=storage)
|
||||
|
||||
:param host: Redis host, default is "localhost".
|
||||
:type host: str
|
||||
|
||||
:param port: Redis port, default is 6379.
|
||||
:type port: int
|
||||
|
||||
:param db: Redis database, default is 0.
|
||||
:type db: int
|
||||
|
||||
:param password: Redis password, default is None.
|
||||
:type password: Optional[str]
|
||||
|
||||
:param prefix: Prefix for keys, default is "telebot".
|
||||
:type prefix: Optional[str]
|
||||
|
||||
:param redis_url: Redis URL, default is None.
|
||||
:type redis_url: Optional[str]
|
||||
|
||||
:param connection_pool: Redis connection pool, default is None.
|
||||
:type connection_pool: Optional[ConnectionPool]
|
||||
|
||||
:param separator: Separator for keys, default is ":".
|
||||
:type separator: Optional[str]
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host="localhost",
|
||||
port=6379,
|
||||
db=0,
|
||||
password=None,
|
||||
prefix="telebot",
|
||||
redis_url=None,
|
||||
connection_pool: "ConnectionPool" = None,
|
||||
separator: Optional[str] = ":",
|
||||
) -> None:
|
||||
|
||||
if not redis_installed:
|
||||
raise ImportError("Please install redis using `pip install redis`")
|
||||
|
||||
self.separator = separator
|
||||
self.prefix = prefix
|
||||
if not self.prefix:
|
||||
raise ValueError("Prefix cannot be empty")
|
||||
|
||||
if redis_url:
|
||||
self.redis = redis.asyncio.from_url(redis_url)
|
||||
elif connection_pool:
|
||||
self.redis = Redis(connection_pool=connection_pool)
|
||||
else:
|
||||
self.redis = Redis(host=host, port=port, db=db, password=password)
|
||||
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
@async_with_lock
|
||||
@async_with_pipeline
|
||||
async def set_state(
|
||||
self,
|
||||
pipe,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
state: str,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
if hasattr(state, "name"):
|
||||
state = state.name
|
||||
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
pipe.hget(_key, "data")
|
||||
result = await pipe.execute()
|
||||
data = result[0]
|
||||
if data is None:
|
||||
pipe.hset(_key, "data", json.dumps({}))
|
||||
|
||||
await pipe.hset(_key, "state", state)
|
||||
|
||||
return True
|
||||
|
||||
async def get_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Union[str, None]:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
state_bytes = await self.redis.hget(_key, "state")
|
||||
return state_bytes.decode("utf-8") if state_bytes else None
|
||||
|
||||
async def delete_state(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
result = await self.redis.delete(_key)
|
||||
return result > 0
|
||||
|
||||
@async_with_lock
|
||||
@async_with_pipeline
|
||||
async def set_data(
|
||||
self,
|
||||
pipe,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
key: str,
|
||||
value: Union[str, int, float, dict],
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await pipe.hget(_key, "data")
|
||||
data = await pipe.execute()
|
||||
data = data[0]
|
||||
if data is None:
|
||||
raise RuntimeError(f"StateRedisStorage: key {_key} does not exist.")
|
||||
else:
|
||||
data = json.loads(data)
|
||||
data[key] = value
|
||||
await pipe.hset(_key, "data", json.dumps(data))
|
||||
return True
|
||||
|
||||
async def get_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> dict:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
data = await self.redis.hget(_key, "data")
|
||||
return json.loads(data) if data else {}
|
||||
|
||||
@async_with_lock
|
||||
@async_with_pipeline
|
||||
async def reset_data(
|
||||
self,
|
||||
pipe,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
if await pipe.exists(_key):
|
||||
await pipe.hset(_key, "data", "{}")
|
||||
else:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_interactive_data(
|
||||
self,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> Optional[dict]:
|
||||
return StateDataContext(
|
||||
self,
|
||||
chat_id=chat_id,
|
||||
user_id=user_id,
|
||||
business_connection_id=business_connection_id,
|
||||
message_thread_id=message_thread_id,
|
||||
bot_id=bot_id,
|
||||
)
|
||||
|
||||
@async_with_lock
|
||||
@async_with_pipeline
|
||||
async def save(
|
||||
self,
|
||||
pipe,
|
||||
chat_id: int,
|
||||
user_id: int,
|
||||
data: dict,
|
||||
business_connection_id: Optional[str] = None,
|
||||
message_thread_id: Optional[int] = None,
|
||||
bot_id: Optional[int] = None,
|
||||
) -> bool:
|
||||
_key = self._get_key(
|
||||
chat_id,
|
||||
user_id,
|
||||
self.prefix,
|
||||
self.separator,
|
||||
business_connection_id,
|
||||
message_thread_id,
|
||||
bot_id,
|
||||
)
|
||||
if await pipe.exists(_key):
|
||||
await pipe.hset(_key, "data", json.dumps(data))
|
||||
else:
|
||||
return False
|
||||
return True
|
||||
|
||||
def migrate_format(self, bot_id: int, prefix: Optional[str] = "telebot_"):
|
||||
"""
|
||||
Migrate from old to new format of keys.
|
||||
Run this function once to migrate all redis existing keys to new format.
|
||||
|
||||
Starting from version 4.23.0, the format of keys has been changed:
|
||||
<key>:value
|
||||
- Old format: {prefix}chat_id: {user_id: {'state': None, 'data': {}}, ...}
|
||||
- New format:
|
||||
{prefix}{separator}{bot_id}{separator}{business_connection_id}{separator}{message_thread_id}{separator}{chat_id}{separator}{user_id}: {'state': ..., 'data': {}}
|
||||
|
||||
This function will help you to migrate from the old format to the new one in order to avoid data loss.
|
||||
|
||||
:param bot_id: Bot ID; To get it, call a getMe request and grab the id from the response.
|
||||
:type bot_id: int
|
||||
|
||||
:param prefix: Prefix for keys, default is "telebot_"(old default value)
|
||||
:type prefix: Optional[str]
|
||||
"""
|
||||
keys = self.redis.keys(f"{prefix}*")
|
||||
|
||||
for key in keys:
|
||||
old_key = key.decode("utf-8")
|
||||
# old: {prefix}chat_id: {user_id: {'state': None, 'data': {}}, ...}
|
||||
value = self.redis.get(old_key)
|
||||
value = json.loads(value)
|
||||
|
||||
chat_id = old_key[len(prefix) :]
|
||||
user_id = list(value.keys())[0]
|
||||
state = value[user_id]["state"]
|
||||
state_data = value[user_id]["data"]
|
||||
|
||||
# set new format
|
||||
new_key = self._get_key(
|
||||
int(chat_id), int(user_id), self.prefix, self.separator, bot_id=bot_id
|
||||
)
|
||||
self.redis.hset(new_key, "state", state)
|
||||
self.redis.hset(new_key, "data", json.dumps(state_data))
|
||||
|
||||
# delete old key
|
||||
self.redis.delete(old_key)
|
||||
|
||||
def __str__(self) -> str:
|
||||
# include some connection info
|
||||
return f"StateRedisStorage({self.redis})"
|
||||
Reference in New Issue
Block a user