diff --git a/db_create.py b/db_create.py index d881377..06a1783 100755 --- a/db_create.py +++ b/db_create.py @@ -1,35 +1,60 @@ +#!/usr/bin/env python3 import sys +import sqlalchemy from nyaa import app, db, models -# Create tables +NYAA_CATEGORIES = [ + ('Anime', ['Anime Music Video', 'English-translated', 'Non-English-translated', 'Raw']), + ('Audio', ['Lossless', 'Lossy']), + ('Literature', ['English-translated', 'Non-English-translated', 'Raw']), + ('Live Action', ['English-translated', 'Idol/Promotional Video', 'Non-English-translated', 'Raw']), + ('Pictures', ['Graphics', 'Photos']), + ('Software', ['Applications', 'Games']), +] -db.create_all() -# Insert categories and insert if it doesn't eixst -existing_cats = models.MainCategory.query.all() -if not existing_cats: - if app.config['SITE_FLAVOR'] == 'nyaa': - CATEGORIES = [ - ('Anime', ['Anime Music Video', 'English-translated', 'Non-English-translated', 'Raw']), - ('Audio', ['Lossless', 'Lossy']), - ('Literature', ['English-translated', 'Non-English-translated', 'Raw']), - ('Live Action', ['English-translated', 'Idol/Promotional Video', 'Non-English-translated', 'Raw']), - ('Pictures', ['Graphics', 'Photos']), - ('Software', ['Applications', 'Games']), - ] - elif app.config['SITE_FLAVOR'] == 'sukebei': - CATEGORIES = [ - ('Art', ['Anime', 'Doujinshi', 'Games', 'Manga', 'Pictures']), - ('Real Life', ['Photobooks / Pictures', 'Videos']), - ] - else: - CATEGORIES = [] +SUKEBEI_CATEGORIES = [ + ('Art', ['Anime', 'Doujinshi', 'Games', 'Manga', 'Pictures']), + ('Real Life', ['Photobooks / Pictures', 'Videos']), +] - for main_cat_name, sub_cat_names in CATEGORIES: - main_cat = models.MainCategory(name=main_cat_name) + +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 = models.SubCategory(id=i+1, name=sub_cat_name, main_category=main_cat) + sub_cat = sub_class(id=i+1, name=sub_cat_name, main_category=main_cat) db.session.add(main_cat) + +if __name__ == '__main__': + # 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 db 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. \ No newline at end of file