init

src/database/database-manager.ts

@@ -0,0 +1,180 @@
+import { PGlite } from '@electric-sql/pglite'
+import { type PGliteWithLive, live } from '@electric-sql/pglite/live'
+// import { PgliteDatabase, drizzle } from 'drizzle-orm/pglite'
+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<DBManager> {
+		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('Smart composer database initialized.')
+		return dbManager
+	}
+
+	// getDb() {
+	// 	return this.db
+	// }
+
+	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<void> {
+		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<void> {
+		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
+		}
+	}
+}