db_create.py

#!/usr/bin/env python3
import sqlalchemy

from nyaa import create_app, models
from nyaa.extensions import db

app = create_app("config")

NYAA_CATEGORIES = [
    ("Unarchived", ["Karaoke", "Copyright", "Yab"]),
    ("Paid", ["Membership", "Voice pack", "Concert", "Songs", "Others"]),
]


SUKEBEI_CATEGORIES = [
    ("Art", ["Anime", "Doujinshi", "Games", "Manga", "Pictures"]),
    ("Real Life", ["Photobooks / Pictures", "Videos"]),
]


def add_categories(categories, main_class, sub_class):
    for main_cat_name, sub_cat_names in categories:
        main_cat = main_class(name=main_cat_name)
        for i, sub_cat_name in enumerate(sub_cat_names):
            # Composite keys can't autoincrement, set sub_cat id manually (1-index)
            sub_cat = sub_class(id=i + 1, name=sub_cat_name, main_category=main_cat)
        db.session.add(main_cat)


if __name__ == "__main__":
    with app.app_context():
        # Test for the user table, assume db is empty if it's not created
        database_empty = False
        try:
            models.User.query.first()
        except (sqlalchemy.exc.ProgrammingError, sqlalchemy.exc.OperationalError):
            database_empty = True

        print("Creating all tables...")
        db.create_all()

        nyaa_category_test = models.NyaaMainCategory.query.first()
        if not nyaa_category_test:
            print("Adding Nyaa categories...")
            add_categories(
                NYAA_CATEGORIES, models.NyaaMainCategory, models.NyaaSubCategory
            )

        sukebei_category_test = models.SukebeiMainCategory.query.first()
        if not sukebei_category_test:
            print("Adding Sukebei categories...")
            add_categories(
                SUKEBEI_CATEGORIES,
                models.SukebeiMainCategory,
                models.SukebeiSubCategory,
            )

        db.session.commit()

        if database_empty:
            print(
                "Remember to run the following to mark the database up-to-date for Alembic:"
            )
            print("./db_migrate.py stamp head")
            # Technically we should be able to do this here, but when you have
            # Flask-Migrate and Flask-SQA and everything... I didn't get it working.