持久化账户生成历史
This commit is contained in:
magicwang1111
@@ -6,3 +6,4 @@ whitelist.txt
|
||||
blacklist.txt
|
||||
.env
|
||||
secret_key.txt
|
||||
history.sqlite3*
|
||||
|
||||
@@ -31,6 +31,7 @@ Edit `config.json` before starting ComfyUI.
|
||||
"secret_key_env": "SECRET_KEY",
|
||||
"secret_key_file": "secret_key.txt",
|
||||
"users_db": "users_db.json",
|
||||
"history_db": "history.sqlite3",
|
||||
"access_token_expiration_hours": 8760,
|
||||
"max_access_token_expiration_hours": 8760,
|
||||
"log": "account_manager.log",
|
||||
@@ -47,6 +48,10 @@ Edit `config.json` before starting ComfyUI.
|
||||
|
||||
`separate_users` enables account-specific inputs, outputs, queue history, and asset visibility. The older misspelled `seperate_users` key is still accepted for compatibility, but new installs should use `separate_users`.
|
||||
|
||||
Completed generation history is persisted in `history.sqlite3` and restored when ComfyUI restarts. The history database keeps the same 10,000-item limit as ComfyUI's in-memory history.
|
||||
|
||||
ComfyUI's optional Assets panel is separate from queue history and requires starting ComfyUI with `--enable-assets`.
|
||||
|
||||
## API Access
|
||||
|
||||
Authenticated API calls can include either:
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ sanitizer = Sanitizer()
|
||||
ip_filter = IPFilter(WHITELIST, BLACKLIST)
|
||||
timeout = Timeout(ip_filter, BLACKLIST_AFTER_ATTEMPTS)
|
||||
users_db = UsersDB(USERS_FILE)
|
||||
access_control = AccessControl(users_db, instance)
|
||||
access_control = AccessControl(users_db, instance, HISTORY_FILE)
|
||||
jwt_auth = JWTAuth(
|
||||
users_db, access_control, logger, SECRET_KEY, TOKEN_EXPIRE_MINUTES, TOKEN_ALGORITHM
|
||||
)
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
"secret_key_env": "SECRET_KEY",
|
||||
"secret_key_file": "secret_key.txt",
|
||||
"users_db": "users_db.json",
|
||||
"history_db": "history.sqlite3",
|
||||
"access_token_expiration_hours": 8760,
|
||||
"max_access_token_expiration_hours": 8760,
|
||||
"log": "account_manager.log",
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
import importlib.util
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
module_path = Path(__file__).parents[1] / "utils" / "history_store.py"
|
||||
spec = importlib.util.spec_from_file_location("history_store", module_path)
|
||||
history_store = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(history_store)
|
||||
HistoryStore = history_store.HistoryStore
|
||||
|
||||
|
||||
class HistoryStoreTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp_dir = tempfile.TemporaryDirectory()
|
||||
self.database = os.path.join(self.temp_dir.name, "history.sqlite3")
|
||||
self.store = HistoryStore(self.database)
|
||||
|
||||
def tearDown(self):
|
||||
self.temp_dir.cleanup()
|
||||
|
||||
@staticmethod
|
||||
def item(owner_id, value):
|
||||
return {
|
||||
"prompt": [0, f"prompt-{value}", {}, {}, []],
|
||||
"outputs": {"1": {"images": [{"filename": f"{value}.png"}]}},
|
||||
"status": {"status_str": "success", "completed": True, "messages": []},
|
||||
"user_id": owner_id,
|
||||
}
|
||||
|
||||
def test_history_survives_new_store_instance(self):
|
||||
self.store.save("one", self.item("user-a", 1), 10)
|
||||
|
||||
restored = HistoryStore(self.database).load(10)
|
||||
|
||||
self.assertEqual(["one"], list(restored))
|
||||
self.assertEqual("user-a", restored["one"]["user_id"])
|
||||
self.assertEqual("1.png", restored["one"]["outputs"]["1"]["images"][0]["filename"])
|
||||
|
||||
def test_limit_and_owner_delete(self):
|
||||
self.store.save("one", self.item("user-a", 1), 2)
|
||||
self.store.save("two", self.item("user-b", 2), 2)
|
||||
self.store.save("three", self.item("user-a", 3), 2)
|
||||
|
||||
self.assertEqual(["two", "three"], list(self.store.load(10)))
|
||||
|
||||
self.store.delete_owner("user-a")
|
||||
|
||||
self.assertEqual(["two"], list(self.store.load(10)))
|
||||
|
||||
def test_delete_and_clear(self):
|
||||
self.store.save("one", self.item("user-a", 1), 10)
|
||||
self.store.save("two", self.item("user-b", 2), 10)
|
||||
|
||||
self.store.delete("one")
|
||||
self.assertEqual(["two"], list(self.store.load(10)))
|
||||
|
||||
self.store.clear()
|
||||
self.assertEqual({}, self.store.load(10))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+57
-2
@@ -3,6 +3,7 @@ import sys
|
||||
import heapq
|
||||
import copy
|
||||
import contextvars
|
||||
import logging
|
||||
import mimetypes
|
||||
from datetime import datetime
|
||||
from aiohttp import web
|
||||
@@ -13,10 +14,16 @@ from server import PromptServer
|
||||
from execution import PromptQueue, MAXIMUM_HISTORY_SIZE
|
||||
|
||||
from .users_db import UsersDB
|
||||
from .history_store import HistoryStore
|
||||
|
||||
|
||||
logger = logging.getLogger("ComfyUI-Account-Manager")
|
||||
|
||||
|
||||
class AccessControl:
|
||||
def __init__(self, users_db: UsersDB, server: PromptServer):
|
||||
def __init__(
|
||||
self, users_db: UsersDB, server: PromptServer, history_file: str = None
|
||||
):
|
||||
self.users_db = users_db
|
||||
self.server = server
|
||||
|
||||
@@ -29,6 +36,12 @@ class AccessControl:
|
||||
self.__get_save_image_path = folder_paths.get_save_image_path
|
||||
|
||||
self.__prompt_queue = self.server.prompt_queue
|
||||
self.__history_store = None
|
||||
if history_file:
|
||||
try:
|
||||
self.__history_store = HistoryStore(history_file)
|
||||
except Exception:
|
||||
logger.exception("Failed to initialize persistent history")
|
||||
self.__prompt_queue_put = self.__prompt_queue.put
|
||||
self.__prompt_queue_get_flags = self.__prompt_queue.get_flags
|
||||
self.__user_manager_get_request_user_id = getattr(
|
||||
@@ -350,7 +363,7 @@ class AccessControl:
|
||||
prompt = self._unwrap_queue_item(wrapped_prompt)
|
||||
user_id = self._queue_item_user_id(wrapped_prompt)
|
||||
|
||||
if len(self.__prompt_queue.history) > MAXIMUM_HISTORY_SIZE:
|
||||
if len(self.__prompt_queue.history) >= MAXIMUM_HISTORY_SIZE:
|
||||
self.__prompt_queue.history.pop(next(iter(self.__prompt_queue.history)))
|
||||
|
||||
status_dict: Optional[dict] = None
|
||||
@@ -367,8 +380,46 @@ class AccessControl:
|
||||
"user_id": user_id,
|
||||
}
|
||||
self.__prompt_queue.history[prompt[1]].update(history_result)
|
||||
self._save_history_item(prompt[1])
|
||||
self.server.queue_updated()
|
||||
|
||||
def _save_history_item(self, prompt_id: str) -> None:
|
||||
if not self.__history_store:
|
||||
return
|
||||
try:
|
||||
self.__history_store.save(
|
||||
prompt_id,
|
||||
self.__prompt_queue.history[prompt_id],
|
||||
MAXIMUM_HISTORY_SIZE,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to persist history item %s", prompt_id)
|
||||
|
||||
def _load_persisted_history(self) -> None:
|
||||
if not self.__history_store:
|
||||
return
|
||||
try:
|
||||
persisted = self.__history_store.load(MAXIMUM_HISTORY_SIZE)
|
||||
persisted.update(self.__prompt_queue.history)
|
||||
self.__prompt_queue.history = persisted
|
||||
except Exception:
|
||||
logger.exception("Failed to restore persisted history")
|
||||
|
||||
def _delete_persisted_history(
|
||||
self, prompt_id: str = None, owner_id: str = None, clear: bool = False
|
||||
) -> None:
|
||||
if not self.__history_store:
|
||||
return
|
||||
try:
|
||||
if clear:
|
||||
self.__history_store.clear()
|
||||
elif owner_id is not None:
|
||||
self.__history_store.delete_owner(owner_id)
|
||||
elif prompt_id is not None:
|
||||
self.__history_store.delete(prompt_id)
|
||||
except Exception:
|
||||
logger.exception("Failed to update persistent history")
|
||||
|
||||
def user_queue_get_current_queue(self):
|
||||
"""Get the current user-specific queue."""
|
||||
with self.__prompt_queue.mutex:
|
||||
@@ -466,21 +517,25 @@ class AccessControl:
|
||||
current_user_id = self.get_current_user_id()
|
||||
if self.is_admin_user(current_user_id):
|
||||
self.__prompt_queue.history = {}
|
||||
self._delete_persisted_history(clear=True)
|
||||
else:
|
||||
self.__prompt_queue.history = {
|
||||
k: v
|
||||
for k, v in self.__prompt_queue.history.items()
|
||||
if v.get("user_id") != current_user_id
|
||||
}
|
||||
self._delete_persisted_history(owner_id=current_user_id)
|
||||
|
||||
def user_queue_delete_history_item(self, id_to_delete):
|
||||
with self.__prompt_queue.mutex:
|
||||
history_item = self.__prompt_queue.history.get(id_to_delete)
|
||||
if history_item and self._can_access_owner(history_item.get("user_id")):
|
||||
self.__prompt_queue.history.pop(id_to_delete, None)
|
||||
self._delete_persisted_history(prompt_id=id_to_delete)
|
||||
|
||||
def patch_prompt_queue(self):
|
||||
"""Patch the prompt queue with user-specific methods."""
|
||||
self._load_persisted_history()
|
||||
self.__prompt_queue.put = self.user_queue_put
|
||||
self.__prompt_queue.get = self.user_queue_get
|
||||
self.__prompt_queue.task_done = self.user_queue_task_done
|
||||
|
||||
@@ -42,6 +42,7 @@ TOKEN_EXPIRE_MINUTES = 60 * config.get("access_token_expiration_hours", 8760)
|
||||
MAX_TOKEN_EXPIRE_MINUTES = 60 * config.get("max_access_token_expiration_hours", 8760)
|
||||
|
||||
USERS_FILE = os.path.join(EXT_PATH, config.get("users_db", "users_db.json"))
|
||||
HISTORY_FILE = os.path.join(EXT_PATH, config.get("history_db", "history.sqlite3"))
|
||||
LOG_FILE = os.path.join(EXT_PATH, config.get("log", "account_manager.log"))
|
||||
LOG_LEVELS = config.get("log_levels", ["INFO"])
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
from contextlib import closing
|
||||
|
||||
|
||||
class HistoryStore:
|
||||
"""Persist completed ComfyUI history entries in a small SQLite database."""
|
||||
|
||||
def __init__(self, database: str):
|
||||
self.database = os.fspath(database)
|
||||
parent = os.path.dirname(os.path.abspath(self.database))
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
self._initialize()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
return sqlite3.connect(self.database, timeout=10)
|
||||
|
||||
def _initialize(self) -> None:
|
||||
with closing(self._connect()) as connection:
|
||||
with connection:
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS history (
|
||||
sequence INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
prompt_id TEXT NOT NULL UNIQUE,
|
||||
owner_id TEXT NOT NULL,
|
||||
data TEXT NOT NULL
|
||||
)
|
||||
"""
|
||||
)
|
||||
connection.execute(
|
||||
"CREATE INDEX IF NOT EXISTS history_owner_id ON history(owner_id)"
|
||||
)
|
||||
|
||||
def load(self, max_items: int) -> dict:
|
||||
with closing(self._connect()) as connection:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT prompt_id, owner_id, data
|
||||
FROM history
|
||||
ORDER BY sequence DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(max_items,),
|
||||
).fetchall()
|
||||
|
||||
history = {}
|
||||
for prompt_id, owner_id, data in reversed(rows):
|
||||
item = json.loads(data)
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
item["user_id"] = owner_id
|
||||
history[prompt_id] = item
|
||||
return history
|
||||
|
||||
def save(self, prompt_id: str, item: dict, max_items: int) -> None:
|
||||
owner_id = item.get("user_id") or ""
|
||||
data = json.dumps(item, ensure_ascii=False, separators=(",", ":"))
|
||||
with closing(self._connect()) as connection:
|
||||
with connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO history (prompt_id, owner_id, data)
|
||||
VALUES (?, ?, ?)
|
||||
ON CONFLICT(prompt_id) DO UPDATE SET
|
||||
owner_id = excluded.owner_id,
|
||||
data = excluded.data
|
||||
""",
|
||||
(prompt_id, owner_id, data),
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
DELETE FROM history
|
||||
WHERE sequence IN (
|
||||
SELECT sequence
|
||||
FROM history
|
||||
ORDER BY sequence DESC
|
||||
LIMIT -1 OFFSET ?
|
||||
)
|
||||
""",
|
||||
(max_items,),
|
||||
)
|
||||
|
||||
def delete(self, prompt_id: str) -> None:
|
||||
with closing(self._connect()) as connection:
|
||||
with connection:
|
||||
connection.execute("DELETE FROM history WHERE prompt_id = ?", (prompt_id,))
|
||||
|
||||
def delete_owner(self, owner_id: str) -> None:
|
||||
with closing(self._connect()) as connection:
|
||||
with connection:
|
||||
connection.execute("DELETE FROM history WHERE owner_id = ?", (owner_id,))
|
||||
|
||||
def clear(self) -> None:
|
||||
with closing(self._connect()) as connection:
|
||||
with connection:
|
||||
connection.execute("DELETE FROM history")
|
||||
Reference in New Issue
Block a user