109 lines
4.1 KiB
Python
109 lines
4.1 KiB
Python
import aiosqlite
|
|
from datetime import datetime
|
|
from typing import Optional, List, Dict, Any
|
|
|
|
DATABASE_PATH = "bot_database.db"
|
|
|
|
|
|
async def init_db():
|
|
"""Инициализация базы данных и создание таблиц"""
|
|
async with aiosqlite.connect(DATABASE_PATH) as db:
|
|
await db.execute("PRAGMA journal_mode=WAL;")
|
|
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
username TEXT,
|
|
first_name TEXT,
|
|
last_name TEXT,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
)
|
|
""")
|
|
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS images (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
original_file_id TEXT NOT NULL,
|
|
original_url TEXT,
|
|
processed_file_id TEXT,
|
|
processed_url TEXT,
|
|
prompt TEXT,
|
|
status TEXT DEFAULT 'processing',
|
|
created_at TEXT NOT NULL,
|
|
completed_at TEXT,
|
|
FOREIGN KEY(user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
|
)
|
|
""")
|
|
|
|
await db.commit()
|
|
|
|
|
|
async def save_user(user_id: int, username: str = None, first_name: str = None, last_name: str = None):
|
|
"""Сохранение или обновление информации о пользователе"""
|
|
now = datetime.now().isoformat()
|
|
async with aiosqlite.connect(DATABASE_PATH) as db:
|
|
await db.execute("""
|
|
INSERT INTO users (user_id, username, first_name, last_name, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(user_id) DO UPDATE SET
|
|
username = excluded.username,
|
|
first_name = excluded.first_name,
|
|
last_name = excluded.last_name,
|
|
updated_at = excluded.updated_at
|
|
""", (user_id, username, first_name, last_name, now, now))
|
|
await db.commit()
|
|
|
|
|
|
async def save_image_record(
|
|
user_id: int,
|
|
original_file_id: str,
|
|
prompt: str = None,
|
|
original_url: str = None
|
|
) -> int:
|
|
"""Сохранение записи об изображении"""
|
|
now = datetime.now().isoformat()
|
|
async with aiosqlite.connect(DATABASE_PATH) as db:
|
|
cursor = await db.execute("""
|
|
INSERT INTO images (user_id, original_file_id, original_url, prompt, created_at, status)
|
|
VALUES (?, ?, ?, ?, ?, 'processing')
|
|
""", (user_id, original_file_id, original_url, prompt, now))
|
|
await db.commit()
|
|
return cursor.lastrowid
|
|
|
|
|
|
async def update_image_record(
|
|
image_id: int,
|
|
processed_file_id: str = None,
|
|
processed_url: str = None,
|
|
status: str = 'completed'
|
|
):
|
|
"""Обновление записи об изображении после обработки"""
|
|
now = datetime.now().isoformat()
|
|
async with aiosqlite.connect(DATABASE_PATH) as db:
|
|
await db.execute("""
|
|
UPDATE images
|
|
SET processed_file_id = COALESCE(?, processed_file_id),
|
|
processed_url = COALESCE(?, processed_url),
|
|
status = ?,
|
|
completed_at = ?
|
|
WHERE id = ?
|
|
""", (processed_file_id, processed_url, status, now, image_id))
|
|
await db.commit()
|
|
|
|
|
|
async def get_user_images(user_id: int, limit: int = 10) -> List[Dict[str, Any]]:
|
|
"""Получение истории обработанных изображений пользователя"""
|
|
async with aiosqlite.connect(DATABASE_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
cursor = await db.execute("""
|
|
SELECT id, original_file_id, processed_file_id, prompt, status, created_at, completed_at
|
|
FROM images
|
|
WHERE user_id = ?
|
|
ORDER BY created_at DESC
|
|
LIMIT ?
|
|
""", (user_id, limit))
|
|
rows = await cursor.fetchall()
|
|
await db.commit()
|
|
return [dict(row) for row in rows] |