""" Cache utilities for SageUtils. Handles persistent storage and retrieval of model metadata, hashes, and info. """ import os import json import pathlib import re import hashlib import folder_paths import datetime users_path = pathlib.Path(folder_paths.get_user_directory()) sage_users_path = users_path / "default" / "SageUtils" sage_backup_path = sage_users_path / "backup" os.makedirs(str(sage_users_path), exist_ok=True) os.makedirs(str(sage_backup_path), exist_ok=True) class SageCache: """ Persistent cache for model metadata, hashes, and info. """ def __init__(self): self.main_path = sage_users_path / "sage_cache.json" self.info_path = sage_users_path / "sage_cache_info.json" self.hash_path = sage_users_path / "sage_cache_hash.json" self.info_backup_path = sage_backup_path / "sage_cache_info.json" self.hash_backup_path = sage_backup_path / "sage_cache_hash.json" self.data = {} self.hash = {} self.info = {} self.last_hash = {} self.last_info = {} def by_path(self, file_path: str) -> dict: """Get cache info by file path.""" the_hash = self.hash.get(file_path, "") if the_hash: return self.info.get(the_hash, {}) print(f"No hash found for file: {file_path}") return {} def by_hash(self, file_hash: str) -> dict: """Get cache info by file hash.""" return self.info.get(file_hash, {}) def convert_old_cache(self): """Convert old cache format to new format, splitting into hash and info.""" print("Converting old cache format to new format.") for key in self.data: current_hash = self.data[key].get("hash", "") if current_hash: self.hash[key] = current_hash # Add the ones not on civitai first if self.data[key].get("civitai", False) == False: self.info[current_hash] = self.data[key] for key in self.data: current_hash = self.data[key].get("hash", "") # Add the ones on civitai, overwriting the previous ones if current_hash and self.data[key].get("civitai", False): self.info[current_hash] = self.data[key] def prune_old_backups(self, prefix, min_count=7, min_days=7): """Prune old backup files, keeping only the newest file for each unique hash, then enforce min_count and min_days.""" print(f"Pruning old backups with prefix: {prefix}") # Updated regex to allow optional microseconds in timestamp pattern = re.compile(rf"{re.escape(prefix)}-(\d{{4}}-\d{{2}}-\d{{2}}T\d{{2}}:\d{{2}}:\d{{2}}(?:\.\d+)?).*\.json$") backups = [] for f in sage_backup_path.iterdir(): if f.is_file() and f.name.startswith(prefix) and f.suffix == ".json": m = pattern.search(f.name) if m: try: dt = datetime.datetime.fromisoformat(m.group(1)) backups.append((dt, f)) except Exception as e: continue # Sort by date, newest first backups.sort(reverse=True) # Deduplicate by hash: keep only the newest file for each hash file_hash_map = {} for dt, f in backups: try: with f.open('rb') as file_obj: file_bytes = file_obj.read() file_hash = hashlib.sha256(file_bytes).hexdigest() file_hash_map[f] = file_hash except Exception: continue hash_to_file = {} deduped_backups = [] for dt, f in backups: file_hash = file_hash_map.get(f) if not file_hash: continue if file_hash not in hash_to_file: hash_to_file[file_hash] = (dt, f) deduped_backups.append((dt, f)) else: # Remove duplicate try: f.unlink() except Exception as e: pass # Now apply min_count and min_days logic to deduped_backups deduped_backups.sort(reverse=True) keep = deduped_backups[:min_count] now = datetime.datetime.now() days_kept = set(dt.date() for dt, _ in keep) for dt, f in deduped_backups[min_count:]: if (now - dt).days < min_days and dt.date() not in days_kept: keep.append((dt, f)) days_kept.add(dt.date()) keep_files = set(f for _, f in keep) for _, f in deduped_backups: if f not in keep_files: try: f.unlink() except Exception as e: pass print(f"Pruned old backups with prefix: {prefix}") def _atomic_write_json(self, path, data): """Write JSON data to a file atomically.""" import tempfile import os temp_dir = path.parent with tempfile.NamedTemporaryFile('w', dir=temp_dir, delete=False, encoding='utf-8') as tf: json.dump(data, tf, separators=(",", ":"), sort_keys=True, indent=4) tf.flush() os.fsync(tf.fileno()) tempname = tf.name os.replace(tempname, path) def _save_json(self, path, data, label): """Save data to a JSON file atomically, backing up and pruning old backups on error.""" try: self._atomic_write_json(path, data) except Exception as e: print(f"Unable to save {label} to {path}: {e}") # If file exists, back it up with an error suffix current_date = datetime.datetime.now().isoformat() if path.is_file(): error_prefix = f"{path.stem}-save-error" error_backup_path = sage_backup_path / f"{error_prefix}-{current_date}.json" try: with path.open("r") as src_file, error_backup_path.open("w") as dst_file: dst_file.write(src_file.read()) print(f"Backed up problematic file to {error_backup_path}") self.prune_old_backups(error_prefix) except Exception as backup_e: print(f"Unable to backup error file {path} to {error_backup_path}: {backup_e}") def backup_json(self, backup_prefix, data, current_date): """Backup data to a JSON file atomically, pruning old backups. Skip if identical backup already exists.""" import hashlib import tempfile import os # Serialize data to JSON and hash it data_json = json.dumps(data, separators=(",", ":"), sort_keys=True, indent=4) data_hash = hashlib.sha256(data_json.encode("utf-8")).hexdigest() # Check for existing identical backup for f in sage_backup_path.iterdir(): if f.is_file() and f.name.startswith(backup_prefix) and f.suffix == ".json": try: with f.open("rb") as file_obj: file_bytes = file_obj.read() file_hash = hashlib.sha256(file_bytes).hexdigest() if file_hash == data_hash: print(f"Identical backup already exists: {f}. Skipping new backup.") return except Exception: continue # No identical backup found, proceed to save atomically backup_path = sage_backup_path / f"{backup_prefix}-{current_date}.json" try: temp_dir = backup_path.parent with tempfile.NamedTemporaryFile('w', dir=temp_dir, delete=False, encoding='utf-8') as tf: tf.write(data_json) tf.flush() os.fsync(tf.fileno()) tempname = tf.name os.replace(tempname, backup_path) self.prune_old_backups(backup_prefix) except Exception as e: print(f"Unable to backup {backup_prefix} to {backup_path}: {e}") def load_json_file(self, path, label, current_date): """Load data from a JSON file, backing up the file if an error occurs.""" try: with path.open("r") as read_file: data = json.load(read_file) return data except Exception as e: print(f"Unable to load {label} from {path}: {e}") # If file exists, back it up with an error suffix if path.is_file(): error_prefix = f"{path.stem}-error" error_backup_path = sage_backup_path / f"{error_prefix}-{current_date}.json" try: with path.open("r") as src_file, error_backup_path.open("w") as dst_file: dst_file.write(src_file.read()) print(f"Backed up problematic file to {error_backup_path}") self.prune_old_backups(error_prefix) except Exception as backup_e: print(f"Unable to backup error file {path} to {error_backup_path}: {backup_e}") return None def load(self): """Load cache from disk only if not already loaded or if file has changed.""" current_date = datetime.datetime.now().isoformat() # Track last modification times if not hasattr(self, 'hash_mtime'): self.hash_mtime = None if not hasattr(self, 'info_mtime'): self.info_mtime = None try: # Try to load new format (hash/info split) hash_needs_reload = False info_needs_reload = False if self.hash_path.is_file(): hash_mtime = self.hash_path.stat().st_mtime if not self.hash or self.hash_mtime != hash_mtime: hash_needs_reload = True if self.info_path.is_file(): info_mtime = self.info_path.stat().st_mtime if not self.info or self.info_mtime != info_mtime: info_needs_reload = True if self.hash_path.is_file() and self.info_path.is_file(): if hash_needs_reload: print("Loading hash cache from disk.") hash_data = self.load_json_file(self.hash_path, "hash cache", current_date) if hash_data is not None: self.hash = hash_data self.last_hash = self.hash.copy() self.hash_mtime = self.hash_path.stat().st_mtime self.backup_json("sage_cache_hash", self.hash, current_date) else: self.hash = {} self.last_hash = {} self.hash_mtime = None if info_needs_reload: print("Loading info cache from disk.") info_data = self.load_json_file(self.info_path, "info cache", current_date) if info_data is not None: self.info = info_data self.last_info = self.info.copy() self.info_mtime = self.info_path.stat().st_mtime self.backup_json("sage_cache_info", self.info, current_date) else: self.info = {} self.last_info = {} self.info_mtime = None # Fallback: try to load old format elif self.main_path.is_file(): data = self.load_json_file(self.main_path, "main cache", current_date) if data is not None: self.data = data self.convert_old_cache() else: self.data = {} except Exception as e: print(f"Unable to load cache: {e}") def save(self): """Save cache to disk.""" saved = False if self.hash and self.hash != self.last_hash: self._save_json(self.hash_path, self.hash, "hash cache") self.last_hash = self.hash saved = True if self.info and self.info != self.last_info: self._save_json(self.info_path, self.info, "info cache") self.last_info = self.info saved = True if saved: print("Saved cache to disk.") def add_entry(self, file_path: str, file_hash: str): self.hash[file_path] = file_hash if file_hash not in self.info: self.info[file_hash] = { "hash": file_hash, "lastUsed": "", "civitai": False, "filePath": file_path } self.save() def add_or_update_entry(self, file_path: str, info_dict: dict): """ Add or update a cache entry for a given file path. Ensures both hash and info are updated together. """ file_hash = info_dict.get("hash") if not file_hash: raise ValueError("info_dict must contain a 'hash' key") self.hash[file_path] = file_hash self.info[file_hash] = info_dict def remove_entry(self, file_path: str): """ Remove a cache entry by file path. Removes both hash and info if no other file uses the same hash. """ file_hash = self.hash.get(file_path) if file_hash: del self.hash[file_path] # Only remove info if no other file_path uses this hash if file_hash not in self.hash.values(): self.info.pop(file_hash, None) def update_last_used_by_path(self, file_path: str): """ Update the 'lastUsed' field for a given file path. """ file_hash = self.hash.get(file_path) if file_hash and file_hash in self.info: self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat() def update_last_used_by_hash(self, file_hash: str): """ Update the 'lastUsed' field for a given file hash. """ if file_hash in self.info: self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat() def get_last_used_by_path(self, file_path: str) -> datetime.datetime | None: """ Get the 'lastUsed' field for a given file path. Returns an empty string if the file path is not found. """ file_hash = self.hash.get(file_path) if file_hash and file_hash in self.info: last_used_str = self.info[file_hash].get('lastUsed', '') if last_used_str: return datetime.datetime.fromisoformat(last_used_str) return None # Global cache instance cache = SageCache()