101 lines
2.6 KiB
Python
101 lines
2.6 KiB
Python
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:
|
|
database = g.pop("db", None)
|
|
if database is not None:
|
|
database.close()
|
|
|
|
|
|
def apply_schema(database: sqlite3.Connection) -> None:
|
|
schema_path = Path(__file__).with_name("schema.sql")
|
|
database.executescript(schema_path.read_text(encoding="utf-8"))
|
|
sync_dayparts(database)
|
|
|
|
|
|
def init_db() -> None:
|
|
database = get_db()
|
|
apply_schema(database)
|
|
database.commit()
|
|
|
|
|
|
def sync_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"]),
|
|
)
|
|
database.execute(
|
|
"""
|
|
UPDATE dayparts
|
|
SET name = ?, sort_order = ?
|
|
WHERE slug = ?
|
|
""",
|
|
(entry["name"], entry["sort_order"], entry["slug"]),
|
|
)
|
|
|
|
|
|
def init_db_if_needed(app: Flask) -> None:
|
|
with app.app_context():
|
|
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)
|