refactor: split main.py into modules (config, database, models, utils, auth, runtime, maintenance)
main.py was ~3000 lines with models, routes, Docker ops, maintenance all mixed. Split into 7 focused modules: - config.py: env vars and constants - database.py: SQLAlchemy engine, SessionLocal, Base, get_db - models.py: ORM models and enums - utils.py: logging, formatting, icon handling, misc helpers - auth.py: password hashing, cookies, CSRF, user dependency - runtime.py: all Docker operations, pool management, session lifecycle - maintenance.py: cleanup loop, schema bootstrap, startup logic - main.py: FastAPI app, middleware, all route handlers only
This commit is contained in:
@@ -0,0 +1,196 @@
|
||||
import datetime as dt
|
||||
import fcntl
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
|
||||
import docker
|
||||
from sqlalchemy import select
|
||||
|
||||
from config import ENABLE_STARTUP_MAINTENANCE, SESSION_IDLE_SECONDS, WEB_POOL_SIZE
|
||||
from database import Base, SessionLocal, engine
|
||||
from models import RdpSlot, Service, ServiceType, SessionModel, SessionStatus, User
|
||||
from utils import ensure_icons_dir, now_utc
|
||||
from auth import hash_password
|
||||
from runtime import (
|
||||
_rdp_slot_container_name,
|
||||
_restart_rdp_slot_bg,
|
||||
docker_client,
|
||||
ensure_schema_compatibility,
|
||||
ensure_universal_pool,
|
||||
ensure_warm_pool,
|
||||
ensure_web_pool,
|
||||
start_rdp_slot_container,
|
||||
stop_runtime_container,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("portal")
|
||||
maintenance_lock_file = None
|
||||
|
||||
|
||||
def cleanup_loop():
|
||||
while True:
|
||||
time.sleep(60)
|
||||
db = SessionLocal()
|
||||
try:
|
||||
ensure_universal_pool()
|
||||
ensure_web_pool()
|
||||
for svc in db.scalars(
|
||||
select(Service).where(
|
||||
Service.active == True,
|
||||
Service.type.in_([ServiceType.WEB, ServiceType.RDP]),
|
||||
)
|
||||
).all():
|
||||
if svc.type == ServiceType.WEB and WEB_POOL_SIZE <= 0:
|
||||
ensure_warm_pool(svc)
|
||||
cutoff = now_utc() - dt.timedelta(seconds=SESSION_IDLE_SECONDS)
|
||||
q = select(SessionModel).where(
|
||||
SessionModel.status == SessionStatus.ACTIVE,
|
||||
SessionModel.last_access_at < cutoff,
|
||||
)
|
||||
stale = db.scalars(q).all()
|
||||
rdp_slots_to_restart: list[int] = []
|
||||
for sess in stale:
|
||||
cid = sess.container_id or ""
|
||||
if cid.startswith("RDPSLOT:"):
|
||||
try:
|
||||
rdp_slots_to_restart.append(int(cid.split(":", 1)[1]))
|
||||
except Exception:
|
||||
pass
|
||||
elif cid and not (
|
||||
cid.startswith("POOL:")
|
||||
or cid.startswith("POOLIDX:")
|
||||
or cid.startswith("WEBPOOLIDX:")
|
||||
):
|
||||
stop_runtime_container(cid)
|
||||
sess.status = SessionStatus.EXPIRED
|
||||
if stale:
|
||||
db.commit()
|
||||
for slot_id in rdp_slots_to_restart:
|
||||
threading.Thread(target=_restart_rdp_slot_bg, args=(slot_id,), daemon=True).start()
|
||||
except Exception:
|
||||
db.rollback()
|
||||
logger.exception("cleanup_loop_failed")
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def bootstrap_admin():
|
||||
admin_user = os.getenv("ADMIN_USERNAME", "admin")
|
||||
admin_password = os.getenv("ADMIN_PASSWORD", "change_me")
|
||||
ttl_days = int(os.getenv("ADMIN_TTL_DAYS", "3650"))
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
existing = db.scalar(select(User).where(User.username == admin_user))
|
||||
if not existing:
|
||||
db.add(
|
||||
User(
|
||||
username=admin_user,
|
||||
password_hash=hash_password(admin_password),
|
||||
active=True,
|
||||
is_admin=True,
|
||||
expires_at=now_utc() + dt.timedelta(days=ttl_days),
|
||||
)
|
||||
)
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def try_acquire_maintenance_leader() -> bool:
|
||||
global maintenance_lock_file
|
||||
if maintenance_lock_file is not None:
|
||||
return True
|
||||
lock_file = open("/tmp/portal-maintenance.lock", "w")
|
||||
try:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
lock_file.close()
|
||||
return False
|
||||
maintenance_lock_file = lock_file
|
||||
return True
|
||||
|
||||
|
||||
def run_maintenance_service() -> None:
|
||||
logger.info("maintenance_service_bootstrap_started")
|
||||
with open("/tmp/portal-schema.lock", "w") as lock_file:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
ensure_schema_compatibility()
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
|
||||
ensure_icons_dir()
|
||||
bootstrap_admin()
|
||||
|
||||
maintenance_lock = open("/tmp/portal-maintenance.lock", "w")
|
||||
fcntl.flock(maintenance_lock.fileno(), fcntl.LOCK_EX)
|
||||
logger.info("maintenance_service_leader_acquired")
|
||||
|
||||
db = SessionLocal()
|
||||
try:
|
||||
ensure_universal_pool()
|
||||
ensure_web_pool()
|
||||
for svc in db.scalars(
|
||||
select(Service).where(
|
||||
Service.active == True,
|
||||
Service.type.in_([ServiceType.WEB, ServiceType.RDP]),
|
||||
)
|
||||
).all():
|
||||
if svc.type == ServiceType.WEB and WEB_POOL_SIZE <= 0:
|
||||
ensure_warm_pool(svc)
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
logger.info("maintenance_service_loop_started")
|
||||
cleanup_loop()
|
||||
|
||||
|
||||
def on_startup() -> None:
|
||||
with open("/tmp/portal-schema.lock", "w") as lock_file:
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
ensure_schema_compatibility()
|
||||
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
||||
ensure_icons_dir()
|
||||
bootstrap_admin()
|
||||
if not try_acquire_maintenance_leader():
|
||||
logger.info("maintenance_leader_skipped")
|
||||
return
|
||||
|
||||
if ENABLE_STARTUP_MAINTENANCE:
|
||||
db = SessionLocal()
|
||||
try:
|
||||
ensure_universal_pool()
|
||||
ensure_web_pool()
|
||||
for svc in db.scalars(
|
||||
select(Service).where(
|
||||
Service.active == True,
|
||||
Service.type.in_([ServiceType.WEB, ServiceType.RDP]),
|
||||
)
|
||||
).all():
|
||||
if svc.type == ServiceType.WEB and WEB_POOL_SIZE <= 0:
|
||||
ensure_warm_pool(svc)
|
||||
elif svc.type == ServiceType.RDP:
|
||||
slots = db.scalars(select(RdpSlot).where(RdpSlot.service_id == svc.id)).all()
|
||||
for slot in slots:
|
||||
try:
|
||||
cname = _rdp_slot_container_name(svc.slug, slot.id)
|
||||
try:
|
||||
c = docker_client().containers.get(cname)
|
||||
if c.status != "running":
|
||||
c.start()
|
||||
except docker.errors.NotFound:
|
||||
start_rdp_slot_container(slot, svc)
|
||||
slot.container_name = cname
|
||||
except Exception:
|
||||
logger.exception("startup_rdp_slot_start_failed slot_id=%s", slot.id)
|
||||
if slots:
|
||||
db.commit()
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
thread = threading.Thread(target=cleanup_loop, daemon=True)
|
||||
thread.start()
|
||||
logger.info("maintenance_leader_started")
|
||||
Reference in New Issue
Block a user