fix(maintenance): restore a backup of any size the disk has room for

The legacy restore refused the maintainer's own 1.6 GB safety backup with
"File too large. Maximum size is 100MB", a fixed cap since #153. Free
disk space now decides: the upload, its staged copy and the swap must fit,
under a 64 GiB ceiling. A request with a Content-Length that does not fit
is refused before anything is written, one without as soon as its running
size does not, with 507 insufficient-space and a message without paths.
The upload still streams and its temp file is removed on every path.

TESTING.md

@@ -2566,6 +2566,19 @@ replaced the database since the review; the plan is dropped). The job:
    prompt ids only), logs the backup by its file name, and prunes safety backups to the
    newest 10. With no backup the log and the status line say "No safety backup".
 
+The legacy restore (`POST /prompt_manager/restore`) takes a backup of any size the disks
+have room for, so a real library's safety backup restores: the upload, its staged copy
+and the swap (the upload twice in the temporary folder, once more plus the current
+database beside the database, added up on one disk, each with the 16 MiB margin) must
+fit, under a 64 GiB ceiling (`restore_max_bytes`). A request whose Content-Length does
+not fit is refused before anything is written, one without is refused as soon as its
+running size does not, both with 507 `insufficient-space` and "Not enough free disk
+space to restore a 1.6 GB backup; there is room for about …. Nothing was changed."
+(no path); past the ceiling it is the legacy 400 "File too large" with code `too-large`.
+The upload streams with `read_chunk`, which aiohttp's `client_max_size` (ComfyUI's
+`--max-upload-size`, 100 MB by default) does not cap, and its temp file is closed and
+removed on every path (`TestLargeRestores`, a 150 MB backup under a 100 MB app limit).
+
 Safety backups and restore backups are one pool: every `<db>.backup_*` file beside the
 database counts toward the same `MAX_SAFETY_BACKUPS` (10), whichever wrote it, and the
 pool is pruned after every backup a job took, also when the apply then failed

py/api/admin.py

@@ -1274,7 +1274,11 @@ class AdminRoutesMixin:
     async def restore_database(self, request):
         """Restore the prompts database from an uploaded SQLite file.
 
-        The upload is streamed to a temp file (capped at restore_max_bytes),
+        The upload is streamed to a temp file. Free disk space decides how
+        large it may be (``maintenance.restore_room``), under a 64 GiB ceiling
+        (``restore_max_bytes``): a request whose Content-Length does not fit
+        is refused before anything is written, and one without is refused as
+        soon as its running size does not. The file is
         verified with PromptModel.verify_database_file and only then copied
         over the live database, which is backed up first. The whole request
         holds the registry (records): it is refused while a scan or another
@@ -1287,10 +1291,44 @@ class AdminRoutesMixin:
             lambda: self._restore_database(request),
         )
 
+    def _restore_refusal(self, size, ceiling, room, *, known):
+        """``400 too-large`` past the ceiling, else ``507 insufficient-space``."""
+        from .maintenance import restore_space_message
+
+        if size > ceiling:
+            return web.json_response(
+                {
+                    "success": False,
+                    "code": "too-large",
+                    "error": (
+                        "File too large. Maximum size is "
+                        f"{ceiling // (1024 * 1024)}MB"
+                    ),
+                },
+                status=400,
+            )
+        return web.json_response(
+            {
+                "success": False,
+                "code": "insufficient-space",
+                "error": restore_space_message(size, room, known=known),
+            },
+            status=507,
+        )
+
     async def _restore_database(self, request):
-        max_bytes = getattr(self, "restore_max_bytes", 100 * 1024 * 1024)
+        from .maintenance import RESTORE_MAX_BYTES, restore_room
+
+        ceiling = getattr(self, "restore_max_bytes", RESTORE_MAX_BYTES)
         temp_path = None
         try:
+            # Free disk space decides how large a backup may be (#313): the
+            # upload, its staged copy and the swap must all fit.
+            room = await self._run_in_executor(restore_room, self.db.model.db_path)
+            max_bytes = min(ceiling, room)
+            length = request.content_length
+            if length is not None and length > max_bytes:
+                return self._restore_refusal(length, ceiling, room, known=True)
             reader = await request.multipart()
             field = await reader.next()
 
@@ -1316,16 +1354,7 @@ class AdminRoutesMixin:
                         break
                     size += len(chunk)
                     if size > max_bytes:
-                        return web.json_response(
-                            {
-                                "success": False,
-                                "error": (
-                                    "File too large. Maximum size is "
-                                    f"{max_bytes // (1024 * 1024)}MB"
-                                ),
-                            },
-                            status=400,
-                        )
+                        return self._restore_refusal(size, ceiling, room, known=False)
                     out.write(chunk)
 
             if size == 0:

py/api/maintenance.py

@@ -207,6 +207,57 @@ def disk_plan(db_path, *, backup, vacuum):
     return public, short, size
 
 
+# A restore upload's hard ceiling; below it, free disk space decides.
+RESTORE_MAX_BYTES = 64 * 2**30
+
+
+def size_label(size):
+    """``1.6 GB`` or ``120.0 MB``: a backup's size for a message."""
+    if size >= 2**30:
+        return f"{size / 2**30:,.1f} GB"
+    return _megabytes(size)
+
+
+def restore_room(db_path):
+    """The largest restore upload the disks have room for (blocking).
+
+    The upload is written to a temp file and staged once more in the
+    temporary folder (2x), then copied over the live database beside it,
+    whose current contents are first kept as a safety backup (1x the upload
+    for the WAL, plus the current database). Needs on one device add up,
+    each with ``DISK_MARGIN_BYTES``.
+    """
+    db_path = os.path.abspath(db_path)
+    folder = os.path.dirname(db_path)
+    current = database_bytes(db_path) if os.path.exists(db_path) else 0
+    devices = {}
+    for path, per_upload, fixed in (
+        (tempfile.gettempdir(), 2, 0),
+        (folder, 1, current),
+    ):
+        entry = devices.setdefault(_device(path), {"path": path, "per": 0, "fixed": 0})
+        entry["per"] += per_upload
+        entry["fixed"] += fixed
+    room = min(
+        (_free_bytes(e["path"]) - e["fixed"] - DISK_MARGIN_BYTES) // e["per"]
+        for e in devices.values()
+    )
+    return max(0, room)
+
+
+def restore_space_message(size, room, *, known):
+    """The truthful refusal of a restore the disks have no room for."""
+    if known:
+        return (
+            f"Not enough free disk space to restore a {size_label(size)} backup; "
+            f"there is room for about {size_label(room)}. Nothing was changed."
+        )
+    return (
+        "Not enough free disk space to restore this backup: it is larger than "
+        f"the {size_label(room)} there is room for. Nothing was changed."
+    )
+
+
 def short_space_message(short):
     """The truthful refusal for ``disk_plan``'s ``short``."""
     place, what, required, free = short

