mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
feat: add database migration system for model update schema
Add migration support to handle schema changes without data loss. Instead of dropping and recreating tables, the system now: - Uses CREATE TABLE IF NOT EXISTS for initial table creation - Adds _apply_migrations method to handle incremental schema updates - Adds _get_table_columns helper to inspect existing table structure - Adds new columns to model_update_versions table (sort_index, name, base_model, released_at, size_bytes, preview_url, is_in_library, should_ignore) - Adds should_ignore_model column to model_update_status table This ensures existing databases are upgraded gracefully while preserving user data.
This commit is contained in:
@@ -77,15 +77,13 @@ class ModelUpdateService:
|
||||
|
||||
_SCHEMA = """
|
||||
PRAGMA foreign_keys = ON;
|
||||
DROP TABLE IF EXISTS model_update_versions;
|
||||
DROP TABLE IF EXISTS model_update_status;
|
||||
CREATE TABLE model_update_status (
|
||||
CREATE TABLE IF NOT EXISTS model_update_status (
|
||||
model_id INTEGER PRIMARY KEY,
|
||||
model_type TEXT NOT NULL,
|
||||
last_checked_at REAL,
|
||||
should_ignore_model INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE TABLE model_update_versions (
|
||||
CREATE TABLE IF NOT EXISTS model_update_versions (
|
||||
version_id INTEGER PRIMARY KEY,
|
||||
model_id INTEGER NOT NULL,
|
||||
sort_index INTEGER NOT NULL DEFAULT 0,
|
||||
@@ -129,11 +127,68 @@ class ModelUpdateService:
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA foreign_keys = ON")
|
||||
conn.executescript(self._SCHEMA)
|
||||
self._apply_migrations(conn)
|
||||
self._schema_initialized = True
|
||||
except Exception as exc: # pragma: no cover - defensive guard
|
||||
logger.error("Failed to initialize update schema: %s", exc, exc_info=True)
|
||||
raise
|
||||
|
||||
def _apply_migrations(self, conn: sqlite3.Connection) -> None:
|
||||
"""Ensure legacy databases match the current schema without dropping data."""
|
||||
|
||||
status_columns = self._get_table_columns(conn, "model_update_status")
|
||||
if "should_ignore_model" not in status_columns:
|
||||
conn.execute(
|
||||
"ALTER TABLE model_update_status "
|
||||
"ADD COLUMN should_ignore_model INTEGER NOT NULL DEFAULT 0"
|
||||
)
|
||||
|
||||
version_columns = self._get_table_columns(conn, "model_update_versions")
|
||||
migrations = {
|
||||
"sort_index": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN sort_index INTEGER NOT NULL DEFAULT 0"
|
||||
),
|
||||
"name": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN name TEXT"
|
||||
),
|
||||
"base_model": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN base_model TEXT"
|
||||
),
|
||||
"released_at": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN released_at TEXT"
|
||||
),
|
||||
"size_bytes": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN size_bytes INTEGER"
|
||||
),
|
||||
"preview_url": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN preview_url TEXT"
|
||||
),
|
||||
"is_in_library": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN is_in_library INTEGER NOT NULL DEFAULT 0"
|
||||
),
|
||||
"should_ignore": (
|
||||
"ALTER TABLE model_update_versions "
|
||||
"ADD COLUMN should_ignore INTEGER NOT NULL DEFAULT 0"
|
||||
),
|
||||
}
|
||||
|
||||
for column, statement in migrations.items():
|
||||
if column not in version_columns:
|
||||
conn.execute(statement)
|
||||
|
||||
def _get_table_columns(self, conn: sqlite3.Connection, table: str) -> set[str]:
|
||||
"""Return the set of existing columns for a table."""
|
||||
|
||||
cursor = conn.execute(f"PRAGMA table_info({table})")
|
||||
return {row["name"] for row in cursor.fetchall()}
|
||||
|
||||
async def refresh_for_model_type(
|
||||
self,
|
||||
model_type: str,
|
||||
|
||||
Reference in New Issue
Block a user