refactor: use sequelize replace better-sqlite3

electron.vite.config.ts

@@ -7,7 +7,8 @@ export default defineConfig({
     plugins: [externalizeDepsPlugin()],
     resolve: {
       alias: {
-        ollama: resolve('ollama/src')
+        '@types': resolve('src/renderer/src/types'),
+        '@main': resolve('src/main')
       }
     }
   },

package.json

@@ -31,11 +31,13 @@
   "dependencies": {
     "@electron-toolkit/preload": "^3.0.0",
     "@electron-toolkit/utils": "^3.0.0",
-    "better-sqlite3": "^11.3.0",
     "electron-log": "^5.1.5",
     "electron-store": "^8.2.0",
     "electron-updater": "^6.1.7",
-    "electron-window-state": "^5.0.3"
+    "electron-window-state": "^5.0.3",
+    "sequelize": "^6.37.3",
+    "sqlite3": "^5.1.7",
+    "umzug": "^3.8.1"
   },
   "devDependencies": {
     "@anthropic-ai/sdk": "^0.24.3",

src/main/config.ts

@@ -1,4 +1,14 @@
+import fs from 'node:fs'
+
+import { app } from 'electron'
 import Store from 'electron-store'
+import path from 'path'
+
+export const DATA_PATH = path.join(app.getPath('userData'), 'Data')
+
+if (!fs.existsSync(DATA_PATH)) {
+  fs.mkdirSync(DATA_PATH, { recursive: true })
+}
 
 export const appConfig = new Store()
 

src/main/database/index.ts

@@ -0,0 +1,34 @@
+import Logger from 'electron-log'
+import path from 'path'
+import { Sequelize } from 'sequelize'
+import { SequelizeStorage, Umzug } from 'umzug'
+
+import { DATA_PATH } from '../config'
+
+const sequelize = new Sequelize({
+  dialect: 'sqlite',
+  storage: path.join(DATA_PATH, 'data.db'),
+  logging: false
+})
+
+const umzug = new Umzug({
+  migrations: { glob: 'src/main/database/migrations/*.js' },
+  context: sequelize.getQueryInterface(),
+  storage: new SequelizeStorage({ sequelize, modelName: 'Migration', tableName: 'migrations' }),
+  logger: Logger
+})
+
+export async function initDatabase() {
+  try {
+    await sequelize.authenticate()
+    Logger.log('Database connection has been established successfully.')
+
+    // Run migrations
+    await umzug.up()
+    Logger.log('Migrations have been executed successfully.')
+  } catch (error) {
+    Logger.error('Migrations failed to execute:', error)
+  }
+}
+
+export default sequelize

src/main/database/migrations/20240912072241-create-files-table.js

@@ -0,0 +1,48 @@
+const { Sequelize } = require('sequelize')
+
+async function up({ context: queryInterface }) {
+  await queryInterface.createTable('files', {
+    id: {
+      type: Sequelize.TEXT,
+      primaryKey: true
+    },
+    name: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    file_name: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    path: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    size: {
+      type: Sequelize.INTEGER,
+      allowNull: false
+    },
+    ext: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    type: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    created_at: {
+      type: Sequelize.TEXT,
+      allowNull: false
+    },
+    count: {
+      type: Sequelize.INTEGER,
+      defaultValue: 1
+    }
+  })
+}
+
+async function down({ context: queryInterface }) {
+  await queryInterface.dropTable('files')
+}
+
+module.exports = { up, down }

src/main/database/models/FileModel.ts

@@ -0,0 +1,41 @@
+import { DataTypes, Model } from 'sequelize'
+
+import { FileMetadata, FileType } from '../../../renderer/src/types'
+import sequelize from '..'
+
+class FileModel extends Model<FileMetadata> implements FileMetadata {
+  public id!: string
+  public name!: string
+  public file_name!: string
+  public path!: string
+  public size!: number
+  public ext!: string
+  public type!: FileType
+  public created_at!: Date
+  public count!: number
+}
+
+FileModel.init(
+  {
+    id: {
+      type: DataTypes.UUID,
+      primaryKey: true
+    },
+    name: DataTypes.STRING,
+    file_name: DataTypes.STRING,
+    path: DataTypes.STRING,
+    size: DataTypes.INTEGER,
+    ext: DataTypes.STRING,
+    type: DataTypes.STRING,
+    created_at: DataTypes.DATE,
+    count: DataTypes.INTEGER
+  },
+  {
+    sequelize,
+    modelName: 'File',
+    tableName: 'files',
+    timestamps: false
+  }
+)
+
+export default FileModel

src/main/db/DatabaseMigrator.ts

@@ -1,95 +0,0 @@
-import Database from 'better-sqlite3'
-import { app } from 'electron'
-import Logger from 'electron-log'
-import * as fs from 'fs'
-import * as path from 'path'
-
-interface Migration {
-  id: number
-  name: string
-  sql: string
-}
-
-export class DatabaseMigrator {
-  private storageDir: string
-  private db: Database.Database
-  private migrationsDir: string
-
-  constructor(migrationsDir: string) {
-    this.storageDir = path.join(app.getPath('userData'), 'Data')
-    this.migrationsDir = migrationsDir
-    this.initStorageDir()
-    this.initDatabase()
-    this.initMigrationsTable()
-  }
-
-  private initStorageDir(): void {
-    if (!fs.existsSync(this.storageDir)) {
-      fs.mkdirSync(this.storageDir, { recursive: true })
-    }
-  }
-
-  private initDatabase(): void {
-    const dbPath = path.join(this.storageDir, 'data.db')
-    this.db = new Database(dbPath)
-  }
-
-  private initMigrationsTable(): void {
-    this.db.exec(`
-      CREATE TABLE IF NOT EXISTS migrations (
-        id INTEGER PRIMARY KEY,
-        name TEXT NOT NULL,
-        applied_at DATETIME DEFAULT CURRENT_TIMESTAMP
-      )
-    `)
-  }
-
-  private getAppliedMigrations(): number[] {
-    const stmt = this.db.prepare('SELECT id FROM migrations ORDER BY id')
-    return stmt.all().map((row: any) => row.id)
-  }
-
-  private loadMigrations(): Migration[] {
-    const files = fs.readdirSync(this.migrationsDir).filter((file) => file.endsWith('.sql'))
-    return files
-      .map((file) => {
-        const [id, ...nameParts] = path.basename(file, '.sql').split('_')
-        return {
-          id: parseInt(id),
-          name: nameParts.join('_'),
-          sql: fs.readFileSync(path.join(this.migrationsDir, file), 'utf-8')
-        }
-      })
-      .sort((a, b) => a.id - b.id)
-  }
-
-  public async migrate(): Promise<void> {
-    const appliedMigrations = this.getAppliedMigrations()
-    const allMigrations = this.loadMigrations()
-
-    const pendingMigrations = allMigrations.filter((migration) => !appliedMigrations.includes(migration.id))
-
-    this.db.exec('BEGIN TRANSACTION')
-
-    try {
-      for (const migration of pendingMigrations) {
-        Logger.log(`Applying migration: ${migration.id}_${migration.name}`)
-        this.db.exec(migration.sql)
-
-        const insertStmt = this.db.prepare('INSERT INTO migrations (id, name) VALUES (?, ?)')
-        insertStmt.run(migration.id, migration.name)
-      }
-
-      this.db.exec('COMMIT')
-      Logger.log('All migrations applied successfully')
-    } catch (error) {
-      this.db.exec('ROLLBACK')
-      Logger.error('Error applying migrations:', error)
-      throw error
-    }
-  }
-
-  public close(): void {
-    this.db.close()
-  }
-}

src/main/index.ts

@@ -1,30 +1,17 @@
 import { electronApp, optimizer } from '@electron-toolkit/utils'
 import { app, BrowserWindow } from 'electron'
-import Logger from 'electron-log'
-import path from 'path'
 
-import { DatabaseMigrator } from './db/DatabaseMigrator'
+import { initDatabase } from './database'
 import { registerIpc } from './ipc'
-import { getResourcePath } from './utils'
 import { updateUserDataPath } from './utils/upgrade'
 import { createMainWindow } from './window'
 
-async function migrateDatabase() {
-  const migrationsDir = path.join(getResourcePath(), 'migrations')
-  const migrator = new DatabaseMigrator(migrationsDir)
-
-  await migrator.migrate()
-  migrator.close()
-
-  Logger.log('Database migration completed successfully.')
-}
-
 // This method will be called when Electron has finished
 // initialization and is ready to create browser windows.
 // Some APIs can only be used after this event occurs.
 app.whenReady().then(async () => {
   await updateUserDataPath()
-  await migrateDatabase()
+  await initDatabase()
 
   // Set app user model id for windows
   electronApp.setAppUserModelId('com.kangfenmao.CherryStudio')

src/main/ipc.ts

@@ -1,12 +1,12 @@
+import { FileMetadata } from '@types'
 import { BrowserWindow, ipcMain, OpenDialogOptions, session, shell } from 'electron'
 import Logger from 'electron-log'
 import fs from 'fs'
 import path from 'path'
 
-import { FileMetadata } from '../renderer/src/types'
 import { appConfig, titleBarOverlayDark, titleBarOverlayLight } from './config'
-import { File } from './file'
+import AppUpdater from './services/AppUpdater'
-import AppUpdater from './updater'
+import File from './services/File'
 import { openFile, saveFile } from './utils/file'
 import { compress, decompress } from './utils/zip'
 import { createMinappWindow } from './window'

src/main/services/File.ts

@@ -1,21 +1,19 @@
-import Database from 'better-sqlite3'
+/* eslint-disable react/no-is-mounted */
+import FileModel from '@main/database/models/FileModel'
+import { getFileType } from '@main/utils/file'
+import { FileMetadata } from '@types'
 import * as crypto from 'crypto'
 import { app, dialog, OpenDialogOptions } from 'electron'
 import * as fs from 'fs'
 import * as path from 'path'
 import { v4 as uuidv4 } from 'uuid'
 
-import { FileMetadata } from '../renderer/src/types'
+class File {
-import { getFileType } from './utils/file'
-
-export class File {
   private storageDir: string
-  private db: Database.Database
 
   constructor() {
     this.storageDir = path.join(app.getPath('userData'), 'Data', 'Files')
     this.initStorageDir()
-    this.initDatabase()
   }
 
   private initStorageDir(): void {
@@ -24,11 +22,6 @@ export class File {
     }
   }
 
-  private initDatabase(): void {
-    const dbPath = path.join(app.getPath('userData'), 'Data', 'data.db')
-    this.db = new Database(dbPath)
-  }
-
   private async getFileHash(filePath: string): Promise<string> {
     return new Promise((resolve, reject) => {
       const hash = crypto.createHash('md5')
@@ -104,11 +97,10 @@ export class File {
 
     if (duplicateFile) {
       // Increment the count for the duplicate file
-      const updateStmt = this.db.prepare('UPDATE files SET count = count + 1 WHERE id = ?')
+      await FileModel.increment('count', { where: { id: duplicateFile.id } })
-      updateStmt.run(duplicateFile.id)
 
       // Fetch the updated file metadata
-      return this.getFile(duplicateFile.id)!
+      return (await this.getFile(duplicateFile.id))!
     }
 
     const uuid = uuidv4()
@@ -132,38 +124,21 @@ export class File {
       count: 1
     }
 
-    const stmt = this.db.prepare(`
+    await FileModel.create(fileMetadata)
-      INSERT INTO files (id, name, file_name, path, size, ext, type, created_at, count)
-      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
-    `)
-
-    stmt.run(
-      fileMetadata.id,
-      fileMetadata.name,
-      fileMetadata.file_name,
-      fileMetadata.path,
-      fileMetadata.size,
-      fileMetadata.ext,
-      fileMetadata.type,
-      fileMetadata.created_at.toISOString(),
-      fileMetadata.count
-    )
 
     return fileMetadata
   }
 
   async deleteFile(fileId: string): Promise<void> {
-    const fileMetadata = this.getFile(fileId)
+    const fileMetadata = await this.getFile(fileId)
     if (fileMetadata) {
       if (fileMetadata.count > 1) {
         // Decrement the count if there are multiple references
-        const updateStmt = this.db.prepare('UPDATE files SET count = count - 1 WHERE id = ?')
+        await FileModel.decrement('count', { where: { id: fileId } })
-        updateStmt.run(fileId)
       } else {
         // Delete the file and database entry if this is the last reference
         await fs.promises.unlink(fileMetadata.path)
-        const deleteStmt = this.db.prepare('DELETE FROM files WHERE id = ?')
+        await FileModel.destroy({ where: { id: fileId } })
-        deleteStmt.run(fileId)
       }
     }
   }
@@ -178,15 +153,15 @@ export class File {
     await Promise.all(deletePromises)
   }
 
-  getFile(id: string): FileMetadata | null {
+  async getFile(id: string): Promise<FileMetadata | null> {
-    const stmt = this.db.prepare('SELECT * FROM files WHERE id = ?')
+    const file = await FileModel.findByPk(id)
-    const row = stmt.get(id) as any
+    return file ? (file.toJSON() as FileMetadata) : null
-    return row ? { ...row, created_at: new Date(row.created_at) } : null
   }
 
-  getAllFiles(): FileMetadata[] {
+  async getAllFiles(): Promise<FileMetadata[]> {
-    const stmt = this.db.prepare('SELECT * FROM files')
+    const files = await FileModel.findAll()
-    const rows = stmt.all() as any[]
+    return files.map((file) => file.toJSON() as FileMetadata)
-    return rows.map((row) => ({ ...row, created_at: new Date(row.created_at) }))
   }
 }
+
+export default File

src/preload/index.d.ts

@@ -1,8 +1,7 @@
 import { ElectronAPI } from '@electron-toolkit/preload'
+import { FileMetadata } from '@renderer/types'
 import type { OpenDialogOptions } from 'electron'
 
-import type FileMetadata from '../main/file'
-
 declare global {
   interface Window {
     electron: ElectronAPI

src/renderer/src/config/models.ts

@@ -1,7 +1,7 @@
 import { Model } from '@renderer/types'
 
 const TEXT_TO_IMAGE_REGEX = /flux|diffusion|stabilityai|sd-turbo|dall|cogview/i
-const VISION_REGEX = /llava|moondream|minicpm|gemini|claude|vision/i
+const VISION_REGEX = /llava|moondream|minicpm|gemini|claude|vision|glm-4v/i
 const EMBEDDING_REGEX = /embedding/i
 
 export const SYSTEM_MODELS: Record<string, Model[]> = {

src/renderer/src/context/ThemeProvider.tsx

@@ -1,5 +1,5 @@
 import { useSettings } from '@renderer/hooks/useSettings'
-import { ThemeMode } from '@renderer/store/settings'
+import { ThemeMode } from '@renderer/types'
 import React, { createContext, PropsWithChildren, useContext, useEffect, useState } from 'react'
 
 interface ThemeContextType {

src/renderer/src/hooks/useSettings.ts

@@ -4,9 +4,9 @@ import {
   setSendMessageShortcut as _setSendMessageShortcut,
   setTheme,
   setTopicPosition,
-  setWindowStyle,
+  setWindowStyle
-  ThemeMode
 } from '@renderer/store/settings'
+import { ThemeMode } from '@renderer/types'
 
 export function useSettings() {
   const settings = useAppSelector((state) => state.settings)

src/renderer/src/init.ts

@@ -2,7 +2,7 @@ import KeyvStorage from '@kangfenmao/keyv-storage'
 import localforage from 'localforage'
 
 import { APP_NAME } from './config/env'
-import { ThemeMode } from './store/settings'
+import { ThemeMode } from './types'
 import { loadScript } from './utils'
 
 export async function initMermaid(theme: ThemeMode) {

src/renderer/src/pages/home/Markdown/CodeBlock.tsx

@@ -2,7 +2,7 @@ import { CheckOutlined } from '@ant-design/icons'
 import CopyIcon from '@renderer/components/Icons/CopyIcon'
 import { useTheme } from '@renderer/context/ThemeProvider'
 import { initMermaid } from '@renderer/init'
-import { ThemeMode } from '@renderer/store/settings'
+import { ThemeMode } from '@renderer/types'
 import React, { useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'

src/renderer/src/pages/settings/GeneralSettings.tsx

@@ -5,8 +5,9 @@ import { useSettings } from '@renderer/hooks/useSettings'
 import i18n from '@renderer/i18n'
 import { backup, reset, restore } from '@renderer/services/backup'
 import { useAppDispatch } from '@renderer/store'
-import { setLanguage, setUserName, ThemeMode } from '@renderer/store/settings'
+import { setLanguage, setUserName } from '@renderer/store/settings'
 import { setProxyUrl as _setProxyUrl } from '@renderer/store/settings'
+import { ThemeMode } from '@renderer/types'
 import { isValidProxyUrl } from '@renderer/utils'
 import { Button, Input, Select } from 'antd'
 import { FC, useState } from 'react'

src/renderer/src/store/settings.ts

@@ -1,13 +1,8 @@
 import { createSlice, PayloadAction } from '@reduxjs/toolkit'
+import { ThemeMode } from '@renderer/types'
 
 export type SendMessageShortcut = 'Enter' | 'Shift+Enter'
 
-export enum ThemeMode {
-  light = 'light',
-  dark = 'dark',
-  auto = 'auto'
-}
-
 export interface SettingsState {
   showAssistants: boolean
   showTopics: boolean

src/renderer/src/types/index.ts

@@ -106,3 +106,9 @@ export enum FileType {
   DOCUMENT = 'document',
   OTHER = 'other'
 }
+
+export enum ThemeMode {
+  light = 'light',
+  dark = 'dark',
+  auto = 'auto'
+}

tsconfig.node.json

@@ -11,6 +11,14 @@
     "composite": true,
     "types": [
       "electron-vite/node"
+    ],
+    "paths": {
+      "@types": [
+        "./src/renderer/src/types/index.ts"
+      ],
+      "@main/*": [
+        "./src/main/*"
       ]
     }
   }
+}