Improve memory management

.pre-commit-config.yaml

@@ -18,7 +18,7 @@ repos:
     hooks:
       - id: pytest
         name: pytest
-        entry: poetry run pytest tests
+        entry: env AIOCACHE_DISABLE=1 poetry run pytest tests
         language: system
         types: [python]
         pass_filenames: false

conftest.py

@@ -10,7 +10,6 @@
 
 from goosebit import app
 from goosebit.models import UpdateModeEnum, UpdateStateEnum
-from goosebit.updater.manager import reset_update_manager
 
 # Configure logging
 logging.basicConfig(level=logging.WARN)
@@ -36,9 +35,7 @@ async def test_app():
 
 @pytest_asyncio.fixture(scope="module")
 async def async_client(test_app):
-    async with AsyncClient(
-        transport=ASGITransport(app=test_app), base_url="http://test"
-    ) as client:
+    async with AsyncClient(transport=ASGITransport(app=test_app), base_url="http://test") as client:
         login_data = {"username": "admin@goosebit.local", "password": "admin"}
         response = await client.post("/login", data=login_data, follow_redirects=True)
         assert response.status_code == 200
@@ -48,8 +45,6 @@ async def async_client(test_app):
 
 @pytest_asyncio.fixture(scope="function")
 async def db():
-    reset_update_manager()
-
     await Tortoise.init(config=TORTOISE_CONF)
     await Tortoise.generate_schemas()
     yield

goosebit/__init__.py

@@ -52,9 +52,7 @@ async def login_ui(request: Request):
 
 
 @app.post("/login", include_in_schema=False, dependencies=[Depends(authenticate_user)])
-async def login(
-    request: Request, form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
-):
+async def login(request: Request, form_data: Annotated[OAuth2PasswordRequestForm, Depends()]):
     resp = RedirectResponse(request.url_for("ui_root"), status_code=302)
     resp.set_cookie(key="session_id", value=create_session(form_data.username))
     return resp

goosebit/api/devices.py

@@ -10,7 +10,7 @@
 from goosebit.auth import validate_user_permissions
 from goosebit.models import Device, Firmware, UpdateModeEnum, UpdateStateEnum
 from goosebit.permissions import Permissions
-from goosebit.updater.manager import delete_device, get_update_manager
+from goosebit.updater.manager import delete_devices, get_update_manager
 
 router = APIRouter(prefix="/devices")
 
@@ -42,22 +42,18 @@ async def parse(device: Device) -> dict:
             "uuid": device.uuid,
             "name": device.name,
             "fw": device.fw_version,
-            "fw_version": (
-                target_firmware.version if target_firmware is not None else None
-            ),
+            "fw_version": (target_firmware.version if target_firmware is not None else None),
             "hw_model": device.hardware.model,
             "hw_revision": device.hardware.revision,
             "feed": device.feed,
             "flavor": device.flavor,
             "progress": device.progress,
             "state": str(device.last_state),
             "update_mode": str(device.update_mode),
-            "force_update": manager.force_update,
+            "force_update": device.force_update,
             "last_ip": device.last_ip,
             "last_seen": last_seen,
-            "online": (
-                last_seen < manager.poll_seconds if last_seen is not None else None
-            ),
+            "online": (last_seen < manager.poll_seconds if last_seen is not None else None),
         }
 
     total_records = await Device.all().count()
@@ -73,9 +69,7 @@ class UpdateDevicesModel(BaseModel):
 
 @router.post(
     "/update",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.DEVICE.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.DEVICE.WRITE])],
 )
 async def devices_update(_: Request, config: UpdateDevicesModel) -> dict:
     for uuid in config.devices:
@@ -101,14 +95,12 @@ class ForceUpdateModel(BaseModel):
 
 @router.post(
     "/force_update",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.DEVICE.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.DEVICE.WRITE])],
 )
 async def devices_force_update(_: Request, config: ForceUpdateModel) -> dict:
     for uuid in config.devices:
         updater = await get_update_manager(uuid)
-        updater.force_update = True
+        await updater.update_force_update(True)
     return {"success": True}
 
 
@@ -130,11 +122,8 @@ class DeleteModel(BaseModel):
 
 @router.post(
     "/delete",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.DEVICE.DELETE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.DEVICE.DELETE])],
 )
 async def devices_delete(_: Request, config: DeleteModel) -> dict:
