refactor: extract _init_schema and _startup_pools to remove duplication in maintenance.py

app/maintenance.py

@@ -120,49 +120,56 @@ def try_acquire_maintenance_leader() -> bool:
     return True
 
 
-def run_maintenance_service() -> None:
-    logger.info("maintenance_service_bootstrap_started")
+def _init_schema() -> 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()
 
+
+def _startup_pools(db) -> None:
+    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()
+
+
+def run_maintenance_service() -> None:
+    logger.info("maintenance_service_bootstrap_started")
+    _init_schema()
+
     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)
-            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()
+        _startup_pools(db)
     finally:
         db.close()
 
@@ -171,13 +178,7 @@ def run_maintenance_service() -> None:
 
 
 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()
+    _init_schema()
     if not try_acquire_maintenance_leader():
         logger.info("maintenance_leader_skipped")
         return
@@ -185,32 +186,7 @@ def on_startup() -> None:
     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()
+            _startup_pools(db)
         finally:
             db.close()