generated from FOSS/Bot-Template
Initial commit
This commit is contained in:
Lolzteam
@@ -0,0 +1,5 @@
|
||||
from .db_settings import SettingsModel, SettingsRepository
|
||||
from .db_users import UserModel, UsersRepository
|
||||
|
||||
ModelSettings = SettingsModel
|
||||
ModelUsers = UserModel
|
||||
@@ -0,0 +1,51 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import event
|
||||
from sqlalchemy.ext.asyncio import AsyncAttrs, AsyncSession, async_sessionmaker, create_async_engine
|
||||
from sqlalchemy.orm import DeclarativeBase
|
||||
|
||||
from tgbot.data.config import PATH_DATABASE
|
||||
|
||||
database_path = Path(PATH_DATABASE)
|
||||
database_url = f"sqlite+aiosqlite:///{database_path.as_posix()}"
|
||||
|
||||
engine = create_async_engine(database_url, echo=False)
|
||||
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
||||
|
||||
|
||||
# Общая база для всех SQLAlchemy-моделей
|
||||
class Base(AsyncAttrs, DeclarativeBase):
|
||||
pass
|
||||
|
||||
|
||||
# SQLite по умолчанию не включает foreign keys, поэтому включаем явно
|
||||
@event.listens_for(engine.sync_engine, "connect")
|
||||
def _enable_sqlite_foreign_keys(dbapi_connection, connection_record) -> None:
|
||||
cursor = dbapi_connection.cursor()
|
||||
cursor.execute("PRAGMA foreign_keys=ON")
|
||||
cursor.close()
|
||||
|
||||
|
||||
# Открываем сессию и сами отвечаем за сохранение или откат
|
||||
@asynccontextmanager
|
||||
async def _session_scope() -> AsyncIterator[AsyncSession]:
|
||||
async with session_factory() as session:
|
||||
try:
|
||||
yield session
|
||||
await session.commit()
|
||||
except Exception:
|
||||
await session.rollback()
|
||||
raise
|
||||
|
||||
|
||||
# Отдельная обертка, чтобы IDE нормально видела асинхронный контекст
|
||||
def session_scope() -> AbstractAsyncContextManager[AsyncSession]:
|
||||
return _session_scope()
|
||||
|
||||
|
||||
# Закрываем пул соединений при остановке приложения
|
||||
async def close_database() -> None:
|
||||
await engine.dispose()
|
||||
@@ -0,0 +1,61 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from sqlalchemy import Boolean, Integer
|
||||
from sqlalchemy import update as sqlalchemy_update
|
||||
from sqlalchemy.dialects.sqlite import insert
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from tgbot.database.core import Base, session_scope
|
||||
from tgbot.database.repository import BaseRepository
|
||||
|
||||
|
||||
# Модель настроек бота
|
||||
class SettingsModel(Base):
|
||||
__tablename__ = "storage_settings"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, default=1)
|
||||
status_work: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
|
||||
|
||||
|
||||
ModelBase = SettingsModel
|
||||
BaseModel = SettingsModel
|
||||
|
||||
|
||||
# Репозиторий настроек бота
|
||||
class SettingsRepository(BaseRepository[SettingsModel]):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.table_model = SettingsModel
|
||||
self.storage_name = SettingsModel.__tablename__
|
||||
|
||||
# Создание строки настроек если их еще нет
|
||||
async def ensure_default(self) -> None:
|
||||
statement = insert(SettingsModel).values(id=1, status_work=False)
|
||||
statement = statement.on_conflict_do_nothing(index_elements=[SettingsModel.id])
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(statement)
|
||||
|
||||
# Настройки должны быть всегда, поэтому при пустой таблице создаем дефолт
|
||||
async def get(self) -> SettingsModel:
|
||||
settings = await super().get(id=1)
|
||||
|
||||
if settings is None:
|
||||
await self.ensure_default()
|
||||
settings = await super().get(id=1)
|
||||
|
||||
if settings is None:
|
||||
raise RuntimeError("Настройки бота по умолчанию не сохранились")
|
||||
|
||||
return settings
|
||||
|
||||
# Обновление единственной строки настроек
|
||||
async def update(self, **kwargs) -> None:
|
||||
if not kwargs:
|
||||
return
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(
|
||||
sqlalchemy_update(SettingsModel)
|
||||
.where(SettingsModel.id == 1)
|
||||
.values(**kwargs)
|
||||
)
|
||||
@@ -0,0 +1,106 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from sqlalchemy import BigInteger, Integer, String, or_
|
||||
from sqlalchemy import update as sqlalchemy_update
|
||||
from sqlalchemy.dialects.sqlite import insert
|
||||
from sqlalchemy.orm import Mapped, mapped_column
|
||||
|
||||
from tgbot.database.core import Base, session_scope
|
||||
from tgbot.database.repository import BaseRepository
|
||||
from tgbot.utils.const_functions import get_unix
|
||||
|
||||
|
||||
# Модель пользователя Telegram
|
||||
class UserModel(Base):
|
||||
__tablename__ = "storage_users"
|
||||
|
||||
increment: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id: Mapped[int] = mapped_column(BigInteger, nullable=False, unique=True, index=True)
|
||||
user_login: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
user_name: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
user_surname: Mapped[str] = mapped_column(String(255), nullable=False, default="")
|
||||
user_fullname: Mapped[str] = mapped_column(String(511), nullable=False, default="")
|
||||
user_unix: Mapped[int] = mapped_column(Integer, nullable=False, default=get_unix)
|
||||
|
||||
|
||||
ModelBase = UserModel
|
||||
BaseModel = UserModel
|
||||
|
||||
|
||||
# Репозиторий пользователей
|
||||
class UsersRepository(BaseRepository[UserModel]):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.table_model = UserModel
|
||||
self.storage_name = UserModel.__tablename__
|
||||
|
||||
# Для совместимости add ведет себя как upsert
|
||||
async def add(
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str,
|
||||
user_fullname: str,
|
||||
) -> UserModel:
|
||||
return await self.upsert(
|
||||
user_id=user_id,
|
||||
user_login=user_login,
|
||||
user_name=user_name,
|
||||
user_surname=user_surname,
|
||||
user_fullname=user_fullname,
|
||||
)
|
||||
|
||||
# Создание пользователя или обновление его данных по user_id
|
||||
async def upsert(
|
||||
self,
|
||||
user_id: int,
|
||||
user_login: str,
|
||||
user_name: str,
|
||||
user_surname: str,
|
||||
user_fullname: str,
|
||||
) -> UserModel:
|
||||
statement = insert(UserModel).values(
|
||||
user_id=user_id,
|
||||
user_login=user_login,
|
||||
user_name=user_name,
|
||||
user_surname=user_surname,
|
||||
user_fullname=user_fullname,
|
||||
user_unix=get_unix(),
|
||||
)
|
||||
statement = statement.on_conflict_do_update(
|
||||
index_elements=[UserModel.user_id],
|
||||
set_={
|
||||
"user_login": statement.excluded.user_login,
|
||||
"user_name": statement.excluded.user_name,
|
||||
"user_surname": statement.excluded.user_surname,
|
||||
"user_fullname": statement.excluded.user_fullname,
|
||||
},
|
||||
where=or_(
|
||||
UserModel.user_login != statement.excluded.user_login,
|
||||
UserModel.user_name != statement.excluded.user_name,
|
||||
UserModel.user_surname != statement.excluded.user_surname,
|
||||
UserModel.user_fullname != statement.excluded.user_fullname,
|
||||
),
|
||||
)
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(statement)
|
||||
|
||||
user = await self.get(user_id=user_id)
|
||||
|
||||
if user is None:
|
||||
raise RuntimeError("Пользователь не сохранился")
|
||||
|
||||
return user
|
||||
|
||||
# Обновление пользователя по Telegram ID
|
||||
async def update(self, user_id: int, **kwargs) -> None:
|
||||
if not kwargs:
|
||||
return
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(
|
||||
sqlalchemy_update(UserModel)
|
||||
.where(UserModel.user_id == user_id)
|
||||
.values(**kwargs)
|
||||
)
|
||||
@@ -0,0 +1,59 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
import asyncio
|
||||
from contextlib import AbstractAsyncContextManager
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
|
||||
|
||||
from tgbot.database.core import database_url
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
||||
ALEMBIC_INI = PROJECT_ROOT / "alembic.ini"
|
||||
MIGRATIONS_DIR = PROJECT_ROOT / "migrations"
|
||||
|
||||
|
||||
# Собираем конфиг Alembic так, чтобы команды работали из любой папки
|
||||
def get_alembic_config(url: str = database_url) -> Config:
|
||||
config = Config(str(ALEMBIC_INI))
|
||||
config.set_main_option("script_location", str(MIGRATIONS_DIR))
|
||||
config.set_main_option("sqlalchemy.url", url)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
# Применяем все миграции до последней версии
|
||||
async def run_migrations(engine: Optional[AsyncEngine] = None) -> None:
|
||||
if engine is None:
|
||||
config = get_alembic_config()
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, command.upgrade, config, "head")
|
||||
return
|
||||
|
||||
config = get_alembic_config(str(engine.url))
|
||||
|
||||
async with engine_context(engine) as connection:
|
||||
await connection.run_sync(_upgrade_with_connection, config)
|
||||
|
||||
|
||||
# Отдельная обертка нужна, чтобы IDE корректно видела async context manager
|
||||
def engine_context(engine: AsyncEngine) -> AbstractAsyncContextManager[AsyncConnection]:
|
||||
return engine.begin()
|
||||
|
||||
|
||||
# Alembic умеет работать с синхронным соединением внутри async-engine
|
||||
def _upgrade_with_connection(connection, config: Config) -> None:
|
||||
config.attributes["connection"] = connection
|
||||
command.upgrade(config, "head")
|
||||
|
||||
|
||||
# Ручной запуск миграций из консоли
|
||||
async def _main() -> None:
|
||||
await run_migrations()
|
||||
print("Миграции базы данных применены")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(_main())
|
||||
@@ -0,0 +1,106 @@
|
||||
# - *- coding: utf- 8 - *-
|
||||
from typing import Generic, List, Optional, Type, TypeVar
|
||||
|
||||
from sqlalchemy import delete as sqlalchemy_delete
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import update as sqlalchemy_update
|
||||
from sqlalchemy.exc import SQLAlchemyError
|
||||
|
||||
from tgbot.database.core import Base, database_path, session_scope
|
||||
from tgbot.utils.misc.bot_logging import bot_logger
|
||||
|
||||
ModelTranslator = TypeVar("ModelTranslator", bound=Base)
|
||||
|
||||
|
||||
# Базовый репозиторий с общими методами работы с БД
|
||||
class BaseRepository(Generic[ModelTranslator]):
|
||||
def __init__(self):
|
||||
self.storage_name = "storage"
|
||||
self.table_model: Optional[Type[ModelTranslator]] = None
|
||||
|
||||
# Без модели репозиторий работать не должен
|
||||
def _model(self) -> Type[ModelTranslator]:
|
||||
if self.table_model is None:
|
||||
raise RuntimeError("Модель базы данных не настроена")
|
||||
|
||||
return self.table_model
|
||||
|
||||
# Удаление только по явному фильтру, без случайной чистки всей таблицы
|
||||
async def delete(self, **kwargs) -> None:
|
||||
if not kwargs:
|
||||
raise ValueError("Для удаления нужен хотя бы один фильтр")
|
||||
|
||||
model = self._model()
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(sqlalchemy_delete(model).filter_by(**kwargs))
|
||||
|
||||
# Полное удаление всех строк таблицы. Название намеренно прямое.
|
||||
async def delete_all_rows(self) -> None:
|
||||
model = self._model()
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(sqlalchemy_delete(model))
|
||||
|
||||
# Возвращение первой записи по фильтру
|
||||
async def get(self, **kwargs) -> Optional[ModelTranslator]:
|
||||
model = self._model()
|
||||
statement = select(model).filter_by(**kwargs)
|
||||
|
||||
async with session_scope() as session:
|
||||
response = await session.execute(statement)
|
||||
|
||||
return response.scalars().first()
|
||||
|
||||
# Возвращение всех записей по фильтру
|
||||
async def gets(self, **kwargs) -> List[ModelTranslator]:
|
||||
model = self._model()
|
||||
statement = select(model).filter_by(**kwargs)
|
||||
|
||||
async with session_scope() as session:
|
||||
response = await session.execute(statement)
|
||||
|
||||
return list(response.scalars().all())
|
||||
|
||||
# Возвращение всей таблицы
|
||||
async def get_all(self) -> List[ModelTranslator]:
|
||||
model = self._model()
|
||||
|
||||
async with session_scope() as session:
|
||||
response = await session.execute(select(model))
|
||||
|
||||
return list(response.scalars().all())
|
||||
|
||||
# Обновляем только записи, которые попали под фильтр
|
||||
async def update(self, filters: dict, **kwargs) -> None:
|
||||
if not filters:
|
||||
raise ValueError("Для обновления нужен хотя бы один фильтр")
|
||||
|
||||
if not kwargs:
|
||||
return
|
||||
|
||||
model = self._model()
|
||||
|
||||
async with session_scope() as session:
|
||||
await session.execute(
|
||||
sqlalchemy_update(model)
|
||||
.filter_by(**filters)
|
||||
.values(**kwargs)
|
||||
)
|
||||
|
||||
|
||||
# Готовим подключение к БД. Сами миграции запускаются отдельно через migrate.py.
|
||||
async def prepare_database() -> None:
|
||||
database_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Импортируем модели, чтобы репозитории работали с уже загруженными таблицами
|
||||
import tgbot.database # noqa: F401
|
||||
|
||||
from tgbot.database.db_settings import SettingsRepository
|
||||
|
||||
try:
|
||||
await SettingsRepository().ensure_default()
|
||||
except SQLAlchemyError as error:
|
||||
raise RuntimeError("База данных не готова. Запусти миграции командой: python migrate.py") from error
|
||||
|
||||
bot_logger.info("База данных готова")
|
||||
Reference in New Issue
Block a user