using LiteDB; using LiteDB.Async; using Microsoft.Extensions.Logging; using Microsoft.Extensions.Options; using StabilityMatrix.Core.Extensions; using StabilityMatrix.Core.Models.Api; using StabilityMatrix.Core.Models.Configs; using StabilityMatrix.Core.Models.Database; using StabilityMatrix.Core.Services; namespace StabilityMatrix.Core.Database; public class LiteDbContext : ILiteDbContext { private readonly ILogger logger; private readonly ISettingsManager settingsManager; private readonly DebugOptions debugOptions; private LiteDatabaseAsync? database; public LiteDatabaseAsync Database => database ??= CreateDatabase(); // Notification events public event EventHandler? CivitModelsChanged; // Collections (Tables) public ILiteCollectionAsync CivitModels => Database.GetCollection("CivitModels"); public ILiteCollectionAsync CivitModelVersions => Database.GetCollection("CivitModelVersions"); public ILiteCollectionAsync CivitModelQueryCache => Database.GetCollection("CivitModelQueryCache"); public ILiteCollectionAsync GithubCache => Database.GetCollection("GithubCache"); public ILiteCollectionAsync LocalModelFiles => Database.GetCollection("LocalModelFiles"); public ILiteCollectionAsync InferenceProjects => Database.GetCollection("InferenceProjects"); public ILiteCollectionAsync LocalImageFiles => Database.GetCollection("LocalImageFiles"); public LiteDbContext( ILogger logger, ISettingsManager settingsManager, IOptions debugOptions ) { this.logger = logger; this.settingsManager = settingsManager; this.debugOptions = debugOptions.Value; } private LiteDatabaseAsync CreateDatabase() { LiteDatabaseAsync? db = null; if (debugOptions.TempDatabase) { db = new LiteDatabaseAsync(":temp:"); } else { // Attempt to create connection, might be in use try { var dbPath = Path.Combine(settingsManager.LibraryDir, "StabilityMatrix.db"); db = new LiteDatabaseAsync( new ConnectionString() { Filename = dbPath, Connection = ConnectionType.Shared, } ); } catch (IOException e) { logger.LogWarning("Database in use or not accessible ({Message}), using temporary database", e.Message); } } // Fallback to temporary database db ??= new LiteDatabaseAsync(":temp:"); // Register reference fields LiteDBExtensions.Register(m => m.ModelVersions, "CivitModelVersions"); LiteDBExtensions.Register(e => e.Items, "CivitModels"); return db; } public async Task<(CivitModel?, CivitModelVersion?)> FindCivitModelFromFileHashAsync(string hashBlake3) { var version = await CivitModelVersions .Query() .Where( mv => mv.Files!.Select(f => f.Hashes).Select(hashes => hashes.BLAKE3).Any(hash => hash == hashBlake3) ) .FirstOrDefaultAsync() .ConfigureAwait(false); if (version is null) return (null, null); var model = await CivitModels .Query() .Include(m => m.ModelVersions) .Where(m => m.ModelVersions!.Select(v => v.Id).Any(id => id == version.Id)) .FirstOrDefaultAsync() .ConfigureAwait(false); return (model, version); } public async Task UpsertCivitModelAsync(CivitModel civitModel) { // Insert model versions first then model var versionsUpdated = await CivitModelVersions.UpsertAsync(civitModel.ModelVersions).ConfigureAwait(false); var updated = await CivitModels.UpsertAsync(civitModel).ConfigureAwait(false); // Notify listeners on any change var anyUpdated = versionsUpdated > 0 || updated; if (anyUpdated) { CivitModelsChanged?.Invoke(this, EventArgs.Empty); } return anyUpdated; } public async Task UpsertCivitModelAsync(IEnumerable civitModels) { var civitModelsArray = civitModels.ToArray(); // Get all model versions then insert models var versions = civitModelsArray.SelectMany(model => model.ModelVersions ?? new()); var versionsUpdated = await CivitModelVersions.UpsertAsync(versions).ConfigureAwait(false); var updated = await CivitModels.UpsertAsync(civitModelsArray).ConfigureAwait(false); // Notify listeners on any change var anyUpdated = versionsUpdated > 0 || updated > 0; if (updated > 0 || versionsUpdated > 0) { CivitModelsChanged?.Invoke(this, EventArgs.Empty); } return anyUpdated; } // Add to cache public async Task UpsertCivitModelQueryCacheEntryAsync(CivitModelQueryCacheEntry entry) { var changed = await CivitModelQueryCache.UpsertAsync(entry).ConfigureAwait(false); if (changed) { CivitModelsChanged?.Invoke(this, EventArgs.Empty); } return changed; } public async Task GetGithubCacheEntry(string? cacheKey) { if (string.IsNullOrEmpty(cacheKey)) return null; if (await GithubCache.FindByIdAsync(cacheKey).ConfigureAwait(false) is { } result) { return result; } return null; } public Task UpsertGithubCacheEntry(GithubCacheEntry cacheEntry) => GithubCache.UpsertAsync(cacheEntry); public void Dispose() { if (database is not null) { try { database.Dispose(); } catch (ObjectDisposedException) { } catch (ApplicationException) { // Ignores a mutex error from library // https://stability-matrix.sentry.io/share/issue/5c62f37462444e7eab18cea314af231f/ } database = null; } GC.SuppressFinalize(this); } }