from __future__ import annotations from collections.abc import Iterable import os from pathlib import Path import subprocess import sys import click from flask import Flask, current_app from flask.cli import with_appcontext from sqlalchemy import inspect from app.extensions import db from app.models import Account, Household, HouseholdMember, OptionItem from app.services import demo_login_lines, seed_demo_accounts, seed_demo_households, seed_stress_households DEFAULT_ADMIN_USERNAME = "admin" DEFAULT_ADMIN_DISPLAY_NAME = "系统管理员" DEFAULT_ADMIN_PASSWORD = "ChangeMe123!" BASE_OPTION_ITEMS: tuple[dict[str, str | int | bool | None], ...] = ( {"option_group": "relation_category", "option_code": "relative", "option_label": "亲戚", "parent_code": None, "sort_order": 10}, {"option_group": "relation_category", "option_code": "friend", "option_label": "朋友", "parent_code": None, "sort_order": 20}, {"option_group": "relation_category", "option_code": "colleague", "option_label": "同事", "parent_code": None, "sort_order": 30}, {"option_group": "gift_method", "option_code": "cash", "option_label": "现金", "parent_code": None, "sort_order": 10}, {"option_group": "gift_method", "option_code": "transfer", "option_label": "转账", "parent_code": None, "sort_order": 20}, {"option_group": "gift_scene", "option_code": "pre_wedding", "option_label": "婚前", "parent_code": None, "sort_order": 10}, {"option_group": "gift_scene", "option_code": "wedding_day", "option_label": "婚礼当天", "parent_code": None, "sort_order": 20}, {"option_group": "gift_scene", "option_code": "post_wedding", "option_label": "婚后补记", "parent_code": None, "sort_order": 30}, ) def register_commands(app: Flask) -> None: app.cli.add_command(init_db_command) app.cli.add_command(seed_admin_command) app.cli.add_command(seed_options_command) app.cli.add_command(seed_all_command) app.cli.add_command(seed_demo_command) app.cli.add_command(seed_stress_command) app.cli.add_command(dev_bootstrap_command) @click.command("init-db") @with_appcontext def init_db_command() -> None: db.create_all() click.echo("Database tables created.") def ensure_seed_tables_exist() -> None: inspector = inspect(db.engine) required_tables = {Account.__tablename__, OptionItem.__tablename__} existing_tables = set(inspector.get_table_names()) if required_tables.issubset(existing_tables): return missing_tables = ", ".join(sorted(required_tables - existing_tables)) raise click.ClickException( f"Missing database tables: {missing_tables}. Run 'flask db upgrade' or 'flask init-db' first.", ) def ensure_demo_seed_tables_exist() -> None: inspector = inspect(db.engine) required_tables = { Account.__tablename__, OptionItem.__tablename__, Household.__tablename__, HouseholdMember.__tablename__, } existing_tables = set(inspector.get_table_names()) if required_tables.issubset(existing_tables): return missing_tables = ", ".join(sorted(required_tables - existing_tables)) raise click.ClickException( f"Missing database tables: {missing_tables}. Run 'flask db upgrade' or 'flask init-db' first.", ) @click.command("seed-admin") @click.option("--username", default=DEFAULT_ADMIN_USERNAME, show_default=True) @click.option("--display-name", default=DEFAULT_ADMIN_DISPLAY_NAME, show_default=True) @click.option("--password", default=DEFAULT_ADMIN_PASSWORD, show_default=True) @with_appcontext def seed_admin_command(username: str, display_name: str, password: str) -> None: ensure_seed_tables_exist() existing_account = db.session.execute( db.select(Account).where(Account.username == username), ).scalar_one_or_none() if existing_account is not None: click.echo(f"Admin account '{username}' already exists.") return account = Account() account.username = username account.display_name = display_name account.set_password(password) account.role = "admin" account.status = "active" db.session.add(account) db.session.commit() click.echo(f"Admin account '{username}' created.") @click.command("seed-options") @with_appcontext def seed_options_command() -> None: ensure_seed_tables_exist() created_count = seed_option_items(BASE_OPTION_ITEMS) db.session.commit() click.echo(f"Seeded {created_count} option items.") @click.command("seed-all") @click.option("--username", default=DEFAULT_ADMIN_USERNAME, show_default=True) @click.option("--display-name", default=DEFAULT_ADMIN_DISPLAY_NAME, show_default=True) @click.option("--password", default=DEFAULT_ADMIN_PASSWORD, show_default=True) @with_appcontext def seed_all_command(username: str, display_name: str, password: str) -> None: ensure_seed_tables_exist() created_options = seed_option_items(BASE_OPTION_ITEMS) existing_account = db.session.execute( db.select(Account).where(Account.username == username), ).scalar_one_or_none() if existing_account is None: admin_account = Account() admin_account.username = username admin_account.display_name = display_name admin_account.set_password(password) admin_account.role = "admin" admin_account.status = "active" db.session.add(admin_account) admin_created = True else: admin_created = False db.session.commit() if admin_created: click.echo(f"Admin account '{username}' created.") else: click.echo(f"Admin account '{username}' already exists.") click.echo(f"Seeded {created_options} option items.") @click.command("seed-demo") @with_appcontext def seed_demo_command() -> None: ensure_demo_seed_tables_exist() created_accounts = seed_demo_accounts() created_households = seed_demo_households() db.session.commit() click.echo(f"Seeded {created_accounts} demo accounts.") click.echo(f"Seeded {created_households} demo households.") click.echo("Demo login accounts:") for login_line in demo_login_lines(): click.echo(f"- {login_line}") @click.command("seed-stress") @click.option("--household-count", default=1000, show_default=True, type=int) @click.option("--members-per-household", default=4, show_default=True, type=int) @click.option("--prefix", default="STRESS-", show_default=True) @with_appcontext def seed_stress_command(household_count: int, members_per_household: int, prefix: str) -> None: ensure_demo_seed_tables_exist() try: created_households, created_members = seed_stress_households( household_count=household_count, members_per_household=members_per_household, household_code_prefix=prefix, ) except ValueError as exc: raise click.ClickException(str(exc)) from exc db.session.commit() click.echo(f"Seeded {created_households} stress households.") click.echo(f"Seeded {created_members} stress household members.") @click.command("dev-bootstrap") @click.option("--host", default="127.0.0.1", show_default=True) @click.option("--port", default=5000, show_default=True, type=int) @click.option("--skip-server", is_flag=True, default=False) @with_appcontext def dev_bootstrap_command(host: str, port: int, skip_server: bool) -> None: app = current_app._get_current_object() project_root = os.fspath(Path(app.root_path).parent) python_executable = sys.executable env = os.environ.copy() env["FLASK_APP"] = "run.py" _run_python_command( [python_executable, "-m", "flask", "db", "upgrade"], env=env, cwd=project_root, label="Applying database migrations", ) _run_python_command( [python_executable, "-m", "flask", "seed-all"], env=env, cwd=project_root, label="Seeding baseline admin and options", ) _run_python_command( [python_executable, "-m", "flask", "seed-demo"], env=env, cwd=project_root, label="Seeding demo accounts and households", ) click.echo(f"Service URL: http://{host}:{port}") click.echo("Available demo logins:") for login_line in demo_login_lines(): click.echo(f"- {login_line}") if skip_server: click.echo("Skipped starting the Flask development server.") return click.echo("Starting Flask development server...") app.run(host=host, port=port, debug=True) def seed_option_items(option_rows: Iterable[dict[str, str | int | bool | None]]) -> int: created_count = 0 created_items_by_code: dict[tuple[str, str], OptionItem] = {} for row in option_rows: option_group = str(row["option_group"]) option_code = str(row["option_code"]) existing_item = db.session.execute( db.select(OptionItem).where( OptionItem.option_group == option_group, OptionItem.option_code == option_code, ), ).scalar_one_or_none() if existing_item is not None: created_items_by_code[(option_group, option_code)] = existing_item continue parent_id = None parent_code = row.get("parent_code") if parent_code is not None: parent_item = created_items_by_code.get((option_group.replace("_detail", "_category"), str(parent_code))) if parent_item is None: parent_item = db.session.execute( db.select(OptionItem).where(OptionItem.option_code == str(parent_code)), ).scalar_one_or_none() if parent_item is not None: parent_id = parent_item.id option_item = OptionItem() option_item.option_group = option_group option_item.option_code = option_code option_item.option_label = str(row["option_label"]) option_item.parent_id = parent_id option_item.sort_order = int(row.get("sort_order", 0) or 0) option_item.is_enabled = bool(row.get("is_enabled", True)) option_item.is_system = bool(row.get("is_system", True)) db.session.add(option_item) db.session.flush() created_items_by_code[(option_group, option_code)] = option_item created_count += 1 return created_count def _run_python_command( args: list[str], *, env: dict[str, str], cwd: str, label: str, ) -> None: click.echo(label) completed_process = subprocess.run( args, cwd=cwd, env=env, check=False, capture_output=True, text=True, ) if completed_process.stdout: click.echo(completed_process.stdout.rstrip()) if completed_process.returncode == 0: return if completed_process.stderr: click.echo(completed_process.stderr.rstrip(), err=True) raise click.ClickException(f"Command failed: {' '.join(args)}")