first commit

nouri/db.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+
+import click
+from flask import Flask, current_app, g
+from flask.cli import with_appcontext
+from werkzeug.security import generate_password_hash
+
+from .constants import DAYPARTS
+
+
+def get_db() -> sqlite3.Connection:
+    if "db" not in g:
+        g.db = sqlite3.connect(
+            current_app.config["DATABASE_PATH"],
+            detect_types=sqlite3.PARSE_DECLTYPES,
+        )
+        g.db.row_factory = sqlite3.Row
+        g.db.execute("PRAGMA foreign_keys = ON")
+    return g.db
+
+
+def close_db(_error=None) -> None:
+    db = g.pop("db", None)
+    if db is not None:
+        db.close()
+
+
+def init_db() -> None:
+    database = get_db()
+    schema_path = Path(__file__).with_name("schema.sql")
+    database.executescript(schema_path.read_text(encoding="utf-8"))
+    seed_dayparts(database)
+    database.commit()
+
+
+def seed_dayparts(database: sqlite3.Connection) -> None:
+    for entry in DAYPARTS:
+        database.execute(
+            """
+            INSERT OR IGNORE INTO dayparts (slug, name, sort_order)
+            VALUES (?, ?, ?)
+            """,
+            (entry["slug"], entry["name"], entry["sort_order"]),
+        )
+
+
+def init_db_if_needed(app: Flask) -> None:
+    db_path = Path(app.config["DATABASE_PATH"])
+    needs_init = not db_path.exists()
+    with app.app_context():
+        if needs_init:
+            init_db()
+            return
+
+        database = get_db()
+        table = database.execute(
+            "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'users'"
+        ).fetchone()
+        if table is None:
+            init_db()
+
+
+def user_count() -> int:
+    row = get_db().execute("SELECT COUNT(*) AS count FROM users").fetchone()
+    return int(row["count"])
+
+
+@click.command("init-db")
+@with_appcontext
+def init_db_command() -> None:
+    init_db()
+    click.echo("Database initialized.")
+
+
+@click.command("create-user")
+@click.argument("username")
+@click.argument("password")
+@click.option("--display-name", default="", help="Friendly display name.")
+@with_appcontext
+def create_user_command(username: str, password: str, display_name: str) -> None:
+    database = get_db()
+    database.execute(
+        """
+        INSERT INTO users (username, display_name, password_hash)
+        VALUES (?, ?, ?)
+        """,
+        (username.strip().lower(), display_name.strip(), generate_password_hash(password)),
+    )
+    database.commit()
+    click.echo(f"User '{username}' created.")
+
+
+def init_app(app: Flask) -> None:
+    app.teardown_appcontext(close_db)
+    app.cli.add_command(init_db_command)
+    app.cli.add_command(create_user_command)