135 lines
4.7 KiB
Python
135 lines
4.7 KiB
Python
import sqlite3
|
||
import random
|
||
from datetime import datetime
|
||
|
||
DB_NAME = "bot.db"
|
||
|
||
|
||
def init_db():
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
# Единая таблица для всех записей
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS records (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
type TEXT NOT NULL,
|
||
phone TEXT,
|
||
qr_code TEXT,
|
||
qr_data TEXT,
|
||
created_at TEXT NOT NULL,
|
||
used_at TEXT,
|
||
usage_count INTEGER DEFAULT 0
|
||
)
|
||
""")
|
||
|
||
# Таблица белого списка
|
||
conn.execute("""
|
||
CREATE TABLE IF NOT EXISTS allowed_users (
|
||
user_id INTEGER PRIMARY KEY,
|
||
username TEXT,
|
||
added_at TEXT NOT NULL,
|
||
added_by TEXT
|
||
)
|
||
""")
|
||
|
||
# Индексы
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_records_used ON records(used_at)")
|
||
conn.execute("CREATE INDEX IF NOT EXISTS idx_records_type ON records(type)")
|
||
|
||
|
||
def add_record(record_type: str, phone: str = None, qr_code: str = None, qr_data: str = None) -> bool:
|
||
"""Добавляет запись"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
try:
|
||
conn.execute("""
|
||
INSERT INTO records (type, phone, qr_code, qr_data, created_at)
|
||
VALUES (?, ?, ?, ?, ?)
|
||
""", (record_type, phone, qr_code, qr_data, datetime.now().isoformat()))
|
||
return True
|
||
except sqlite3.IntegrityError:
|
||
return False
|
||
|
||
|
||
def get_random_record():
|
||
"""Возвращает случайную запись из тех, у кого минимальный usage_count (с ограничением выборки)"""
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
|
||
# Общее количество записей
|
||
total_count = conn.execute("SELECT COUNT(*) as count FROM records").fetchone()["count"]
|
||
|
||
if total_count == 0:
|
||
return None
|
||
|
||
# Определяем лимит для выборки
|
||
if total_count <= 50:
|
||
limit = 10
|
||
else:
|
||
limit = max(10, min(50, int(total_count * 0.2)))
|
||
|
||
# Находим минимальный usage_count
|
||
min_usage = conn.execute("SELECT MIN(usage_count) as min_usage FROM records").fetchone()["min_usage"]
|
||
|
||
# Берём записи с минимальным usage_count, но не больше limit
|
||
candidates = conn.execute("""
|
||
SELECT * FROM records
|
||
WHERE usage_count = ?
|
||
LIMIT ?
|
||
""", (min_usage, limit)).fetchall()
|
||
|
||
if not candidates:
|
||
return None
|
||
|
||
# Из них выбираем случайную
|
||
selected = random.choice(candidates)
|
||
|
||
# Обновляем счётчик и дату использования
|
||
conn.execute("""
|
||
UPDATE records SET usage_count = usage_count + 1, used_at = ?
|
||
WHERE id = ?
|
||
""", (datetime.now().isoformat(), selected["id"]))
|
||
|
||
return dict(selected)
|
||
|
||
|
||
def get_all_records():
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
rows = conn.execute("SELECT * FROM records ORDER BY created_at DESC").fetchall()
|
||
return [dict(row) for row in rows]
|
||
|
||
|
||
def delete_record(record_id: int):
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.execute("DELETE FROM records WHERE id = ?", (record_id,))
|
||
|
||
|
||
# ========== Функции для белого списка ==========
|
||
def add_allowed_user(user_id: int, username: str = None, added_by: str = "admin") -> bool:
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
try:
|
||
conn.execute(
|
||
"INSERT INTO allowed_users (user_id, username, added_at, added_by) VALUES (?, ?, ?, ?)",
|
||
(user_id, username, datetime.now().isoformat(), added_by)
|
||
)
|
||
return True
|
||
except sqlite3.IntegrityError:
|
||
return False
|
||
|
||
|
||
def remove_allowed_user(user_id: int) -> bool:
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
cursor = conn.execute("DELETE FROM allowed_users WHERE user_id = ?", (user_id,))
|
||
return cursor.rowcount > 0
|
||
|
||
|
||
def is_user_allowed(user_id: int) -> bool:
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
cursor = conn.execute("SELECT 1 FROM allowed_users WHERE user_id = ?", (user_id,))
|
||
return cursor.fetchone() is not None
|
||
|
||
|
||
def get_allowed_users():
|
||
with sqlite3.connect(DB_NAME) as conn:
|
||
conn.row_factory = sqlite3.Row
|
||
rows = conn.execute("SELECT * FROM allowed_users ORDER BY added_at DESC").fetchall()
|
||
return [dict(row) for row in rows] |