fix: export every prompt by paging instead of truncating at 10000
/prompt_manager/export fetched a single page of 10000 prompts and silently dropped the rest. It now walks the table in EXPORT_PAGE_SIZE pages (1000) inside one executor job until a short page is returned.
This commit is contained in:
+15
-1
@@ -11,6 +11,10 @@ from aiohttp import web
|
||||
MAX_PAGE_LIMIT = 500
|
||||
MAX_PAGE_OFFSET = 10_000_000
|
||||
|
||||
# The export walks the whole table in pages of this size instead of asking
|
||||
# for one giant result set (and silently truncating large libraries).
|
||||
EXPORT_PAGE_SIZE = 1000
|
||||
|
||||
|
||||
def parse_page_params(query, default_limit=50, max_limit=MAX_PAGE_LIMIT):
|
||||
"""Return ``(limit, offset)`` from a query mapping, clamped to safe bounds.
|
||||
@@ -1099,10 +1103,20 @@ class PromptRoutesMixin:
|
||||
status=500,
|
||||
)
|
||||
|
||||
def _all_prompts_for_export(self):
|
||||
"""Every prompt, fetched in EXPORT_PAGE_SIZE pages (blocking)."""
|
||||
prompts, offset = [], 0
|
||||
while True:
|
||||
page = self.db.search_prompts(limit=EXPORT_PAGE_SIZE, offset=offset)
|
||||
prompts.extend(page)
|
||||
if len(page) < EXPORT_PAGE_SIZE:
|
||||
return prompts
|
||||
offset += EXPORT_PAGE_SIZE
|
||||
|
||||
async def export_prompts(self, request):
|
||||
"""Export all prompts to JSON."""
|
||||
try:
|
||||
prompts = await self._run_in_executor(self.db.search_prompts, limit=10000)
|
||||
prompts = await self._run_in_executor(self._all_prompts_for_export)
|
||||
|
||||
export_data = {
|
||||
"export_date": datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
|
||||
@@ -7,7 +7,7 @@ import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
@@ -19,6 +19,7 @@ from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
|
||||
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from py.api import PromptManagerAPI # noqa: E402
|
||||
import py.api.prompts as prompts_module # noqa: E402
|
||||
from py.api.prompts import ( # noqa: E402
|
||||
MAX_PAGE_LIMIT,
|
||||
MAX_PAGE_OFFSET,
|
||||
@@ -504,6 +505,30 @@ class TestTagInputValidation(RouteCoverageCase):
|
||||
await self._assert_all_400("bad\x00tag")
|
||||
|
||||
|
||||
class TestExportPagesThroughAllPrompts(RouteCoverageCase):
|
||||
"""GET /prompt_manager/export never truncates: it pages until exhausted."""
|
||||
|
||||
async def test_export_is_complete_beyond_one_page(self):
|
||||
ids = {self._save_prompt(f"export {i}") for i in range(5)}
|
||||
real_search = self.api.db.search_prompts
|
||||
calls = []
|
||||
|
||||
def spy(*args, **kwargs):
|
||||
calls.append((kwargs.get("limit"), kwargs.get("offset")))
|
||||
return real_search(*args, **kwargs)
|
||||
|
||||
self.api.db.search_prompts = spy
|
||||
|
||||
with patch.object(prompts_module, "EXPORT_PAGE_SIZE", 2):
|
||||
resp = await self.client.request("GET", "/prompt_manager/export")
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = json.loads(await resp.text())
|
||||
self.assertEqual(data["total_prompts"], 5)
|
||||
self.assertEqual({p["id"] for p in data["prompts"]}, ids)
|
||||
self.assertEqual(calls, [(2, 0), (2, 2), (2, 4)])
|
||||
|
||||
|
||||
class TestTagMutationRoutes(RouteCoverageCase):
|
||||
|
||||
async def test_rename_rejects_bad_json_and_empty_name(self):
|
||||
|
||||
Reference in New Issue
Block a user