feat(proj): init
This commit is contained in:
vl.arkhangelskii
@@ -0,0 +1,400 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user