401 lines
12 KiB
Python
401 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import date
|
|
from typing import Any
|
|
|
|
from bot.db.database import Database
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Budget:
|
|
id: int
|
|
user_id: int
|
|
name: str
|
|
total_amount: float
|
|
start_date: date
|
|
end_date: date
|
|
currency: str
|
|
is_active: bool
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Expense:
|
|
id: int
|
|
user_id: int
|
|
budget_id: int
|
|
amount: float
|
|
note: str | None
|
|
spent_at: date
|
|
|
|
|
|
def _parse_date(value: str) -> date:
|
|
return date.fromisoformat(value[:10])
|
|
|
|
|
|
def _row_to_budget(row: Any) -> Budget:
|
|
return Budget(
|
|
id=int(row["id"]),
|
|
user_id=int(row["user_id"]),
|
|
name=str(row["name"] or "Бюджет"),
|
|
total_amount=float(row["total_amount"]),
|
|
start_date=_parse_date(row["start_date"]),
|
|
end_date=_parse_date(row["end_date"]),
|
|
currency=str(row["currency"] or "RUB"),
|
|
is_active=bool(row["is_active"]),
|
|
)
|
|
|
|
|
|
def _row_to_expense(row: Any) -> Expense:
|
|
return Expense(
|
|
id=int(row["id"]),
|
|
user_id=int(row["user_id"]),
|
|
budget_id=int(row["budget_id"]),
|
|
amount=float(row["amount"]),
|
|
note=row["note"],
|
|
spent_at=_parse_date(row["spent_at"]),
|
|
)
|
|
|
|
|
|
class BudgetRepository:
|
|
def __init__(self, db: Database) -> None:
|
|
self._db = db
|
|
|
|
async def ensure_user(self, user_id: int) -> None:
|
|
await self._db.conn.execute(
|
|
"INSERT OR IGNORE INTO users(user_id) VALUES (?)",
|
|
(user_id,),
|
|
)
|
|
await self._db.conn.commit()
|
|
|
|
async def get_selected_budget_id(self, user_id: int) -> int | None:
|
|
cursor = await self._db.conn.execute(
|
|
"SELECT selected_budget_id FROM users WHERE user_id = ?",
|
|
(user_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
if row is None or row["selected_budget_id"] is None:
|
|
return None
|
|
return int(row["selected_budget_id"])
|
|
|
|
async def set_selected_budget_id(self, user_id: int, budget_id: int | None) -> None:
|
|
await self.ensure_user(user_id)
|
|
await self._db.conn.execute(
|
|
"UPDATE users SET selected_budget_id = ? WHERE user_id = ?",
|
|
(budget_id, user_id),
|
|
)
|
|
await self._db.conn.commit()
|
|
|
|
async def create_budget(
|
|
self,
|
|
user_id: int,
|
|
*,
|
|
name: str,
|
|
total_amount: float,
|
|
start_date: date,
|
|
end_date: date,
|
|
currency: str = "RUB",
|
|
is_active: bool = True,
|
|
select: bool = True,
|
|
) -> Budget:
|
|
await self.ensure_user(user_id)
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
INSERT INTO budgets(
|
|
user_id, name, total_amount, start_date, end_date, currency, is_active
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
RETURNING *
|
|
""",
|
|
(
|
|
user_id,
|
|
name.strip() or "Бюджет",
|
|
total_amount,
|
|
start_date.isoformat(),
|
|
end_date.isoformat(),
|
|
currency,
|
|
1 if is_active else 0,
|
|
),
|
|
)
|
|
row = await cursor.fetchone()
|
|
await self._db.conn.commit()
|
|
if row is None:
|
|
raise RuntimeError("Failed to create budget")
|
|
budget = _row_to_budget(row)
|
|
if select:
|
|
await self.set_selected_budget_id(user_id, budget.id)
|
|
return budget
|
|
|
|
async def update_budget(
|
|
self,
|
|
budget_id: int,
|
|
user_id: int,
|
|
*,
|
|
name: str | None = None,
|
|
total_amount: float | None = None,
|
|
start_date: date | None = None,
|
|
end_date: date | None = None,
|
|
currency: str | None = None,
|
|
) -> Budget:
|
|
budget = await self.get_budget_for_user(budget_id, user_id)
|
|
if budget is None:
|
|
raise ValueError("Бюджет не найден")
|
|
|
|
next_name = name.strip() if name is not None else budget.name
|
|
next_total = total_amount if total_amount is not None else budget.total_amount
|
|
next_start = start_date if start_date is not None else budget.start_date
|
|
next_end = end_date if end_date is not None else budget.end_date
|
|
next_currency = currency if currency is not None else budget.currency
|
|
|
|
await self._db.conn.execute(
|
|
"""
|
|
UPDATE budgets
|
|
SET name = ?, total_amount = ?, start_date = ?, end_date = ?, currency = ?
|
|
WHERE id = ? AND user_id = ?
|
|
""",
|
|
(
|
|
next_name or "Бюджет",
|
|
next_total,
|
|
next_start.isoformat(),
|
|
next_end.isoformat(),
|
|
next_currency,
|
|
budget_id,
|
|
user_id,
|
|
),
|
|
)
|
|
await self._db.conn.commit()
|
|
updated = await self.get_budget_for_user(budget_id, user_id)
|
|
if updated is None:
|
|
raise RuntimeError("Failed to update budget")
|
|
return updated
|
|
|
|
async def set_budget_active(
|
|
self,
|
|
budget_id: int,
|
|
user_id: int,
|
|
is_active: bool,
|
|
) -> Budget:
|
|
budget = await self.get_budget_for_user(budget_id, user_id)
|
|
if budget is None:
|
|
raise ValueError("Бюджет не найден")
|
|
await self._db.conn.execute(
|
|
"UPDATE budgets SET is_active = ? WHERE id = ? AND user_id = ?",
|
|
(1 if is_active else 0, budget_id, user_id),
|
|
)
|
|
await self._db.conn.commit()
|
|
updated = await self.get_budget_for_user(budget_id, user_id)
|
|
if updated is None:
|
|
raise RuntimeError("Failed to update budget activity")
|
|
return updated
|
|
|
|
async def get_budget_by_id(self, budget_id: int) -> Budget | None:
|
|
cursor = await self._db.conn.execute(
|
|
"SELECT * FROM budgets WHERE id = ?",
|
|
(budget_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return _row_to_budget(row) if row else None
|
|
|
|
async def get_budget_for_user(self, budget_id: int, user_id: int) -> Budget | None:
|
|
cursor = await self._db.conn.execute(
|
|
"SELECT * FROM budgets WHERE id = ? AND user_id = ?",
|
|
(budget_id, user_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return _row_to_budget(row) if row else None
|
|
|
|
async def list_budgets_for_user(self, user_id: int) -> list[Budget]:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM budgets
|
|
WHERE user_id = ?
|
|
ORDER BY is_active DESC, id DESC
|
|
""",
|
|
(user_id,),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [_row_to_budget(row) for row in rows]
|
|
|
|
async def resolve_budget(
|
|
self,
|
|
user_id: int,
|
|
budget_id: int | None = None,
|
|
*,
|
|
require_active: bool = False,
|
|
) -> Budget | None:
|
|
if budget_id is not None:
|
|
budget = await self.get_budget_for_user(budget_id, user_id)
|
|
if budget is None:
|
|
return None
|
|
if require_active and not budget.is_active:
|
|
raise ValueError("Бюджет неактивен — включи его или выбери другой")
|
|
return budget
|
|
|
|
selected_id = await self.get_selected_budget_id(user_id)
|
|
if selected_id is not None:
|
|
selected = await self.get_budget_for_user(selected_id, user_id)
|
|
if selected is not None:
|
|
if not require_active or selected.is_active:
|
|
return selected
|
|
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM budgets
|
|
WHERE user_id = ?
|
|
ORDER BY is_active DESC, id DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
budget = _row_to_budget(row)
|
|
if require_active and not budget.is_active:
|
|
raise ValueError("Нет активного бюджета — создай или включи существующий")
|
|
return budget
|
|
|
|
async def list_all_budgets(self) -> list[Budget]:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM budgets
|
|
ORDER BY user_id ASC, is_active DESC, id DESC
|
|
"""
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [_row_to_budget(row) for row in rows]
|
|
|
|
async def add_expense(
|
|
self,
|
|
user_id: int,
|
|
budget_id: int,
|
|
amount: float,
|
|
note: str | None = None,
|
|
spent_at: date | None = None,
|
|
) -> Expense:
|
|
await self.ensure_user(user_id)
|
|
spent = spent_at or date.today()
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
INSERT INTO expenses(user_id, budget_id, amount, note, spent_at)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
RETURNING *
|
|
""",
|
|
(user_id, budget_id, amount, note, spent.isoformat()),
|
|
)
|
|
row = await cursor.fetchone()
|
|
await self._db.conn.commit()
|
|
if row is None:
|
|
raise RuntimeError("Failed to insert expense")
|
|
return _row_to_expense(row)
|
|
|
|
async def delete_last_expense(
|
|
self,
|
|
user_id: int,
|
|
budget_id: int | None = None,
|
|
) -> Expense | None:
|
|
if budget_id is None:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM expenses
|
|
WHERE user_id = ?
|
|
ORDER BY id DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id,),
|
|
)
|
|
else:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM expenses
|
|
WHERE user_id = ? AND budget_id = ?
|
|
ORDER BY id DESC
|
|
LIMIT 1
|
|
""",
|
|
(user_id, budget_id),
|
|
)
|
|
row = await cursor.fetchone()
|
|
if row is None:
|
|
return None
|
|
await self._db.conn.execute("DELETE FROM expenses WHERE id = ?", (row["id"],))
|
|
await self._db.conn.commit()
|
|
return _row_to_expense(row)
|
|
|
|
async def spent_on_date(
|
|
self,
|
|
budget_id: int,
|
|
day: date,
|
|
) -> float:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT COALESCE(SUM(amount), 0) AS total
|
|
FROM expenses
|
|
WHERE budget_id = ? AND spent_at = ?
|
|
""",
|
|
(budget_id, day.isoformat()),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return float(row["total"]) if row else 0.0
|
|
|
|
async def list_expenses_on_date(
|
|
self,
|
|
budget_id: int,
|
|
day: date,
|
|
) -> list[Expense]:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM expenses
|
|
WHERE budget_id = ? AND spent_at = ?
|
|
ORDER BY id ASC
|
|
""",
|
|
(budget_id, day.isoformat()),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [_row_to_expense(row) for row in rows]
|
|
|
|
async def count_expenses_for_budget(self, budget_id: int) -> int:
|
|
cursor = await self._db.conn.execute(
|
|
"SELECT COUNT(*) AS cnt FROM expenses WHERE budget_id = ?",
|
|
(budget_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return int(row["cnt"]) if row else 0
|
|
|
|
async def sum_expenses_for_budget(self, budget_id: int) -> float:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT COALESCE(SUM(amount), 0) AS total
|
|
FROM expenses
|
|
WHERE budget_id = ?
|
|
""",
|
|
(budget_id,),
|
|
)
|
|
row = await cursor.fetchone()
|
|
return float(row["total"]) if row else 0.0
|
|
|
|
async def list_expenses_for_budget(
|
|
self,
|
|
budget_id: int,
|
|
*,
|
|
limit: int,
|
|
offset: int,
|
|
) -> list[Expense]:
|
|
cursor = await self._db.conn.execute(
|
|
"""
|
|
SELECT * FROM expenses
|
|
WHERE budget_id = ?
|
|
ORDER BY spent_at DESC, id DESC
|
|
LIMIT ? OFFSET ?
|
|
""",
|
|
(budget_id, limit, offset),
|
|
)
|
|
rows = await cursor.fetchall()
|
|
return [_row_to_expense(row) for row in rows]
|
|
|
|
async def clear_expenses_for_budget_id(self, budget_id: int) -> None:
|
|
await self._db.conn.execute(
|
|
"DELETE FROM expenses WHERE budget_id = ?",
|
|
(budget_id,),
|
|
)
|
|
await self._db.conn.commit()
|