Files
agentbox-test/models.py
T

305 lines
11 KiB
Python

"""Database models and helpers for the meal tracker."""
import sqlite3
import os
from datetime import date, datetime
def _db_path() -> str:
return os.environ.get(
"DB_PATH", os.path.join(os.path.dirname(__file__), "meals.db")
)
def get_db() -> sqlite3.Connection:
"""Return a connection to the SQLite database with foreign keys enabled."""
conn = sqlite3.connect(_db_path())
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA foreign_keys = ON")
return conn
def init_db() -> None:
"""Create tables if they don't exist, including migration for household support."""
with get_db() as db:
db.executescript("""
CREATE TABLE IF NOT EXISTS households (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT UNIQUE NOT NULL,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
username TEXT UNIQUE NOT NULL,
password_hash TEXT NOT NULL,
is_admin INTEGER NOT NULL DEFAULT 0,
household_id INTEGER REFERENCES households(id),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS meal_periods (
id INTEGER PRIMARY KEY AUTOINCREMENT,
date TEXT NOT NULL,
meal_type TEXT NOT NULL CHECK(meal_type IN ('lunch', 'dinner')),
UNIQUE(date, meal_type)
);
CREATE TABLE IF NOT EXISTS responses (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL REFERENCES users(id),
period_id INTEGER NOT NULL REFERENCES meal_periods(id),
status TEXT NOT NULL DEFAULT 'not_answered'
CHECK(status IN ('yes', 'no', 'not_answered')),
changed_at TIMESTAMP,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
UNIQUE(user_id, period_id)
);
CREATE INDEX IF NOT EXISTS idx_responses_user
ON responses(user_id);
CREATE INDEX IF NOT EXISTS idx_responses_period
ON responses(period_id);
CREATE INDEX IF NOT EXISTS idx_meal_periods_date
ON meal_periods(date);
""")
# Migration: if users table exists but lacks household_id column, add it.
cols = [r["name"] for r in db.execute("PRAGMA table_info(users)").fetchall()]
if "household_id" not in cols:
db.execute("ALTER TABLE users ADD COLUMN household_id INTEGER REFERENCES households(id)")
db.execute("CREATE INDEX IF NOT EXISTS idx_users_household ON users(household_id)")
# ── Households ────────────────────────────────────────────────────────────────
def create_household(name: str) -> dict:
"""Create a new household. Returns the row as a dict."""
with get_db() as db:
cur = db.execute("INSERT INTO households (name) VALUES (?)", (name,))
db.commit()
row = db.execute("SELECT * FROM households WHERE id = ?", (cur.lastrowid,)).fetchone()
return dict(row)
def get_household_by_id(household_id: int) -> dict | None:
"""Fetch a household by ID."""
with get_db() as db:
row = db.execute("SELECT * FROM households WHERE id = ?", (household_id,)).fetchone()
return dict(row) if row else None
def get_household_by_name(name: str) -> dict | None:
"""Fetch a household by name."""
with get_db() as db:
row = db.execute("SELECT * FROM households WHERE name = ?", (name,)).fetchone()
return dict(row) if row else None
def get_all_households() -> list[dict]:
"""Return all households."""
with get_db() as db:
rows = db.execute("SELECT * FROM households ORDER BY name").fetchall()
return [dict(r) for r in rows]
def count_household_users(household_id: int) -> int:
"""Return number of users in a household."""
with get_db() as db:
return db.execute(
"SELECT COUNT(*) FROM users WHERE household_id = ?", (household_id,)
).fetchone()[0]
def delete_household(household_id: int) -> None:
"""Delete a household, all its users, and all their responses."""
with get_db() as db:
# Delete responses for all users in the household
db.execute(
"DELETE FROM responses WHERE user_id IN "
"(SELECT id FROM users WHERE household_id = ?)",
(household_id,),
)
db.execute("DELETE FROM users WHERE household_id = ?", (household_id,))
db.execute("DELETE FROM households WHERE id = ?", (household_id,))
# ── Users ─────────────────────────────────────────────────────────────────────
def get_user_by_id(user_id: int) -> dict | None:
"""Fetch a single user by ID."""
with get_db() as db:
row = db.execute("SELECT * FROM users WHERE id = ?", (user_id,)).fetchone()
return dict(row) if row else None
def get_user_by_username(username: str) -> dict | None:
"""Fetch a single user by username."""
with get_db() as db:
row = db.execute(
"SELECT * FROM users WHERE username = ?", (username,)
).fetchone()
return dict(row) if row else None
def create_user(username: str, password_hash: str, household_id: int,
is_admin: bool = False) -> dict:
"""Insert a new user into a household. Returns the user row as a dict."""
with get_db() as db:
cur = db.execute(
"INSERT INTO users (username, password_hash, is_admin, household_id) "
"VALUES (?, ?, ?, ?)",
(username, password_hash, int(is_admin), household_id),
)
user_id = cur.lastrowid
db.commit()
return get_user_by_id(user_id)
def delete_user(user_id: int) -> None:
"""Delete a user and all their responses."""
with get_db() as db:
db.execute("DELETE FROM responses WHERE user_id = ?", (user_id,))
db.execute("DELETE FROM users WHERE id = ?", (user_id,))
def get_household_users(household_id: int) -> list[dict]:
"""Return all users in a given household."""
with get_db() as db:
rows = db.execute(
"SELECT id, username, is_admin, created_at "
"FROM users WHERE household_id = ? ORDER BY username",
(household_id,),
).fetchall()
return [dict(r) for r in rows]
# ── Meal periods ──────────────────────────────────────────────────────────────
def ensure_meal_periods(for_date: date) -> list[dict]:
"""Ensure lunch and dinner periods exist for a given date.
Returns the two period rows (as dicts) for that date.
"""
date_str = for_date.isoformat()
with get_db() as db:
db.execute(
"INSERT OR IGNORE INTO meal_periods (date, meal_type) VALUES (?, 'lunch')",
(date_str,),
)
db.execute(
"INSERT OR IGNORE INTO meal_periods (date, meal_type) VALUES (?, 'dinner')",
(date_str,),
)
rows = db.execute(
"SELECT id, date, meal_type FROM meal_periods WHERE date = ? ORDER BY meal_type",
(date_str,),
).fetchall()
return [dict(r) for r in rows]
# ── Dashboard ─────────────────────────────────────────────────────────────────
def get_dashboard_data(for_date: date, household_id: int) -> dict:
"""Return all users and their responses for a given date, scoped to a household.
Returns:
{
"date": "YYYY-MM-DD",
"periods": [{"id": 1, "date": "...", "meal_type": "lunch"}, ...],
"users": [
{
"id": 1,
"username": "alice",
"responses": {
"lunch": {"status": "yes", "changed_at": None},
"dinner": {"status": "no", "changed_at": "2025-01-15 16:30:00"},
}
},
...
]
}
"""
periods = ensure_meal_periods(for_date)
users = get_household_users(household_id)
with get_db() as db:
period_ids = [p["id"] for p in periods]
if not period_ids:
return {"date": for_date.isoformat(), "periods": periods, "users": []}
placeholders = ",".join("?" * len(period_ids))
resp_rows = db.execute(
f"SELECT user_id, period_id, status, changed_at "
f"FROM responses "
f"WHERE period_id IN ({placeholders})",
period_ids,
).fetchall()
resp_map = {}
for r in resp_rows:
resp_map[(r["user_id"], r["period_id"])] = {
"status": r["status"],
"changed_at": r["changed_at"],
}
period_map = {p["meal_type"]: p["id"] for p in periods}
enriched_users = []
for u in users:
user_responses = {}
for meal_type, pid in period_map.items():
key = (u["id"], pid)
if key in resp_map:
user_responses[meal_type] = resp_map[key]
else:
user_responses[meal_type] = {"status": "not_answered", "changed_at": None}
enriched_users.append({**u, "responses": user_responses})
return {
"date": for_date.isoformat(),
"periods": periods,
"users": enriched_users,
}
# ── Responses ─────────────────────────────────────────────────────────────────
def upsert_response(user_id: int, period_id: int, new_status: str) -> dict | None:
"""Insert or update a response. Handles changed_at logic.
changed_at is set when flipping between 'yes' and 'no'.
It stays NULL on first answer or when status doesn't change.
"""
with get_db() as db:
existing = db.execute(
"SELECT status, changed_at FROM responses WHERE user_id = ? AND period_id = ?",
(user_id, period_id),
).fetchone()
if existing is None:
db.execute(
"INSERT INTO responses (user_id, period_id, status) VALUES (?, ?, ?)",
(user_id, period_id, new_status),
)
elif existing["status"] == new_status:
pass
elif existing["status"] in ("yes", "no") and new_status in ("yes", "no"):
now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
db.execute(
"UPDATE responses SET status = ?, changed_at = ? "
"WHERE user_id = ? AND period_id = ?",
(new_status, now, user_id, period_id),
)
else:
db.execute(
"UPDATE responses SET status = ? WHERE user_id = ? AND period_id = ?",
(new_status, user_id, period_id),
)
row = db.execute(
"SELECT * FROM responses WHERE user_id = ? AND period_id = ?",
(user_id, period_id),
).fetchone()
return dict(row) if row else None