f1tness_parser/backend/app/cli.py

"""Command line helpers.

    python -m app.cli import [--replace] [--trainee-email EMAIL]
    python -m app.cli set-password EMAIL [--superuser]

Imports all trainings from the Obsidian and Apple notes into the database as trainings of one
trainee. The trainer from a note header is matched with a user of the trainer role by
"First Last" ("соло" means no trainer). The schema must already exist (`alembic upgrade head`).

`set-password` asks for a new password of the user (and with --superuser makes them a
superuser). It is the way to get a working login: users are created only by superusers, and
the users made by migrations have random passwords.
"""

import argparse
import asyncio
import getpass
import logging
from typing import Any, Dict, List, Optional, Set

from app.config import settings
from app.core.database.connection import create_engine, create_session_factory
from app.core.database.repository import TrainingRepository
from app.core.database.token_repository import RefreshTokenRepository
from app.core.database.user_repository import UserRepository
from app.core.dto.training import TrainingDTO
from app.core.dto.user import UserRole
from app.core.security import hash_password
from app.core.parsers.apple import AppleNotesParser
from app.core.parsers.obsidian import ObsidianNotesParser


SOLO_NAMES = {"solo", "соло"}


async def assign_participants(
    session, trainings: List[TrainingDTO], trainee_email: Optional[str]
) -> None:
    """Set trainee_id and trainer_id on the parsed trainings."""
    users = await UserRepository(session).list()
    trainees = [user for user in users if user.role == UserRole.trainee]
    if trainee_email:
        matching = [u for u in trainees if u.email == trainee_email.strip().lower()]
        if not matching:
            raise SystemExit(f"No trainee with email {trainee_email}")
        trainee = matching[0]
    elif len(trainees) == 1:
        trainee = trainees[0]
    else:
        raise SystemExit("There is more than one trainee, pass --trainee-email")

    trainer_ids = {
        f"{user.first_name} {user.last_name}".strip().casefold(): user.id
        for user in users
        if user.role == UserRole.trainer
    }
    unknown: Set[str] = set()
    for training in trainings:
        training.trainee_id = trainee.id
        name = (training.trainer or "").strip().casefold()
        if name in SOLO_NAMES or not name:
            training.trainer_id = None
        elif name in trainer_ids:
            training.trainer_id = trainer_ids[name]
        else:
            unknown.add(training.trainer or "")
    if unknown:
        raise SystemExit(f"No trainer users for: {', '.join(sorted(unknown))}")


async def import_notes(replace: bool, trainee_email: Optional[str]) -> None:
    parsers = [ObsidianNotesParser(), AppleNotesParser()]
    trainings = []
    for parser in parsers:
        parsed = parser.parse_file()
        print(
            f"{type(parser).__name__}: {len(parsed)} trainings, "
            f"{len(parser.errors)} skipped lines"
        )
        trainings.extend(parsed)

    engine = create_engine(settings.get_postgres_database_url())
    try:
        async with create_session_factory(engine)() as session:
            repository = TrainingRepository(session)
            await assign_participants(session, trainings, trainee_email)
            existing = await repository.count()
            if existing and not replace:
                raise SystemExit(
                    f"Database already has {existing} trainings, "
                    "use --replace to overwrite them"
                )
            if replace:
                await repository.delete_all()
            await repository.save_many(trainings)
    finally:
        await engine.dispose()
    print(f"Imported {len(trainings)} trainings")


async def set_password(email: str, password: str, superuser: bool) -> None:
    engine = create_engine(settings.get_postgres_database_url())
    try:
        async with create_session_factory(engine)() as session:
            users = UserRepository(session)
            found = await users.get_credentials(email.strip().lower())
            if found is None:
                raise SystemExit(f"No user with email {email}")
            user, _ = found
            values: Dict[str, Any] = {"password_hash": hash_password(password)}
            if superuser:
                values["is_superuser"] = True
            await users.update(user.id, values)
            # Whoever was logged in with the old password is logged out
            await RefreshTokenRepository(session).revoke_all_of_user(user.id)
    finally:
        await engine.dispose()
    print(f"Password of {email} is changed")


def main() -> None:
    logging.basicConfig(level=logging.WARNING)
    arg_parser = argparse.ArgumentParser(prog="app.cli")
    subcommands = arg_parser.add_subparsers(dest="command", required=True)
    import_parser = subcommands.add_parser("import", help="import notes into the DB")
    import_parser.add_argument(
        "--replace", action="store_true", help="delete existing trainings first"
    )
    import_parser.add_argument(
        "--trainee-email",
        help="whose trainings these are (not needed when there is a single trainee)",
    )
    password_parser = subcommands.add_parser("set-password", help="set a user's password")
    password_parser.add_argument("email")
    password_parser.add_argument(
        "--superuser", action="store_true", help="also make the user a superuser"
    )
    args = arg_parser.parse_args()

    if args.command == "import":
        asyncio.run(import_notes(args.replace, args.trainee_email))
    elif args.command == "set-password":
        password = getpass.getpass("New password: ")
        if len(password) < 8:
            raise SystemExit("The password must be at least 8 characters long")
        if password != getpass.getpass("Repeat: "):
            raise SystemExit("The passwords differ")
        asyncio.run(set_password(args.email, password, args.superuser))


if __name__ == "__main__":
    main()