223 lines
7.7 KiB
Python
223 lines
7.7 KiB
Python
# database.py
|
||
from typing import Any, Coroutine
|
||
|
||
import aiosqlite
|
||
from datetime import timedelta
|
||
import pytz
|
||
|
||
DB_NAME = "bot_database.db"
|
||
|
||
|
||
async def init_db():
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute("""
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
user_id INTEGER PRIMARY KEY,
|
||
username TEXT,
|
||
full_name TEXT,
|
||
is_allowed BOOLEAN DEFAULT 0,
|
||
last_notification_time TIMESTAMP DEFAULT NULL
|
||
)
|
||
""")
|
||
|
||
# Добавляем колонку если её нет
|
||
try:
|
||
await db.execute("ALTER TABLE users ADD COLUMN last_notification_time TIMESTAMP DEFAULT NULL")
|
||
except aiosqlite.OperationalError:
|
||
pass
|
||
|
||
await db.execute("""
|
||
CREATE TABLE IF NOT EXISTS settings (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT
|
||
)
|
||
""")
|
||
|
||
await db.execute(
|
||
"INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)",
|
||
("target_email", "")
|
||
)
|
||
await db.execute(
|
||
"INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)",
|
||
("admin_id", "")
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def add_user(user_id: int, username: str, full_name: str):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"INSERT OR IGNORE INTO users (user_id, username, full_name, is_allowed, last_notification_time) VALUES (?, ?, ?, 0, NULL)",
|
||
(user_id, username, full_name)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def set_user_allowed(user_id: int, allowed: bool):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"UPDATE users SET is_allowed = ? WHERE user_id = ?",
|
||
(1 if allowed else 0, user_id)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def is_user_allowed(user_id: int) -> bool:
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute("SELECT is_allowed FROM users WHERE user_id = ?", (user_id,)) as cursor:
|
||
row = await cursor.fetchone()
|
||
return row is not None and row[0] == 1
|
||
|
||
|
||
async def get_target_email() -> str:
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute("SELECT value FROM settings WHERE key = 'target_email'") as cursor:
|
||
row = await cursor.fetchone()
|
||
return row[0] if row else ""
|
||
|
||
|
||
async def set_target_email(email: str):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"UPDATE settings SET value = ? WHERE key = 'target_email'",
|
||
(email,)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def get_admin_id() -> int:
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute("SELECT value FROM settings WHERE key = 'admin_id'") as cursor:
|
||
row = await cursor.fetchone()
|
||
if row and row[0]:
|
||
return int(row[0])
|
||
return None
|
||
|
||
|
||
async def set_admin_id(admin_id: int):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"UPDATE settings SET value = ? WHERE key = 'admin_id'",
|
||
(str(admin_id),)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def get_all_users():
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute("SELECT user_id, username, full_name, is_allowed FROM users ORDER BY user_id") as cursor:
|
||
return await cursor.fetchall()
|
||
|
||
|
||
async def delete_user(user_id: int):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute("DELETE FROM users WHERE user_id = ?", (user_id,))
|
||
await db.commit()
|
||
return True
|
||
|
||
|
||
async def get_user_by_id(user_id: int):
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute("SELECT user_id, username, full_name, is_allowed FROM users WHERE user_id = ?",
|
||
(user_id,)) as cursor:
|
||
return await cursor.fetchone()
|
||
|
||
|
||
async def can_send_notification(user_id: int) -> bool:
|
||
"""Проверяет, можно ли отправить уведомление (с учётом часового пояса Екатеринбурга)"""
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute(
|
||
"SELECT is_allowed, last_notification_time FROM users WHERE user_id = ?",
|
||
(user_id,)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
|
||
if not row:
|
||
return True
|
||
|
||
is_allowed, last_time = row
|
||
|
||
# Если пользователь уже имеет доступ, не отправляем уведомления
|
||
if is_allowed == 1:
|
||
return False
|
||
|
||
# Если нет доступа и не было уведомлений
|
||
if last_time is None:
|
||
return True
|
||
|
||
from timezone_utils import parse_time_from_utc, get_current_time
|
||
|
||
# Преобразуем время из БД в Екатеринбург
|
||
last_notification = parse_time_from_utc(last_time)
|
||
|
||
if last_notification is None:
|
||
return True
|
||
|
||
# Проверяем, прошёл ли час
|
||
now = get_current_time()
|
||
return (now - last_notification) >= timedelta(hours=1)
|
||
|
||
|
||
async def update_notification_time(user_id: int):
|
||
"""Обновляет время последнего уведомления (храним в UTC)"""
|
||
from timezone_utils import get_current_time, parse_time_to_utc
|
||
|
||
local_time = get_current_time()
|
||
utc_time = parse_time_to_utc(local_time)
|
||
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"UPDATE users SET last_notification_time = ? WHERE user_id = ?",
|
||
(utc_time.isoformat(), user_id)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def reset_notification_timer(user_id: int):
|
||
"""Сбрасывает таймер уведомлений"""
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
await db.execute(
|
||
"UPDATE users SET last_notification_time = NULL WHERE user_id = ?",
|
||
(user_id,)
|
||
)
|
||
await db.commit()
|
||
|
||
|
||
async def get_last_notification_time(user_id: int) -> Any | None:
|
||
"""Получить время последнего уведомления в формате Екатеринбурга"""
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute(
|
||
"SELECT last_notification_time FROM users WHERE user_id = ?",
|
||
(user_id,)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
if row and row[0]:
|
||
from timezone_utils import parse_time_from_utc, format_time
|
||
local_time = parse_time_from_utc(row[0])
|
||
if local_time:
|
||
return format_time(local_time)
|
||
return None
|
||
|
||
|
||
async def get_next_allowed_time(user_id: int) -> Any | None:
|
||
"""
|
||
Возвращает время (в Екатеринбурге), когда пользователь сможет отправить следующий запрос
|
||
"""
|
||
async with aiosqlite.connect(DB_NAME) as db:
|
||
async with db.execute(
|
||
"SELECT last_notification_time FROM users WHERE user_id = ?",
|
||
(user_id,)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
|
||
if not row or row[0] is None:
|
||
return None
|
||
|
||
from timezone_utils import parse_time_from_utc, format_time
|
||
from datetime import timedelta
|
||
|
||
last_time = parse_time_from_utc(row[0])
|
||
if last_time:
|
||
next_time = last_time + timedelta(hours=1)
|
||
return format_time(next_time)
|
||
return None |