ComfyUI-WorkflowPublisher/workflow_service/database.py

59 lines
2.4 KiB
Python

import aiosqlite
import json
import re
DATABASE_FILE = "/db/workflows_service.sqlite"
async def init_db():
async with aiosqlite.connect(DATABASE_FILE) as db:
await db.execute("""
CREATE TABLE IF NOT EXISTS workflows (
name TEXT PRIMARY KEY,
base_name TEXT NOT NULL,
version TEXT NOT NULL,
workflow_json TEXT NOT NULL
)
""")
await db.commit()
print(f"数据库 '{DATABASE_FILE}' 已准备就绪。")
async def save_workflow(name: str, workflow_json: str):
version_match = re.search(r"\[(20\d{12})\]$", name)
if not version_match:
raise ValueError("Workflow name must have a version suffix like [YYYYMMDDHHMMSS]")
version = version_match.group(1)
base_name = re.sub(r"\s*\[(20\d{12})\]$", "", name).strip()
async with aiosqlite.connect(DATABASE_FILE) as db:
await db.execute(
"INSERT OR REPLACE INTO workflows (name, base_name, version, workflow_json) VALUES (?, ?, ?, ?)",
(name, base_name, version, workflow_json)
)
await db.commit()
async def get_all_workflows() -> list[dict]:
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("SELECT name, workflow_json FROM workflows")
rows = await cursor.fetchall()
return [{"name": row["name"], "workflow": json.loads(row["workflow_json"])} for row in rows]
async def get_latest_workflow_by_base_name(base_name: str) -> dict | None:
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("SELECT * FROM workflows WHERE base_name = ? ORDER BY version DESC LIMIT 1", (base_name,))
row = await cursor.fetchone()
return dict(row) if row else None
async def get_workflow_by_version(base_name: str, version: str) -> dict | None:
name = f"{base_name} [{version}]"
async with aiosqlite.connect(DATABASE_FILE) as db:
db.row_factory = aiosqlite.Row
cursor = await db.execute("SELECT * FROM workflows WHERE name = ?", (name,))
row = await cursor.fetchone()
return dict(row) if row else None
async def delete_workflow(name: str) -> bool:
async with aiosqlite.connect(DATABASE_FILE) as db:
cursor = await db.execute("DELETE FROM workflows WHERE name = ?", (name,))
await db.commit()
return cursor.rowcount > 0