Update db_create.py to add categories for both Nyaa and Sukebei

Also includes a reminder to stamp the database version for Alembic.

db_create.py

@@ -1,15 +1,9 @@
+#!/usr/bin/env python3
 import sys
+import sqlalchemy
 from nyaa import app, db, models
 
-# Create tables
-
-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 = [
+NYAA_CATEGORIES = [
     ('Anime', ['Anime Music Video', 'English-translated', 'Non-English-translated', 'Raw']),
     ('Audio', ['Lossless', 'Lossy']),
     ('Literature', ['English-translated', 'Non-English-translated', 'Raw']),
@@ -17,19 +11,50 @@ if not existing_cats:
     ('Pictures', ['Graphics', 'Photos']),
     ('Software', ['Applications', 'Games']),
 ]
-    elif app.config['SITE_FLAVOR'] == 'sukebei':
-        CATEGORIES = [
+
+
+SUKEBEI_CATEGORIES = [
     ('Art', ['Anime', 'Doujinshi', 'Games', 'Manga', 'Pictures']),
     ('Real Life', ['Photobooks / Pictures', 'Videos']),
 ]
-    else:
-        CATEGORIES = []
 
-    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.