-    for uuid in config.devices:
-        await delete_device(uuid)
+    await delete_devices(config.devices)
     return {"success": True}

goosebit/api/firmware.py

@@ -14,9 +14,7 @@
 
 @router.get(
     "/all",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.READ])],
 )
 async def firmware_get_all(
     request: Request,
@@ -42,9 +40,7 @@ async def parse(f):
 
 @router.post(
     "/delete",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.DELETE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.DELETE])],
 )
 async def firmware_delete(_: Request, files: list[int] = Body()) -> dict:
     success = False

goosebit/api/rollouts.py

@@ -13,19 +13,13 @@
 
 @router.get(
     "/all",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.READ])],
 )
 async def rollouts_get_all(request: Request) -> dict[str, int | list[dict]]:
     query = Rollout.all().prefetch_related("firmware")
 
     def search_filter(search_value):
-        return (
-            Q(name__icontains=search_value)
-            | Q(feed__icontains=search_value)
-            | Q(flavor__icontains=search_value)
-        )
+        return Q(name__icontains=search_value) | Q(feed__icontains=search_value) | Q(flavor__icontains=search_value)
 
     async def parse(rollout: Rollout) -> dict:
         return {
@@ -53,9 +47,7 @@ class CreateRolloutsModel(BaseModel):
 
 @router.post(
     "/",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.WRITE])],
 )
 async def rollouts_create(_: Request, rollout: CreateRolloutsModel) -> dict:
     await Rollout.create(
@@ -74,9 +66,7 @@ class UpdateRolloutsModel(BaseModel):
 
 @router.post(
     "/update",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.WRITE])],
 )
 async def rollouts_update(_: Request, rollouts: UpdateRolloutsModel) -> dict:
     await Rollout.filter(id__in=rollouts.ids).update(paused=rollouts.paused)
@@ -89,9 +79,7 @@ class DeleteRolloutsModel(BaseModel):
 
 @router.post(
     "/delete",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.DELETE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.DELETE])],
 )
 async def rollouts_delete(_: Request, rollouts: DeleteRolloutsModel) -> dict:
     await Rollout.filter(id__in=rollouts.ids).delete()

goosebit/api/routes.py

@@ -4,9 +4,7 @@
 from goosebit.auth import authenticate_api_session
 
 # main router that requires authentication
-main_router = APIRouter(
-    prefix="/api", dependencies=[Depends(authenticate_api_session)], tags=["api"]
-)
+main_router = APIRouter(prefix="/api", dependencies=[Depends(authenticate_api_session)], tags=["api"])
 main_router.include_router(firmware.router)
 main_router.include_router(devices.router)
 main_router.include_router(rollouts.router)

goosebit/auth/__init__.py

@@ -37,9 +37,7 @@ async def authenticate_user(request: Request):
 
 
 def create_session(username: str) -> str:
-    return jwt.encode(
-        header={"alg": "HS256"}, claims={"username": username}, key=SECRET
-    )
+    return jwt.encode(header={"alg": "HS256"}, claims={"username": username}, key=SECRET)
 
 
 def authenticate_session(request: Request):

goosebit/models.py

@@ -53,23 +53,21 @@ class Tag(Model):
 class Device(Model):
     uuid = fields.CharField(max_length=255, primary_key=True)
     name = fields.CharField(max_length=255, null=True)
-    assigned_firmware = fields.ForeignKeyField(
-        "models.Firmware", related_name="assigned_devices", null=True
-    )
+    assigned_firmware = fields.ForeignKeyField("models.Firmware", related_name="assigned_devices", null=True)
+    force_update = fields.BooleanField(default=False)
     fw_version = fields.CharField(max_length=255, null=True)
     hardware = fields.ForeignKeyField("models.Hardware", related_name="devices")
     feed = fields.CharField(max_length=255, default="default")
     flavor = fields.CharField(max_length=255, default="default")
     update_mode = fields.IntEnumField(UpdateModeEnum, default=UpdateModeEnum.ROLLOUT)
     last_state = fields.IntEnumField(UpdateStateEnum, default=UpdateStateEnum.UNKNOWN)
     progress = fields.IntField(null=True)
