Refactored DB handlers

app/server.py

@@ -8,40 +8,28 @@ from quart_cors import cors
 from routers.systems import systems_bp
 from routers.nodes import nodes_bp, register_client, unregister_client
 from routers.bot import bot_bp
-from routers.auth import auth_bp # ONLY import auth_bp, not jwt instance from auth.py
-from internal.db_wrappers import SystemDbController, DiscordIdDbController
-from internal.auth_wrappers import UserDbController
-# Import the JWTManager instance and its configuration function
-from config.jwt_config import jwt, configure_jwt # Import the actual jwt instance and the config function
+from routers.auth import auth_bp
+from internal.db_wrappers import SystemDbController, DiscordIdDbController #
+from internal.auth_wrappers import UserDbController #
+from config.jwt_config import jwt, configure_jwt
 
 
 # --- WebSocket Server Components ---
-# Dictionary to store active clients: {client_id: websocket}
 active_clients = {}
 
-
 async def websocket_server_handler(websocket):
-    """Handles incoming WebSocket connections and messages from clients."""
     client_id = None
     try:
-        # Handshake: Receive the first message which should contain the client ID
         handshake_message = await websocket.recv()
         handshake_data = json.loads(handshake_message)
-
         if handshake_data.get("type") == "handshake" and "id" in handshake_data:
             client_id = handshake_data["id"]
             await register_client(websocket, client_id)
-            await websocket.send(json.dumps({"type": "handshake_ack", "status": "success"})) # Acknowledge handshake
-
-            # Keep the connection alive and listen for potential messages from the client
-            # (Though in this server-commanded model, clients might not send much)
-            # We primarily wait for the client to close the connection
+            await websocket.send(json.dumps({"type": "handshake_ack", "status": "success"}))
             await websocket.wait_closed()
-
         else:
             print(f"Received invalid handshake from {websocket.remote_address}. Closing connection.")
             await websocket.close()
-
     except websockets.exceptions.ConnectionClosedError:
         print(f"Client connection closed unexpectedly for {client_id}.")
     except json.JSONDecodeError:
@@ -56,62 +44,70 @@ async def websocket_server_handler(websocket):
 app = Quart(__name__)
 app = cors(app, allow_origin="*")
 
-# Store the websocket server instance
 websocket_server_instance = None
-
-# Make active_clients accessible via the app instance.
 app.active_clients = active_clients
 
 # Create and attach the DB wrappers
-app.sys_db_h = SystemDbController()
-app.d_id_db_h = DiscordIdDbController()
-app.user_db_h = UserDbController()
+app.sys_db_h = SystemDbController() #
+app.d_id_db_h = DiscordIdDbController() #
+app.user_db_h = UserDbController() #
 
-# Configure JWT settings and initialize the JWTManager instance with the app
 configure_jwt(app)
-jwt.init_app(app) # Crucial: This initializes the global 'jwt' instance with your app
+jwt.init_app(app)
 
 
 @app.before_serving
-async def startup_websocket_server():
-    """Starts the WebSocket server when the Quart app starts."""
+async def startup_tasks(): # Combined startup logic
+    """Starts the WebSocket server and prepares other resources."""
     global websocket_server_instance
     websocket_server_address = "0.0.0.0"
     websocket_server_port = 8765
-
-    # Start the WebSocket server task
     websocket_server_instance = await websockets.serve(
         websocket_server_handler,
         websocket_server_address,
         websocket_server_port
     )
     print(f"WebSocket server started on ws://{websocket_server_address}:{websocket_server_port}")
+    # Database connections are now established on first use by MongoHandler's __aenter__/connect
+    # No explicit connect calls needed here unless desired for early failure detection.
+    print("Application startup complete. DB connections will be initialized on first use.")
 
 @app.after_serving
-async def shutdown_websocket_server():
-    """Shuts down the WebSocket server when the Quart app stops."""
+async def shutdown_tasks(): # Combined shutdown logic
+    """Shuts down services and closes connections."""
     global websocket_server_instance
     if websocket_server_instance:
         websocket_server_instance.close()
         await websocket_server_instance.wait_closed()
         print("WebSocket server shut down.")
 
+    # Close database connections
+    if hasattr(app, 'user_db_h') and app.user_db_h:
+        print("Closing User DB connection...")
+        await app.user_db_h.close_db_connection() #
+    if hasattr(app, 'sys_db_h') and app.sys_db_h:
+        print("Closing System DB connection...")
+        await app.sys_db_h.close_db_connection() #
+    if hasattr(app, 'd_id_db_h') and app.d_id_db_h:
+        print("Closing Discord ID DB connection...")
+        await app.d_id_db_h.close_db_connection() #
+    print("All database connections have been signaled to close.")
+
 
 app.register_blueprint(systems_bp, url_prefix="/systems")
 app.register_blueprint(nodes_bp, url_prefix="/nodes")
 app.register_blueprint(bot_bp, url_prefix="/bots")
-app.register_blueprint(auth_bp, url_prefix="/auth") # Register the auth blueprint
+app.register_blueprint(auth_bp, url_prefix="/auth")
 
 @app.route('/')
 async def index():
     return "Welcome to the Radio App Server API!"
 
-# --- Main Execution ---
 if __name__ == "__main__":
     print("Starting Quart API server...")
     app.run(
         host="0.0.0.0",
         port=5000,
-        debug=False # Set to True for development
+        debug=False
     )
     print("Quart API server stopped.")