mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 15:15:44 -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 = """
|
_SCHEMA = """
|
||||||
PRAGMA foreign_keys = ON;
|
PRAGMA foreign_keys = ON;
|
||||||
DROP TABLE IF EXISTS model_update_versions;
|
CREATE TABLE IF NOT EXISTS model_update_status (
|
||||||
DROP TABLE IF EXISTS model_update_status;
|
|
||||||
CREATE TABLE model_update_status (
|
|
||||||
model_id INTEGER PRIMARY KEY,
|
model_id INTEGER PRIMARY KEY,
|
||||||
model_type TEXT NOT NULL,
|
model_type TEXT NOT NULL,
|
||||||
last_checked_at REAL,
|
last_checked_at REAL,
|
||||||
should_ignore_model INTEGER NOT NULL DEFAULT 0
|
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,
|
version_id INTEGER PRIMARY KEY,
|
||||||
model_id INTEGER NOT NULL,
|
model_id INTEGER NOT NULL,
|
||||||
sort_index INTEGER NOT NULL DEFAULT 0,
|
sort_index INTEGER NOT NULL DEFAULT 0,
|
||||||
@@ -129,11 +127,68 @@ class ModelUpdateService:
|
|||||||
conn.execute("PRAGMA journal_mode=WAL")
|
conn.execute("PRAGMA journal_mode=WAL")
|
||||||
conn.execute("PRAGMA foreign_keys = ON")
|
conn.execute("PRAGMA foreign_keys = ON")
|
||||||
conn.executescript(self._SCHEMA)
|
conn.executescript(self._SCHEMA)
|
||||||
|
self._apply_migrations(conn)
|
||||||
self._schema_initialized = True
|
self._schema_initialized = True
|
||||||
except Exception as exc: # pragma: no cover - defensive guard
|
except Exception as exc: # pragma: no cover - defensive guard
|
||||||
logger.error("Failed to initialize update schema: %s", exc, exc_info=True)
|
logger.error("Failed to initialize update schema: %s", exc, exc_info=True)
|
||||||
raise
|
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(
|
async def refresh_for_model_type(
|
||||||
self,
|
self,
|
||||||
model_type: str,
|
model_type: str,
|
||||||
|
|||||||
Reference in New Issue
Block a user