import { PGlite } from '@electric-sql/pglite' // @ts-expect-error import { type PGliteWithLive, live } from '@electric-sql/pglite/live' import { App, normalizePath } from 'obsidian' import { PGLITE_DB_PATH } from '../constants' import { ConversationManager } from './modules/conversation/conversation-manager' import { TemplateManager } from './modules/template/template-manager' import { VectorManager } from './modules/vector/vector-manager' import { pgliteResources } from './pglite-resources' import { migrations } from './sql' export class DBManager { private app: App private dbPath: string private db: PGliteWithLive | null = null // private db: PgliteDatabase | null = null private vectorManager: VectorManager private templateManager: TemplateManager private conversationManager: ConversationManager constructor(app: App, dbPath: string) { this.app = app this.dbPath = dbPath } static async create(app: App): Promise { const dbManager = new DBManager(app, normalizePath(PGLITE_DB_PATH)) await dbManager.loadExistingDatabase() if (!dbManager.db) { await dbManager.createNewDatabase() } await dbManager.migrateDatabase() await dbManager.save() dbManager.vectorManager = new VectorManager(app, dbManager) dbManager.templateManager = new TemplateManager(app, dbManager) dbManager.conversationManager = new ConversationManager(app, dbManager) console.log('infio database initialized.') return dbManager } getPgClient() { return this.db } getVectorManager() { return this.vectorManager } getTemplateManager() { return this.templateManager } getConversationManager() { return this.conversationManager } private async createNewDatabase() { const { fsBundle, wasmModule, vectorExtensionBundlePath } = await this.loadPGliteResources() this.db = await PGlite.create({ fsBundle: fsBundle, wasmModule: wasmModule, extensions: { vector: vectorExtensionBundlePath, live, }, }) } private async loadExistingDatabase() { try { const databaseFileExists = await this.app.vault.adapter.exists( this.dbPath, ) if (!databaseFileExists) { return null } const fileBuffer = await this.app.vault.adapter.readBinary(this.dbPath) const fileBlob = new Blob([fileBuffer], { type: 'application/x-gzip' }) const { fsBundle, wasmModule, vectorExtensionBundlePath } = await this.loadPGliteResources() this.db = await PGlite.create({ loadDataDir: fileBlob, fsBundle: fsBundle, wasmModule: wasmModule, extensions: { vector: vectorExtensionBundlePath, live }, }) // return drizzle(this.pgClient) } catch (error) { console.error('Error loading database:', error) return null } } private async migrateDatabase(): Promise { if (!this.db) { throw new Error('Database client not initialized'); } try { // Execute SQL migrations for (const [_key, migration] of Object.entries(migrations)) { // Split SQL into individual commands and execute them one by one const commands = migration.sql.split('\n\n').filter(cmd => cmd.trim()); for (const command of commands) { console.log('Executing SQL migration:', command); await this.db.query(command); } } } catch (error) { console.error('Error executing SQL migrations:', error); throw error; } } async save(): Promise { if (!this.db) { return } try { const blob: Blob = await this.db.dumpDataDir('gzip') await this.app.vault.adapter.writeBinary( this.dbPath, Buffer.from(await blob.arrayBuffer()), ) } catch (error) { console.error('Error saving database:', error) } } async cleanup() { this.db?.close() this.db = null } private async loadPGliteResources(): Promise<{ fsBundle: Blob wasmModule: WebAssembly.Module vectorExtensionBundlePath: URL }> { try { // Convert base64 to binary data const wasmBinary = Buffer.from(pgliteResources.wasmBase64, 'base64') const dataBinary = Buffer.from(pgliteResources.dataBase64, 'base64') const vectorBinary = Buffer.from(pgliteResources.vectorBase64, 'base64') // Create blobs from binary data const fsBundle = new Blob([dataBinary], { type: 'application/octet-stream', }) const wasmModule = await WebAssembly.compile(wasmBinary) // Create a blob URL for the vector extension const vectorBlob = new Blob([vectorBinary], { type: 'application/gzip', }) const vectorExtensionBundlePath = URL.createObjectURL(vectorBlob) return { fsBundle, wasmModule, vectorExtensionBundlePath: new URL(vectorExtensionBundlePath), } } catch (error) { console.error('Error loading PGlite resources:', error) throw error } } }