tests/test_api_backup_restore.py

@@ -47,11 +47,13 @@ def _texts(path):
 
 
 class BackupRestoreTestCase(AioHTTPTestCase):
+    app_options = {}
+
     async def get_application(self):
         self.tmpdir = tempfile.mkdtemp()
         self.db_path = os.path.join(self.tmpdir, "prompts.db")
 
-        app = web.Application()
+        app = web.Application(**self.app_options)
         routes = web.RouteTableDef()
         self.api = PromptManagerAPI(db=PromptDatabase(self.db_path))
         self.api.add_routes(routes)
@@ -327,6 +329,15 @@ class TestRestoreRoute(BackupRestoreTestCase):
         self.assertIn("too large", data["error"].lower())
         self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
 
+    async def test_restore_cap_keeps_the_legacy_shape(self):
+        self.api.restore_max_bytes = 1024
+        upload = self._write("big.db", b"SQLite format 3\x00" + b"\x00" * 8192)
+        resp = await self._restore(upload)
+        data = await resp.json()
+        self.assertEqual(resp.status, 400)
+        self.assertEqual(data["code"], "too-large")
+        self.assertNotIn(self.tmpdir, data["error"])
+
     async def test_non_multipart_request_is_rejected(self):
         resp = await self.client.post("/prompt_manager/restore", data=b"raw bytes")
         self.assertIn(resp.status, (400, 500))
