bot_pozhitki/app/db.py
2026-09-07 21:25:44 +00:00

443 lines
16 KiB
Python

from __future__ import annotations
import sqlite3
from contextlib import closing
from dataclasses import dataclass
from datetime import UTC, datetime
from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
from pathlib import Path
from typing import Any, Iterable
SCHEMA = """
PRAGMA journal_mode = WAL;
PRAGMA foreign_keys = ON;
CREATE TABLE IF NOT EXISTS users (
user_id INTEGER PRIMARY KEY,
first_name TEXT NOT NULL,
last_name TEXT,
nickname TEXT NOT NULL,
username TEXT,
avatar_file_id TEXT,
avatar_path TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS transactions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL REFERENCES users(user_id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('income', 'expense')),
amount_cents INTEGER NOT NULL CHECK (amount_cents > 0),
category TEXT NOT NULL,
note TEXT,
occurred_at TEXT NOT NULL,
created_at TEXT NOT NULL,
cancelled_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_transactions_user_id
ON transactions(user_id);
CREATE TABLE IF NOT EXISTS friendships (
user_id_low INTEGER NOT NULL REFERENCES users(user_id) ON DELETE CASCADE,
user_id_high INTEGER NOT NULL REFERENCES users(user_id) ON DELETE CASCADE,
created_at TEXT NOT NULL,
PRIMARY KEY (user_id_low, user_id_high),
CHECK (user_id_low < user_id_high)
);
CREATE TABLE IF NOT EXISTS user_categories (
user_id INTEGER NOT NULL REFERENCES users(user_id) ON DELETE CASCADE,
kind TEXT NOT NULL CHECK (kind IN ('income', 'expense')),
name TEXT NOT NULL,
created_at TEXT NOT NULL,
PRIMARY KEY (user_id, kind, name)
);
CREATE TABLE IF NOT EXISTS app_meta (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
);
"""
@dataclass(frozen=True, slots=True)
class TransactionInput:
kind: str
amount_cents: int
category: str
note: str | None
occurred_at: str
def utc_now() -> str:
return datetime.now(UTC).isoformat()
def amount_to_cents(value: Any) -> int:
try:
amount = Decimal(str(value)).quantize(Decimal("0.01"), rounding=ROUND_HALF_UP)
except (InvalidOperation, ValueError, TypeError) as exc:
raise ValueError("Некорректная сумма") from exc
cents = int(amount * 100)
if cents <= 0:
raise ValueError("Сумма должна быть больше нуля")
if cents > 100_000_000_000_00:
raise ValueError("Сумма слишком большая")
return cents
class Database:
def __init__(self, path: Path):
self.path = path
def connect(self) -> sqlite3.Connection:
connection = sqlite3.connect(self.path, timeout=15)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA foreign_keys = ON")
connection.execute("PRAGMA busy_timeout = 15000")
return connection
def initialize(self) -> None:
self.path.parent.mkdir(parents=True, exist_ok=True)
with closing(self.connect()) as connection:
connection.executescript(SCHEMA)
columns = {
row["name"]
for row in connection.execute("PRAGMA table_info(transactions)")
}
if "cancelled_at" not in columns:
connection.execute(
"ALTER TABLE transactions ADD COLUMN cancelled_at TEXT"
)
connection.execute(
"""
INSERT OR IGNORE INTO user_categories (user_id, kind, name, created_at)
SELECT user_id, kind, category, MIN(created_at)
FROM transactions
GROUP BY user_id, kind, category
"""
)
connection.commit()
def upsert_user(
self,
*,
user_id: int,
first_name: str,
last_name: str | None,
username: str | None,
avatar_file_id: str | None = None,
avatar_path: str | None = None,
) -> None:
now = utc_now()
nickname = " ".join(part for part in (first_name, last_name) if part).strip()
nickname = nickname or f"Пользователь {user_id}"
with closing(self.connect()) as connection:
connection.execute(
"""
INSERT INTO users (
user_id, first_name, last_name, nickname, username,
avatar_file_id, avatar_path, created_at, updated_at
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(user_id) DO UPDATE SET
first_name = excluded.first_name,
last_name = excluded.last_name,
nickname = excluded.nickname,
username = excluded.username,
avatar_file_id = COALESCE(excluded.avatar_file_id, users.avatar_file_id),
avatar_path = COALESCE(excluded.avatar_path, users.avatar_path),
updated_at = excluded.updated_at
""",
(
user_id,
first_name,
last_name,
nickname,
username,
avatar_file_id,
avatar_path,
now,
now,
),
)
connection.commit()
def get_avatar_file_id(self, user_id: int) -> str | None:
with closing(self.connect()) as connection:
row = connection.execute(
"SELECT avatar_file_id FROM users WHERE user_id = ?", (user_id,)
).fetchone()
return row["avatar_file_id"] if row else None
def add_transaction(self, user_id: int, item: TransactionInput) -> int:
with closing(self.connect()) as connection:
connection.execute(
"""
INSERT OR IGNORE INTO user_categories (user_id, kind, name, created_at)
VALUES (?, ?, ?, ?)
""",
(user_id, item.kind, item.category, utc_now()),
)
cursor = connection.execute(
"""
INSERT INTO transactions (
user_id, kind, amount_cents, category, note,
occurred_at, created_at
) VALUES (?, ?, ?, ?, ?, ?, ?)
""",
(
user_id,
item.kind,
item.amount_cents,
item.category,
item.note,
item.occurred_at,
utc_now(),
),
)
connection.commit()
return int(cursor.lastrowid)
def add_category(self, user_id: int, kind: str, name: str) -> bool:
with closing(self.connect()) as connection:
cursor = connection.execute(
"""
INSERT OR IGNORE INTO user_categories (user_id, kind, name, created_at)
VALUES (?, ?, ?, ?)
""",
(user_id, kind, name, utc_now()),
)
connection.commit()
return cursor.rowcount > 0
def cancel_transaction(self, user_id: int, transaction_id: int) -> bool:
with closing(self.connect()) as connection:
cursor = connection.execute(
"""
UPDATE transactions
SET cancelled_at = ?
WHERE id = ? AND user_id = ? AND cancelled_at IS NULL
""",
(utc_now(), transaction_id, user_id),
)
connection.commit()
return cursor.rowcount > 0
def transactions(self, user_id: int, limit: int = 50) -> list[dict[str, Any]]:
with closing(self.connect()) as connection:
rows = connection.execute(
"""
SELECT id, kind, amount_cents, category, note, occurred_at
FROM transactions
WHERE user_id = ? AND cancelled_at IS NULL
ORDER BY occurred_at DESC, id DESC
LIMIT ?
""",
(user_id, limit),
).fetchall()
return [dict(row) for row in rows]
def user_categories(self, user_id: int) -> list[dict[str, str]]:
with closing(self.connect()) as connection:
rows = connection.execute(
"""
SELECT kind, name
FROM user_categories
WHERE user_id = ?
ORDER BY kind, name COLLATE NOCASE
""",
(user_id,),
).fetchall()
return [dict(row) for row in rows]
def global_categories(self) -> list[dict[str, str]]:
with closing(self.connect()) as connection:
rows = connection.execute(
"""
SELECT kind, name
FROM user_categories
GROUP BY kind, name
ORDER BY name COLLATE NOCASE, kind
"""
).fetchall()
return [dict(row) for row in rows]
def add_friendship(self, first_user_id: int, second_user_id: int) -> bool:
if first_user_id == second_user_id:
return False
low, high = sorted((first_user_id, second_user_id))
with closing(self.connect()) as connection:
cursor = connection.execute(
"""
INSERT OR IGNORE INTO friendships (user_id_low, user_id_high, created_at)
SELECT ?, ?, ?
WHERE EXISTS (SELECT 1 FROM users WHERE user_id = ?)
AND EXISTS (SELECT 1 FROM users WHERE user_id = ?)
""",
(low, high, utc_now(), low, high),
)
connection.commit()
return cursor.rowcount > 0
def friend_ids(self, user_id: int) -> list[int]:
with closing(self.connect()) as connection:
rows = connection.execute(
"""
SELECT CASE
WHEN user_id_low = ? THEN user_id_high
ELSE user_id_low
END AS friend_id
FROM friendships
WHERE user_id_low = ? OR user_id_high = ?
ORDER BY created_at DESC
""",
(user_id, user_id, user_id),
).fetchall()
return [int(row["friend_id"]) for row in rows]
def remove_friendship(self, first_user_id: int, second_user_id: int) -> bool:
if first_user_id == second_user_id:
return False
low, high = sorted((first_user_id, second_user_id))
with closing(self.connect()) as connection:
cursor = connection.execute(
"DELETE FROM friendships WHERE user_id_low = ? AND user_id_high = ?",
(low, high),
)
connection.commit()
return cursor.rowcount > 0
def set_meta(self, key: str, value: str) -> None:
with closing(self.connect()) as connection:
connection.execute(
"""
INSERT INTO app_meta (key, value) VALUES (?, ?)
ON CONFLICT(key) DO UPDATE SET value = excluded.value
""",
(key, value),
)
connection.commit()
def get_meta(self, key: str) -> str | None:
with closing(self.connect()) as connection:
row = connection.execute(
"SELECT value FROM app_meta WHERE key = ?", (key,)
).fetchone()
return row["value"] if row else None
def summary(self, user_id: int) -> dict[str, Any] | None:
with closing(self.connect()) as connection:
row = connection.execute(
"""
SELECT
u.user_id,
u.nickname,
u.username,
u.avatar_path,
COALESCE(SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END), 0) AS income_cents,
COALESCE(SUM(CASE WHEN t.kind = 'expense' THEN t.amount_cents ELSE 0 END), 0) AS expense_cents
FROM users u
LEFT JOIN transactions t
ON t.user_id = u.user_id AND t.cancelled_at IS NULL
WHERE u.user_id = ?
GROUP BY u.user_id
""",
(user_id,),
).fetchone()
if not row:
return None
categories = connection.execute(
"""
SELECT kind, category, SUM(amount_cents) AS amount_cents
FROM transactions
WHERE user_id = ? AND cancelled_at IS NULL
GROUP BY kind, category
ORDER BY amount_cents DESC
""",
(user_id,),
).fetchall()
return self._serialize_totals(dict(row), categories)
def leaderboard(
self,
*,
user_ids: Iterable[int] | None = None,
categories: Iterable[str] | None = None,
limit: int = 100,
) -> list[dict[str, Any]]:
ids = list(dict.fromkeys(user_ids or []))
category_names = list(dict.fromkeys(categories or []))
where = ""
where_params: list[Any] = []
if ids:
placeholders = ",".join("?" for _ in ids)
where = f"WHERE u.user_id IN ({placeholders})"
where_params.extend(ids)
join_conditions = ["t.user_id = u.user_id", "t.cancelled_at IS NULL"]
join_params: list[Any] = []
if category_names:
placeholders = ",".join("?" for _ in category_names)
join_conditions.append(f"t.category IN ({placeholders})")
join_params.extend(category_names)
params = [*join_params, *where_params, limit]
with closing(self.connect()) as connection:
rows = connection.execute(
f"""
SELECT
u.user_id,
u.nickname,
u.username,
u.avatar_path,
COALESCE(SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END), 0) AS income_cents,
COALESCE(SUM(CASE WHEN t.kind = 'expense' THEN t.amount_cents ELSE 0 END), 0) AS expense_cents
FROM users u
LEFT JOIN transactions t ON {' AND '.join(join_conditions)}
{where}
GROUP BY u.user_id
ORDER BY
CASE
WHEN SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END) > 0
THEN 1.0 * (
SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END) -
SUM(CASE WHEN t.kind = 'expense' THEN t.amount_cents ELSE 0 END)
) / SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END)
ELSE 0
END DESC,
(SUM(CASE WHEN t.kind = 'income' THEN t.amount_cents ELSE 0 END) -
SUM(CASE WHEN t.kind = 'expense' THEN t.amount_cents ELSE 0 END)) DESC,
u.user_id ASC
LIMIT ?
""",
params,
).fetchall()
result = []
for rank, row in enumerate(rows, start=1):
item = self._serialize_totals(dict(row), [])
item["rank"] = rank
result.append(item)
return result
@staticmethod
def _serialize_totals(
row: dict[str, Any], categories: Iterable[sqlite3.Row]
) -> dict[str, Any]:
income = int(row["income_cents"])
expense = int(row["expense_cents"])
balance = income - expense
percent = round(balance * 100 / income, 1) if income else 0.0
return {
"user_id": int(row["user_id"]),
"nickname": row["nickname"],
"username": row["username"],
"avatar_url": (
f"/api/avatars/{row['user_id']}" if row.get("avatar_path") else None
),
"income_cents": income,
"expense_cents": expense,
"balance_cents": balance,
"saved_percent": percent,
"categories": [dict(category) for category in categories],
}