+    log_complete = fields.BooleanField(default=False)
     last_log = fields.TextField(null=True)
     last_seen = fields.BigIntField(null=True)
     last_ip = fields.CharField(max_length=15, null=True)
     last_ipv6 = fields.CharField(max_length=40, null=True)
-    tags = fields.ManyToManyField(
-        "models.Tag", related_name="devices", through="device_tags"
-    )
+    tags = fields.ManyToManyField("models.Tag", related_name="devices", through="device_tags")
 
 
 class Rollout(Model):

goosebit/realtime/logs.py

@@ -18,17 +18,16 @@ class RealtimeLogModel(BaseModel):
 
 @router.websocket(
     "/{dev_id}",
-    dependencies=[
-        Security(validate_ws_user_permissions, scopes=[Permissions.HOME.READ])
-    ],
+    dependencies=[Security(validate_ws_user_permissions, scopes=[Permissions.HOME.READ])],
 )
 async def device_logs(websocket: WebSocket, dev_id: str):
     await websocket.accept()
 
     manager = await get_update_manager(dev_id)
+    device = await manager.get_device()
 
     async def callback(log_update):
-        data = RealtimeLogModel(log=log_update, progress=manager.device.progress)
+        data = RealtimeLogModel(log=log_update, progress=device.progress)
         if log_update is None:
             data.clear = True
             data.log = ""

goosebit/telemetry/prometheus.py

@@ -3,9 +3,7 @@
 
 from goosebit import settings
 
-PROMETHEUS_PORT = (
-    settings.config.get("metrics", {}).get("prometheus", {}).get("port", 9090)
-)
+PROMETHEUS_PORT = settings.config.get("metrics", {}).get("prometheus", {}).get("port", 9090)
 
 # separate file to enable it as a feature later.
 reader = PrometheusMetricReader()

goosebit/ui/routes.py

@@ -13,9 +13,7 @@
 
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")
 
-router = APIRouter(
-    prefix="/ui", dependencies=[Depends(authenticate_session)], include_in_schema=False
-)
+router = APIRouter(prefix="/ui", dependencies=[Depends(authenticate_session)], include_in_schema=False)
 
 
 @router.get("/")
@@ -25,21 +23,15 @@ async def ui_root(request: Request):
 
 @router.get(
     "/firmware",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.READ])],
 )
 async def firmware_ui(request: Request):
-    return templates.TemplateResponse(
-        request, "firmware.html", context={"title": "Firmware"}
-    )
+    return templates.TemplateResponse(request, "firmware.html", context={"title": "Firmware"})
 
 
 @router.post(
     "/upload/local",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.WRITE])],
 )
 async def upload_update_local(
     request: Request,
@@ -67,9 +59,7 @@ async def upload_update_local(
 
 @router.post(
     "/upload/remote",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.WRITE])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.FIRMWARE.WRITE])],
 )
 async def upload_update_remote(request: Request, url: str = Form(...)):
     firmware = await Firmware.get_or_none(uri=url)
@@ -89,35 +79,23 @@ async def home_ui(request: Request):
 
 @router.get(
     "/devices",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.DEVICE.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.DEVICE.READ])],
 )
 async def devices_ui(request: Request):
-    return templates.TemplateResponse(
-        request, "devices.html", context={"title": "Devices"}
-    )
+    return templates.TemplateResponse(request, "devices.html", context={"title": "Devices"})
 
 
 @router.get(
     "/rollouts",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.ROLLOUT.READ])],
 )
 async def rollouts_ui(request: Request):
-    return templates.TemplateResponse(
-        request, "rollouts.html", context={"title": "Rollouts"}
-    )
+    return templates.TemplateResponse(request, "rollouts.html", context={"title": "Rollouts"})
 
 
 @router.get(
     "/logs/{dev_id}",
-    dependencies=[
-        Security(validate_user_permissions, scopes=[Permissions.DEVICE.READ])
-    ],
+    dependencies=[Security(validate_user_permissions, scopes=[Permissions.DEVICE.READ])],
 )
 async def logs_ui(request: Request, dev_id: str):
-    return templates.TemplateResponse(
-        request, "logs.html", context={"title": "Log", "device": dev_id}
-    )
+    return templates.TemplateResponse(request, "logs.html", context={"title": "Log", "device": dev_id})