@@ -336,3 +347,132 @@ class TestRestoreRoute(BackupRestoreTestCase):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+MIB = 2**20
+
+
+class TestLargeRestores(BackupRestoreTestCase):
+    """A real library's backup restores: room on disk decides, not 100 MB."""
+
+    # As ComfyUI builds its app: client_max_size is --max-upload-size (100
+    # MB by default). The restore streams with read_chunk, which it does not
+    # cap.
+    app_options = {"client_max_size": 100 * 2**20}
+
+    def _large_db_file(self, texts, pad_bytes):
+        path = self._make_db_file(texts, name="large.db")
+        conn = sqlite3.connect(path)
+        try:
+            conn.execute("CREATE TABLE padding(b BLOB)")
+            conn.execute("INSERT INTO padding VALUES (zeroblob(?))", (pad_bytes,))
+            conn.commit()
+        finally:
+            conn.close()
+        return path
+
+    async def _stream_restore(self, path, chunked=False):
+        with open(path, "rb") as fh:
+            data = fh.read()
+        form = aiohttp.FormData()
+        form.add_field(
+            "database_file",
+            data,
+            filename=os.path.basename(path),
+            content_type="application/octet-stream",
+        )
+        # chunked=False must not be passed: aiohttp 3.14 then frames the body
+        # as chunked while also sending a Content-Length.
+        options = {"chunked": True} if chunked else {}
+        return await self.client.post("/prompt_manager/restore", data=form, **options)
+
+    def _spy_uploads(self):
+        """Patch the upload's mkstemp; the list collects each temp path made."""
+        created = []
+        real_mkstemp = tempfile.mkstemp
+
+        def spy(*args, **kwargs):
+            result = real_mkstemp(*args, **kwargs)
+            created.append(result[1])
+            return result
+
+        return patch("py.api.admin.tempfile.mkstemp", spy), created
+
+    def _assert_removed(self, created):
+        for path in created:
+            for suffix in ("", "-wal", "-shm", "-journal"):
+                self.assertFalse(os.path.exists(path + suffix), path + suffix)
+
+    async def test_a_150_mb_backup_restores(self):
+        self._save("before")
+        upload = self._large_db_file(["restored"], 150 * MIB)
+        self.assertGreater(os.path.getsize(upload), 150 * MIB)
+        spy, created = self._spy_uploads()
+        with spy:
+            resp = await self._stream_restore(upload)
+        data = await resp.json()
+        self.assertEqual(resp.status, 200, data)
+        self.assertEqual(
+            [p["text"] for p in self.api.db.search_prompts()], ["restored"]
+        )
+        self.assertEqual(len(created), 1)
+        self._assert_removed(created)
+
+    async def test_a_short_disk_is_refused_up_front(self):
+        from py.api import maintenance
+
+        self._save("keep me")
+        upload = self._make_db_file(["restored"])
+        spy, created = self._spy_uploads()
+        with spy, patch.object(maintenance, "_free_bytes", return_value=MIB):
+            resp = await self._restore(upload)
+        data = await resp.json()
+        self.assertEqual(resp.status, 507)
+        self.assertEqual(data["code"], "insufficient-space")
+        self.assertTrue(
+            data["error"].startswith("Not enough free disk space to restore a "),
+            data["error"],
+        )
+        self.assertTrue(data["error"].endswith("Nothing was changed."))
+        self.assertNotIn(self.tmpdir, data["error"])
+        self.assertNotIn(tempfile.gettempdir(), data["error"])
+        self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
+        self.assertEqual(created, [])  # refused before anything was written
+
+    async def test_a_short_disk_is_refused_while_streaming_without_a_length(self):
+        from py.api import maintenance
+
+        self._save("keep me")
+        upload = self._large_db_file(["restored"], 4 * MIB)
+        spy, created = self._spy_uploads()
+        # Room for 1 MiB only, and no Content-Length to refuse up front.
+        with spy, patch.object(maintenance, "restore_room", return_value=MIB):
+            resp = await self._stream_restore(upload, chunked=True)
+        data = await resp.json()
+        self.assertEqual(resp.status, 507)
+        self.assertEqual(data["code"], "insufficient-space")
+        self.assertTrue(data["error"].endswith("Nothing was changed."))
+        self.assertEqual([p["text"] for p in self.api.db.search_prompts()], ["keep me"])
+        self.assertEqual(len(created), 1)
+        self._assert_removed(created)
+
+
+class TestRestoreRoom(unittest.TestCase):
+    def test_room_counts_the_upload_its_staged_copy_and_the_swap(self):
+        from py.api import maintenance
+
+        tmp = tempfile.mkdtemp()
+        try:
+            db_path = os.path.join(tmp, "prompts.db")
+            free = 10 * 2**30
+            with (
+                patch.object(maintenance, "_free_bytes", return_value=free),
+                patch.object(maintenance, "_device", return_value="one disk"),
+            ):
+                room = maintenance.restore_room(db_path)
+            # Upload + staged copy + WAL growth on one disk, no database yet.
+            self.assertEqual(room, (free - maintenance.DISK_MARGIN_BYTES) // 3)
+            self.assertEqual(maintenance.size_label(int(1.6 * 2**30)), "1.6 GB")
+            self.assertEqual(maintenance.size_label(120 * 2**20), "120.0 MB")
+        finally:
+            os.rmdir(tmp)