extending setup

setup.py

@@ -15,39 +15,46 @@ from web.utils import get_mysql_config
 
 cmd = sys.argv[1] if len(sys.argv) > 1 else "help"
 
-if cmd == "mysql-init":
-    cfg = get_mysql_config()
-    root_url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/"
-    dbname = cfg["database"]
-    root_engine = create_engine(root_url, future=True)
-    with root_engine.begin() as conn:
-        conn.execute(text(
-            f"CREATE DATABASE IF NOT EXISTS `{dbname}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"))
-    mysql_url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/{dbname}?charset=utf8mb4"
-    mysql_engine = create_engine(mysql_url, future=True)
-    from web.db import Base
-    Base.metadata.create_all(mysql_engine)
-    print("MySQL database and tables ensured")
-elif cmd == "counts":
-    cfg = get_mysql_config()
-    url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/{cfg['database']}?charset=utf8mb4"
-    engine = create_engine(url, future=True)
-    with engine.begin() as conn:
-        for table in [
-            "users",
-            "regions",
-            "keywords",
-            "user_regions",
-            "user_keywords",
-            "job_listings",
-            "job_descriptions",
-            "cached_pages",
-            "user_interactions",
-        ]:
-            try:
-                n = conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
-                print(f"{table}: {list(n)[0][0]}")
-            except Exception as e:
-                print(f"{table}: error {e}")
-else:
-    print(__doc__)
+try:
+    if cmd == "mysql-init":
+        cfg = get_mysql_config()
+        root_url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/"
+        dbname = cfg["database"]
+        root_engine = create_engine(root_url, future=True)
+        with root_engine.begin() as conn:
+            conn.execute(text(
+                f"CREATE DATABASE IF NOT EXISTS `{dbname}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"))
+        mysql_url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/{dbname}?charset=utf8mb4"
+        mysql_engine = create_engine(mysql_url, future=True)
+        from web.db import Base
+        Base.metadata.create_all(mysql_engine)
+        print("MySQL database and tables ensured")
+        sys.exit(0)
+    elif cmd == "counts":
+        cfg = get_mysql_config()
+        url = f"mysql+pymysql://{cfg['user']}:{cfg['password']}@{cfg['host']}:{cfg['port']}/{cfg['database']}?charset=utf8mb4"
+        engine = create_engine(url, future=True)
+        with engine.begin() as conn:
+            for table in [
+                "users",
+                "regions",
+                "keywords",
+                "user_regions",
+                "user_keywords",
+                "job_listings",
+                "job_descriptions",
+                "cached_pages",
+                "user_interactions",
+            ]:
+                try:
+                    n = conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
+                    print(f"{table}: {list(n)[0][0]}")
+                except Exception as e:
+                    print(f"{table}: error {e}")
+        sys.exit(0)
+    else:
+        print(__doc__)
+        sys.exit(0)
+except Exception as e:
+    print(f"Error: {e}")
+    sys.exit(1)