Files
drb-core-server/app/server.py
2025-05-26 02:36:21 -04:00

113 lines
4.1 KiB
Python

# server.py
import asyncio
import websockets
import json
import uuid
from quart import Quart, jsonify, request
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
from internal.db_wrappers import SystemDbController, DiscordIdDbController #
from internal.auth_wrappers import UserDbController #
from config.jwt_config import jwt, configure_jwt
# --- WebSocket Server Components ---
active_clients = {}
async def websocket_server_handler(websocket):
client_id = None
try:
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"}))
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:
print(f"Received invalid JSON from {client_id or 'an unknown client'}.")
except Exception as e:
print(f"An error occurred with client {client_id}: {e}")
finally:
if client_id:
await unregister_client(client_id)
# --- Quart API Components ---
app = Quart(__name__)
app = cors(app, allow_origin="*")
websocket_server_instance = None
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() #
configure_jwt(app)
jwt.init_app(app)
@app.before_serving
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
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_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")
@app.route('/')
async def index():
return "Welcome to the Radio App Server API!"
if __name__ == "__main__":
print("Starting Quart API server...")
app.run(
host="0.0.0.0",
port=5000,
debug=False
)
print("Quart API server stopped.")