diff --git a/tracker.py b/tracker.py index dc33817..82cf7e9 100644 --- a/tracker.py +++ b/tracker.py @@ -7,7 +7,23 @@ from datetime import datetime, timezone, timedelta logger = logging.getLogger(__name__) -DB_PATH = os.path.join(os.path.dirname(__file__), "usage_stats.db") + +def _default_db_path(): + """Store the DB in ComfyUI's user directory so it survives extension reinstalls.""" + try: + import folder_paths + + user_dir = folder_paths.get_user_directory() + db_dir = os.path.join(user_dir, "nodes_stats") + os.makedirs(db_dir, exist_ok=True) + return os.path.join(db_dir, "usage_stats.db") + except Exception: + # Fallback to extension directory if folder_paths is unavailable + return os.path.join(os.path.dirname(__file__), "usage_stats.db") + + +DB_PATH = _default_db_path() +_OLD_DB_PATH = os.path.join(os.path.dirname(__file__), "usage_stats.db") SCHEMA = """ CREATE TABLE IF NOT EXISTS node_usage ( @@ -47,6 +63,7 @@ class UsageTracker: """Create tables on first use. Called under self._lock.""" if self._initialized: return + self._migrate_old_db() conn = sqlite3.connect(self._db_path) try: conn.executescript(SCHEMA) @@ -55,6 +72,22 @@ class UsageTracker: conn.close() self._initialized = True + def _migrate_old_db(self): + """Move DB from old extension-local path to the new user directory.""" + if self._db_path == _OLD_DB_PATH: + return + if not os.path.exists(_OLD_DB_PATH): + return + if os.path.exists(self._db_path): + # New location already has data, skip migration + return + try: + import shutil + shutil.move(_OLD_DB_PATH, self._db_path) + logger.info("nodes-stats: migrated DB to %s", self._db_path) + except Exception: + logger.warning("nodes-stats: failed to migrate old DB", exc_info=True) + def _connect(self): return sqlite3.connect(self._db_path)