244 lines
8.4 KiB
Python
244 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import aiosqlite
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
user_id INTEGER PRIMARY KEY,
|
|
selected_budget_id INTEGER,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS budgets (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
name TEXT NOT NULL DEFAULT '',
|
|
total_amount REAL NOT NULL,
|
|
start_date TEXT NOT NULL,
|
|
end_date TEXT NOT NULL,
|
|
currency TEXT NOT NULL DEFAULT 'RUB',
|
|
is_active INTEGER NOT NULL DEFAULT 1,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS expenses (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
budget_id INTEGER NOT NULL,
|
|
amount REAL NOT NULL,
|
|
note TEXT,
|
|
spent_at TEXT NOT NULL,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE,
|
|
FOREIGN KEY (budget_id) REFERENCES budgets(id) ON DELETE CASCADE
|
|
);
|
|
"""
|
|
|
|
|
|
async def _ensure_indexes(conn: aiosqlite.Connection) -> None:
|
|
await conn.executescript(
|
|
"""
|
|
CREATE INDEX IF NOT EXISTS idx_budgets_user
|
|
ON budgets(user_id, is_active, id);
|
|
CREATE INDEX IF NOT EXISTS idx_expenses_budget_spent_at
|
|
ON expenses(budget_id, spent_at);
|
|
CREATE INDEX IF NOT EXISTS idx_expenses_user_spent_at
|
|
ON expenses(user_id, spent_at);
|
|
"""
|
|
)
|
|
|
|
|
|
async def _table_columns(conn: aiosqlite.Connection, table: str) -> set[str]:
|
|
cursor = await conn.execute(f"PRAGMA table_info({table})")
|
|
rows = await cursor.fetchall()
|
|
return {str(row[1]) for row in rows}
|
|
|
|
|
|
async def _has_unique_user_on_budgets(conn: aiosqlite.Connection) -> bool:
|
|
cursor = await conn.execute("PRAGMA index_list(budgets)")
|
|
indexes = await cursor.fetchall()
|
|
for idx in indexes:
|
|
# (seq, name, unique, origin, partial)
|
|
if not idx[2]:
|
|
continue
|
|
name = idx[1]
|
|
info = await conn.execute(f"PRAGMA index_info({name})")
|
|
cols = [row[2] for row in await info.fetchall()]
|
|
if cols == ["user_id"]:
|
|
return True
|
|
return False
|
|
|
|
|
|
async def migrate_schema(conn: aiosqlite.Connection) -> None:
|
|
"""Upgrade legacy one-budget-per-user schema in place."""
|
|
tables = {
|
|
row[0]
|
|
for row in await (
|
|
await conn.execute(
|
|
"SELECT name FROM sqlite_master WHERE type='table'"
|
|
)
|
|
).fetchall()
|
|
}
|
|
if "budgets" not in tables:
|
|
return
|
|
|
|
budget_cols = await _table_columns(conn, "budgets")
|
|
if "name" not in budget_cols:
|
|
await conn.execute(
|
|
"ALTER TABLE budgets ADD COLUMN name TEXT NOT NULL DEFAULT ''"
|
|
)
|
|
if "is_active" not in budget_cols:
|
|
await conn.execute(
|
|
"ALTER TABLE budgets ADD COLUMN is_active INTEGER NOT NULL DEFAULT 1"
|
|
)
|
|
|
|
user_cols = await _table_columns(conn, "users")
|
|
if "selected_budget_id" not in user_cols:
|
|
await conn.execute(
|
|
"ALTER TABLE users ADD COLUMN selected_budget_id INTEGER"
|
|
)
|
|
|
|
if await _has_unique_user_on_budgets(conn):
|
|
await conn.executescript(
|
|
"""
|
|
CREATE TABLE budgets_migrated (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
name TEXT NOT NULL DEFAULT '',
|
|
total_amount REAL NOT NULL,
|
|
start_date TEXT NOT NULL,
|
|
end_date TEXT NOT NULL,
|
|
currency TEXT NOT NULL DEFAULT 'RUB',
|
|
is_active INTEGER NOT NULL DEFAULT 1,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE
|
|
);
|
|
INSERT INTO budgets_migrated (
|
|
id, user_id, name, total_amount, start_date, end_date,
|
|
currency, is_active, created_at
|
|
)
|
|
SELECT
|
|
id, user_id,
|
|
COALESCE(NULLIF(name, ''), 'Бюджет'),
|
|
total_amount, start_date, end_date, currency,
|
|
COALESCE(is_active, 1), created_at
|
|
FROM budgets;
|
|
DROP TABLE budgets;
|
|
ALTER TABLE budgets_migrated RENAME TO budgets;
|
|
"""
|
|
)
|
|
|
|
expense_cols = await _table_columns(conn, "expenses")
|
|
if "budget_id" not in expense_cols:
|
|
await conn.execute("ALTER TABLE expenses ADD COLUMN budget_id INTEGER")
|
|
await conn.execute(
|
|
"""
|
|
UPDATE expenses
|
|
SET budget_id = (
|
|
SELECT b.id FROM budgets b
|
|
WHERE b.user_id = expenses.user_id
|
|
ORDER BY b.id DESC
|
|
LIMIT 1
|
|
)
|
|
WHERE budget_id IS NULL
|
|
"""
|
|
)
|
|
# Drop orphan expenses that have no budget (should be rare)
|
|
await conn.execute("DELETE FROM expenses WHERE budget_id IS NULL")
|
|
await conn.executescript(
|
|
"""
|
|
CREATE TABLE expenses_migrated (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
budget_id INTEGER NOT NULL,
|
|
amount REAL NOT NULL,
|
|
note TEXT,
|
|
spent_at TEXT NOT NULL,
|
|
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
|
FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE,
|
|
FOREIGN KEY (budget_id) REFERENCES budgets(id) ON DELETE CASCADE
|
|
);
|
|
INSERT INTO expenses_migrated (
|
|
id, user_id, budget_id, amount, note, spent_at, created_at
|
|
)
|
|
SELECT id, user_id, budget_id, amount, note, spent_at, created_at
|
|
FROM expenses;
|
|
DROP TABLE expenses;
|
|
ALTER TABLE expenses_migrated RENAME TO expenses;
|
|
"""
|
|
)
|
|
|
|
# Backfill selected budget for users who have budgets
|
|
await conn.execute(
|
|
"""
|
|
UPDATE users
|
|
SET selected_budget_id = (
|
|
SELECT b.id FROM budgets b
|
|
WHERE b.user_id = users.user_id
|
|
ORDER BY b.is_active DESC, b.id DESC
|
|
LIMIT 1
|
|
)
|
|
WHERE selected_budget_id IS NULL
|
|
AND EXISTS (SELECT 1 FROM budgets b WHERE b.user_id = users.user_id)
|
|
"""
|
|
)
|
|
await conn.execute(
|
|
"""
|
|
UPDATE budgets
|
|
SET name = 'Бюджет'
|
|
WHERE name IS NULL OR TRIM(name) = ''
|
|
"""
|
|
)
|
|
await _ensure_indexes(conn)
|
|
|
|
|
|
class Database:
|
|
def __init__(self, path: Path, *, read_only: bool = False) -> None:
|
|
self.path = path
|
|
self.read_only = read_only
|
|
self._conn: aiosqlite.Connection | None = None
|
|
|
|
async def connect(self) -> None:
|
|
if self.read_only:
|
|
if not self.path.exists():
|
|
raise FileNotFoundError(
|
|
f"Database not found: {self.path}. "
|
|
"Сначала запусти API, чтобы создался data/budget.db"
|
|
)
|
|
uri = f"file:{self.path.resolve().as_posix()}?mode=ro"
|
|
self._conn = await aiosqlite.connect(uri, uri=True)
|
|
self._conn.row_factory = aiosqlite.Row
|
|
await self._conn.execute("PRAGMA foreign_keys = ON")
|
|
return
|
|
|
|
self.path.parent.mkdir(parents=True, exist_ok=True)
|
|
self._conn = await aiosqlite.connect(self.path)
|
|
self._conn.row_factory = aiosqlite.Row
|
|
await self._conn.execute("PRAGMA foreign_keys = ON")
|
|
try:
|
|
await self._conn.execute("PRAGMA journal_mode=WAL")
|
|
except aiosqlite.OperationalError:
|
|
pass
|
|
await self._conn.executescript(SCHEMA)
|
|
await migrate_schema(self._conn)
|
|
# Fresh DBs: migrate may no-op early if tables were just created with
|
|
# full columns — still ensure indexes exist.
|
|
budget_cols = await _table_columns(self._conn, "budgets")
|
|
if "is_active" in budget_cols:
|
|
await _ensure_indexes(self._conn)
|
|
await self._conn.commit()
|
|
|
|
async def close(self) -> None:
|
|
if self._conn is not None:
|
|
await self._conn.close()
|
|
self._conn = None
|
|
|
|
@property
|
|
def conn(self) -> aiosqlite.Connection:
|
|
if self._conn is None:
|
|
raise RuntimeError("Database is not connected")
|
|
return self._conn
|