Fix for #107 - Main functions not running when using uvicorn

wg-manager-backend/database/util.py

@@ -1,4 +1,3 @@
-import contextlib
 import os
 
 import alembic.command
@@ -27,6 +26,7 @@ def perform_migrations():
 def setup_initial_database():
     if not database_exists(engine.url):
         logger.info("Database does not exists. Creating initial database...")
+
         # Create database from metadata
         Base.metadata.create_all(engine)
         logger.info("Database creation done!")
@@ -34,25 +34,23 @@ def setup_initial_database():
     # Create default user
     _db: Session = SessionLocal()
 
-    admin_exists = (
-                       _db.query(models.User.id)
-                           .filter_by(role="admin")
-                           .first()
-                   ) is not None
+    # Retrieve user with admin role
+    admin_exists = _db.query(models.User.id).filter_by(role="admin").first() is not None
 
     if not admin_exists:
         logger.info("Admin user does not exists. Creating with env variables ADMIN_USERNAME, ADMIN_PASSWORD")
-        ADMIN_USERNAME = os.getenv("ADMIN_USERNAME")
-        ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD")
+        env_admin_username = os.getenv("ADMIN_USERNAME")
+        env_admin_password = os.getenv("ADMIN_PASSWORD")
 
-        if not ADMIN_USERNAME:
+        if not env_admin_username:
             raise RuntimeError("Database does not exist and the environment variable ADMIN_USERNAME is set")
-        if not ADMIN_PASSWORD:
+
+        if not env_admin_password:
             raise RuntimeError("Database does not exist and the environment variable ADMIN_PASSWORD is set")
 
         _db.merge(models.User(
-            username=ADMIN_USERNAME,
-            password=middleware.get_password_hash(ADMIN_PASSWORD),
+            username=env_admin_username,
+            password=middleware.get_password_hash(env_admin_password),
             full_name="Admin",
             role="admin",
             email=""

wg-manager-backend/main.py

@@ -67,17 +67,7 @@ def root():
 app.mount("/", StaticFiles(directory=pkg_resources.resource_filename(__name__, 'build')), name="static")
 
 
-@app.on_event("startup")
-async def startup():
-    pass
-
-
-@app.on_event("shutdown")
-async def shutdown():
-    pass
-
-
-if __name__ == "__main__":
+def main():
     # Sleep the wait timer.
     time.sleep(const.INIT_SLEEP)
 
@@ -90,6 +80,22 @@ if __name__ == "__main__":
     # Configure wireguard
     script.wireguard_startup.setup_on_start()
 
+
+@app.on_event("startup")
+async def startup():
+    if __name__ != "__main__":
+        main()
+
+
+@app.on_event("shutdown")
+async def shutdown():
+    pass
+
+
+if __name__ == "__main__":
+
+    main()
+
     run_uvicorn_loguru(
         uvicorn.Config(
             "__main__:app",

wg-manager-backend/migrations/versions/4ac3e58519eb_base.py

@@ -8,7 +8,6 @@ Create Date: 2021-03-13 20:29:10.062757
 from alembic import op
 import sqlalchemy as sa
 
-
 # revision identifiers, used by Alembic.
 from sqlalchemy.exc import OperationalError
 
@@ -17,6 +16,7 @@ down_revision = None
 branch_labels = None
 depends_on = None
 
+
 def ignore_duplicate(fn):
     try:
         fn()
@@ -24,18 +24,20 @@ def ignore_duplicate(fn):
         if "duplicate" in str(e):
             pass
 
+
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
     try:
         op.create_table('api_keys',
-        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
-        sa.Column('key', sa.String(), nullable=True),
-        sa.Column('user_id', sa.Integer(), nullable=True),
-        sa.Column('created_date', sa.DateTime(), nullable=True),
-        sa.ForeignKeyConstraint(['user_id'], ['users.id'], name='fk_user_api_key_user_id', onupdate='CASCADE', ondelete='CASCADE'),
-        sa.PrimaryKeyConstraint('id'),
-        sa.UniqueConstraint('key')
-        )
+                        sa.Column('id', sa.Integer(), autoincrement=True, nullable=False),
+                        sa.Column('key', sa.String(), nullable=True),
+                        sa.Column('user_id', sa.Integer(), nullable=True),
+                        sa.Column('created_date', sa.DateTime(), nullable=True),
+                        sa.ForeignKeyConstraint(['user_id'], ['users.id'], name='fk_user_api_key_user_id',
+                                                onupdate='CASCADE', ondelete='CASCADE'),
+                        sa.PrimaryKeyConstraint('id'),
+                        sa.UniqueConstraint('key')
+                        )
     except OperationalError:
         pass
 
@@ -51,7 +53,8 @@ def upgrade():
         batch_op.drop_constraint("fk_peer_server_id_server", type_="foreignkey")
 
     with op.batch_alter_table('peer', schema=None) as batch_op:
-        batch_op.create_foreign_key('fk_peer_server_id_server', 'server', ['server_id'], ['id'], onupdate='CASCADE', ondelete='CASCADE')
+        batch_op.create_foreign_key('fk_peer_server_id_server', 'server', ['server_id'], ['id'], onupdate='CASCADE',
+                                    ondelete='CASCADE')
 
     ignore_duplicate(lambda: op.add_column('peer', sa.Column('configuration', sa.Text(), nullable=True)))
     ignore_duplicate(lambda: op.add_column('peer', sa.Column('keep_alive', sa.Integer(), nullable=True)))
@@ -60,11 +63,9 @@ def upgrade():
     ignore_duplicate(lambda: op.add_column('peer', sa.Column('shared_key', sa.Text(), nullable=True)))
     ignore_duplicate(lambda: op.add_column('peer', sa.Column('v6_address', sa.String(), nullable=True)))
 
-
-    #op.drop_constraint(None, 'peer', type_='foreignkey')
+    # op.drop_constraint(None, 'peer', type_='foreignkey')
     #
-    #op.drop_column('peer', 'server')
-
+    # op.drop_column('peer', 'server')
 
     try:
         with op.batch_alter_table('peer', schema=None) as batch_op:
@@ -80,8 +81,7 @@ def upgrade():
     ignore_duplicate(lambda: op.add_column('server', sa.Column('subnet', sa.Integer(), nullable=False)))
     ignore_duplicate(lambda: op.add_column('server', sa.Column('v6_address', sa.String(), nullable=True)))
     ignore_duplicate(lambda: op.add_column('server', sa.Column('v6_subnet', sa.Integer(), nullable=False)))
-    #op.create_unique_constraint(None, 'server', ['v6_address'])
-
+    # op.create_unique_constraint(None, 'server', ['v6_address'])
 
     try:
         with op.batch_alter_table('server', schema=None) as batch_op: