Initial commit
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
dinlo
@@ -0,0 +1,497 @@
|
||||
"""Модуль работы с базой данных SQLite."""
|
||||
|
||||
import aiosqlite
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Optional, List
|
||||
|
||||
from config import settings
|
||||
|
||||
|
||||
class Database:
|
||||
"""Класс для работы с SQLite."""
|
||||
|
||||
def __init__(self, db_path: str):
|
||||
self.db_path = db_path
|
||||
|
||||
async def initialize(self) -> None:
|
||||
"""Инициализация базы данных и создание таблиц."""
|
||||
os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
# Таблица авторизованных пользователей бота
|
||||
await db.execute("""
|
||||
CREATE TABLE IF NOT EXISTS bot_users (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL UNIQUE,
|
||||
username TEXT,
|
||||
display_name TEXT,
|
||||
added_by INTEGER,
|
||||
added_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
access_expires_at TIMESTAMP,
|
||||
max_generations INTEGER,
|
||||
used_generations INTEGER NOT NULL DEFAULT 0,
|
||||
is_active INTEGER NOT NULL DEFAULT 1,
|
||||
is_admin INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
""")
|
||||
|
||||
await db.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_bot_users_user_id
|
||||
ON bot_users(user_id)
|
||||
""")
|
||||
|
||||
# Таблица профилей
|
||||
await db.execute("""
|
||||
CREATE TABLE IF NOT EXISTS profiles (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
is_default INTEGER NOT NULL DEFAULT 0,
|
||||
width INTEGER NOT NULL DEFAULT 512,
|
||||
height INTEGER NOT NULL DEFAULT 512,
|
||||
steps INTEGER NOT NULL DEFAULT 20,
|
||||
cfg_scale REAL NOT NULL DEFAULT 7.0,
|
||||
sampler TEXT NOT NULL DEFAULT 'Euler a',
|
||||
scheduler TEXT NOT NULL DEFAULT 'automatic',
|
||||
model TEXT NOT NULL DEFAULT '',
|
||||
lora TEXT NOT NULL DEFAULT '',
|
||||
lora_strength REAL NOT NULL DEFAULT 0.8,
|
||||
negative_prompt TEXT NOT NULL DEFAULT '',
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
""")
|
||||
|
||||
# Таблица пользовательских настроек
|
||||
await db.execute(f"""
|
||||
CREATE TABLE IF NOT EXISTS user_settings (
|
||||
user_id INTEGER PRIMARY KEY,
|
||||
image_ttl_hours INTEGER NOT NULL DEFAULT {settings.DEFAULT_IMAGE_TTL_HOURS},
|
||||
active_profile_id INTEGER,
|
||||
FOREIGN KEY (active_profile_id) REFERENCES profiles(id)
|
||||
)
|
||||
""")
|
||||
|
||||
# Таблица сгенерированных изображений
|
||||
await db.execute("""
|
||||
CREATE TABLE IF NOT EXISTS generated_images (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
user_id INTEGER NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
prompt TEXT NOT NULL,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
expires_at TIMESTAMP NOT NULL
|
||||
)
|
||||
""")
|
||||
|
||||
await db.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_images_expires
|
||||
ON generated_images(expires_at)
|
||||
""")
|
||||
|
||||
await db.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_profiles_user
|
||||
ON profiles(user_id)
|
||||
""")
|
||||
|
||||
await db.commit()
|
||||
|
||||
async def get_user_settings(self, user_id: int) -> dict:
|
||||
"""Получить настройки пользователя."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT image_ttl_hours, active_profile_id FROM user_settings WHERE user_id = ?",
|
||||
(user_id,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if row:
|
||||
return {
|
||||
"image_ttl_hours": row[0],
|
||||
"active_profile_id": row[1]
|
||||
}
|
||||
return {
|
||||
"image_ttl_hours": settings.DEFAULT_IMAGE_TTL_HOURS,
|
||||
"active_profile_id": None
|
||||
}
|
||||
|
||||
async def update_user_settings(self, user_id: int, **kwargs) -> None:
|
||||
"""Обновить настройки пользователя."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
# Сначала проверяем, есть ли уже запись
|
||||
async with db.execute(
|
||||
"SELECT user_id FROM user_settings WHERE user_id = ?", (user_id,)
|
||||
) as cursor:
|
||||
exists = await cursor.fetchone()
|
||||
|
||||
if exists:
|
||||
set_clause = ", ".join(f"{k} = ?" for k in kwargs.keys())
|
||||
values = list(kwargs.values()) + [user_id]
|
||||
await db.execute(
|
||||
f"UPDATE user_settings SET {set_clause} WHERE user_id = ?",
|
||||
values
|
||||
)
|
||||
else:
|
||||
# Создаём с значениями по умолчанию + новые
|
||||
current_settings = await self.get_user_settings(user_id)
|
||||
current_settings.update(kwargs)
|
||||
await db.execute(
|
||||
"INSERT INTO user_settings (user_id, image_ttl_hours, active_profile_id) VALUES (?, ?, ?)",
|
||||
(user_id, current_settings["image_ttl_hours"], current_settings.get("active_profile_id"))
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def create_profile(
|
||||
self,
|
||||
user_id: int,
|
||||
name: str,
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
steps: int = 20,
|
||||
cfg_scale: float = 7.0,
|
||||
sampler: str = "Euler a",
|
||||
scheduler: str = "automatic",
|
||||
model: str = "",
|
||||
lora: str = "",
|
||||
lora_strength: float = 0.8,
|
||||
negative_prompt: str = "",
|
||||
is_default: bool = False
|
||||
) -> int:
|
||||
"""Создать новый профиль."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
if is_default:
|
||||
# Сбрасываем флаг default у всех профилей пользователя
|
||||
await db.execute(
|
||||
"UPDATE profiles SET is_default = 0 WHERE user_id = ?", (user_id,)
|
||||
)
|
||||
|
||||
cursor = await db.execute("""
|
||||
INSERT INTO profiles (
|
||||
user_id, name, is_default, width, height, steps,
|
||||
cfg_scale, sampler, scheduler, model, lora,
|
||||
lora_strength, negative_prompt
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
user_id, name, int(is_default), width, height, steps,
|
||||
cfg_scale, sampler, scheduler, model, lora,
|
||||
lora_strength, negative_prompt
|
||||
))
|
||||
profile_id = cursor.lastrowid
|
||||
await db.commit()
|
||||
return profile_id
|
||||
|
||||
async def get_profile(self, profile_id: int) -> Optional[dict]:
|
||||
"""Получить профиль по ID."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM profiles WHERE id = ?", (profile_id,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
columns = [description[0] for description in cursor.description]
|
||||
return dict(zip(columns, row))
|
||||
|
||||
async def get_user_profiles(self, user_id: int) -> list[dict]:
|
||||
"""Получить все профили пользователя."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM profiles WHERE user_id = ? ORDER BY is_default DESC, name",
|
||||
(user_id,)
|
||||
) as cursor:
|
||||
rows = await cursor.fetchall()
|
||||
columns = [description[0] for description in cursor.description]
|
||||
return [dict(zip(columns, row)) for row in rows]
|
||||
|
||||
async def get_default_profile(self, user_id: int) -> Optional[dict]:
|
||||
"""Получить профиль по умолчанию."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM profiles WHERE user_id = ? AND is_default = 1",
|
||||
(user_id,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
columns = [description[0] for description in cursor.description]
|
||||
return dict(zip(columns, row))
|
||||
|
||||
async def set_default_profile(self, user_id: int, profile_id: int) -> None:
|
||||
"""Установить профиль по умолчанию."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"UPDATE profiles SET is_default = 0 WHERE user_id = ?", (user_id,)
|
||||
)
|
||||
await db.execute(
|
||||
"UPDATE profiles SET is_default = 1 WHERE id = ? AND user_id = ?",
|
||||
(profile_id, user_id)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def update_profile(self, profile_id: int, user_id: int, **kwargs) -> None:
|
||||
"""Обновить профиль."""
|
||||
if not kwargs:
|
||||
return
|
||||
set_clause = ", ".join(f"{k} = ?" for k in kwargs.keys())
|
||||
values = list(kwargs.values()) + [profile_id, user_id]
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
f"UPDATE profiles SET {set_clause} WHERE id = ? AND user_id = ?",
|
||||
values
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def delete_profile(self, profile_id: int, user_id: int) -> bool:
|
||||
"""Удалить профиль."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute(
|
||||
"DELETE FROM profiles WHERE id = ? AND user_id = ?",
|
||||
(profile_id, user_id)
|
||||
)
|
||||
await db.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def add_generated_image(
|
||||
self,
|
||||
user_id: int,
|
||||
file_path: str,
|
||||
prompt: str,
|
||||
ttl_hours: int
|
||||
) -> int:
|
||||
"""Добавить запись о сгенерированном изображении."""
|
||||
expires_at = datetime.now() + timedelta(hours=ttl_hours)
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute(
|
||||
"INSERT INTO generated_images (user_id, file_path, prompt, expires_at) VALUES (?, ?, ?, ?)",
|
||||
(user_id, file_path, prompt, expires_at.isoformat())
|
||||
)
|
||||
image_id = cursor.lastrowid
|
||||
await db.commit()
|
||||
return image_id
|
||||
|
||||
async def get_expired_images(self) -> list[dict]:
|
||||
"""Получить список просроченных изображений."""
|
||||
now = datetime.now().isoformat()
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT id, file_path FROM generated_images WHERE expires_at < ?",
|
||||
(now,)
|
||||
) as cursor:
|
||||
rows = await cursor.fetchall()
|
||||
return [{"id": row[0], "file_path": row[1]} for row in rows]
|
||||
|
||||
async def delete_image_record(self, image_id: int) -> None:
|
||||
"""Удалить запись об изображении."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"DELETE FROM generated_images WHERE id = ?", (image_id,)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
# --- Методы для работы с авторизованными пользователями ---
|
||||
|
||||
async def ensure_admin(self, admin_user_id: int, username: str = None, display_name: str = None) -> None:
|
||||
"""Убедиться, что админ существует в таблице bot_users. Если нет — создать."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT id FROM bot_users WHERE user_id = ?", (admin_user_id,)
|
||||
) as cursor:
|
||||
exists = await cursor.fetchone()
|
||||
if not exists:
|
||||
await db.execute("""
|
||||
INSERT INTO bot_users (user_id, username, display_name, is_admin, is_active)
|
||||
VALUES (?, ?, ?, 1, 1)
|
||||
""", (admin_user_id, username, display_name))
|
||||
else:
|
||||
await db.execute(
|
||||
"UPDATE bot_users SET is_admin = 1 WHERE user_id = ?", (admin_user_id,)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def add_user(
|
||||
self,
|
||||
user_id: int,
|
||||
added_by: int,
|
||||
username: str = None,
|
||||
display_name: str = None,
|
||||
access_expires_at: Optional[datetime] = None,
|
||||
max_generations: Optional[int] = None,
|
||||
) -> bool:
|
||||
"""Добавить нового пользователя. Возвращает True если успешно."""
|
||||
try:
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
INSERT INTO bot_users (
|
||||
user_id, username, display_name, added_by,
|
||||
access_expires_at, max_generations, used_generations, is_active
|
||||
) VALUES (?, ?, ?, ?, ?, ?, 0, 1)
|
||||
""", (
|
||||
user_id, username, display_name, added_by,
|
||||
access_expires_at.isoformat() if access_expires_at else None,
|
||||
max_generations,
|
||||
))
|
||||
await db.commit()
|
||||
return True
|
||||
except aiosqlite.IntegrityError:
|
||||
# Пользователь уже существует
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE bot_users
|
||||
SET added_by = ?, access_expires_at = ?, max_generations = ?,
|
||||
used_generations = 0, is_active = 1
|
||||
WHERE user_id = ?
|
||||
""", (
|
||||
added_by,
|
||||
access_expires_at.isoformat() if access_expires_at else None,
|
||||
max_generations,
|
||||
user_id,
|
||||
))
|
||||
await db.commit()
|
||||
return True
|
||||
|
||||
async def remove_user(self, user_id: int) -> bool:
|
||||
"""Удалить пользователя (деактивировать)."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"UPDATE bot_users SET is_active = 0 WHERE user_id = ? AND is_admin = 0",
|
||||
(user_id,)
|
||||
)
|
||||
await db.commit()
|
||||
return True
|
||||
|
||||
async def check_user_access(self, user_id: int) -> dict:
|
||||
"""
|
||||
Проверить доступ пользователя.
|
||||
Возвращает dict с полями:
|
||||
- has_access: bool
|
||||
- is_admin: bool
|
||||
- reason: str (причина отказа, если нет доступа)
|
||||
- user_info: dict (информация о пользователе)
|
||||
"""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM bot_users WHERE user_id = ?", (user_id,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if not row:
|
||||
return {
|
||||
"has_access": False,
|
||||
"is_admin": False,
|
||||
"reason": "not_registered",
|
||||
"user_info": None,
|
||||
}
|
||||
|
||||
columns = [description[0] for description in cursor.description]
|
||||
user_info = dict(zip(columns, row))
|
||||
|
||||
# Админ всегда имеет доступ
|
||||
if user_info["is_admin"]:
|
||||
return {
|
||||
"has_access": True,
|
||||
"is_admin": True,
|
||||
"reason": None,
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
# Проверяем активность
|
||||
if not user_info["is_active"]:
|
||||
return {
|
||||
"has_access": False,
|
||||
"is_admin": False,
|
||||
"reason": "deactivated",
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
# Проверяем срок доступа
|
||||
if user_info["access_expires_at"]:
|
||||
expires_at = datetime.fromisoformat(user_info["access_expires_at"])
|
||||
if datetime.now() > expires_at:
|
||||
return {
|
||||
"has_access": False,
|
||||
"is_admin": False,
|
||||
"reason": "expired",
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
# Проверяем лимит генераций
|
||||
if user_info["max_generations"] is not None:
|
||||
if user_info["used_generations"] >= user_info["max_generations"]:
|
||||
return {
|
||||
"has_access": False,
|
||||
"is_admin": False,
|
||||
"reason": "generations_limit",
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
return {
|
||||
"has_access": True,
|
||||
"is_admin": False,
|
||||
"reason": None,
|
||||
"user_info": user_info,
|
||||
}
|
||||
|
||||
async def increment_generation_count(self, user_id: int) -> None:
|
||||
"""Увеличить счётчик использованных генераций."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"UPDATE bot_users SET used_generations = used_generations + 1 WHERE user_id = ?",
|
||||
(user_id,)
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def get_all_users(self) -> List[dict]:
|
||||
"""Получить список всех пользователей."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM bot_users ORDER BY is_admin DESC, added_at DESC"
|
||||
) as cursor:
|
||||
rows = await cursor.fetchall()
|
||||
columns = [description[0] for description in cursor.description]
|
||||
return [dict(zip(columns, row)) for row in rows]
|
||||
|
||||
async def get_user_by_id(self, user_id: int) -> Optional[dict]:
|
||||
"""Получить информацию о пользователе по user_id."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
async with db.execute(
|
||||
"SELECT * FROM bot_users WHERE user_id = ?", (user_id,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
if not row:
|
||||
return None
|
||||
columns = [description[0] for description in cursor.description]
|
||||
return dict(zip(columns, row))
|
||||
|
||||
async def update_user_access(
|
||||
self,
|
||||
user_id: int,
|
||||
access_expires_at: Optional[datetime] = None,
|
||||
max_generations: Optional[int] = None,
|
||||
) -> None:
|
||||
"""Обновить параметры доступа пользователя."""
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute("""
|
||||
UPDATE bot_users
|
||||
SET access_expires_at = ?, max_generations = ?, is_active = 1
|
||||
WHERE user_id = ?
|
||||
""", (
|
||||
access_expires_at.isoformat() if access_expires_at else None,
|
||||
max_generations,
|
||||
user_id,
|
||||
))
|
||||
await db.commit()
|
||||
|
||||
async def deactivate_expired_users(self) -> int:
|
||||
"""Деактивировать пользователей с истёкшим сроком доступа. Возвращает количество."""
|
||||
now = datetime.now().isoformat()
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute("""
|
||||
UPDATE bot_users SET is_active = 0
|
||||
WHERE access_expires_at IS NOT NULL
|
||||
AND access_expires_at < ?
|
||||
AND is_active = 1
|
||||
AND is_admin = 0
|
||||
""", (now,))
|
||||
await db.commit()
|
||||
return cursor.rowcount
|
||||
|
||||
|
||||
# Глобальный экземпляр
|
||||
db = Database(settings.DB_PATH)
|
||||
Reference in New Issue
Block a user