From be0693da721fbf80fd01ad0d41b1cf9ecf9955e3 Mon Sep 17 00:00:00 2001
From: an-lee <an.lee.work@gmail.com>
Date: Sun, 8 Sep 2024 09:52:43 +0800
Subject: [PATCH] refactor db

---
 enjoy/src/main/db/index.ts | 266 +++++++++++++++++++------------------
 1 file changed, 136 insertions(+), 130 deletions(-)

diff --git a/enjoy/src/main/db/index.ts b/enjoy/src/main/db/index.ts
index 0b785996..f972d656 100644
--- a/enjoy/src/main/db/index.ts
+++ b/enjoy/src/main/db/index.ts
@@ -43,166 +43,172 @@ import path from "path";
 import url from "url";
 import { i18n } from "@main/i18n";
 import { UserSettingKeyEnum } from "@/types/enums";
+import log from "@main/logger";
 
 const __filename = url.fileURLToPath(import.meta.url);
 const __dirname = path.dirname(__filename);
+const logger = log.scope("DB");
 
 const db = {
   connection: null as Sequelize | null,
   connect: async () => {},
   disconnect: async () => {},
   registerIpcHandlers: () => {},
+  isConnecting: false,
 };
 
+const handlers = [
+  audiosHandler,
+  cacheObjectsHandler,
+  chatAgentsHandler,
+  chatMembersHandler,
+  chatMessagesHandler,
+  chatsHandler,
+  conversationsHandler,
+  messagesHandler,
+  notesHandler,
+  pronunciationAssessmentsHandler,
+  recordingsHandler,
+  segmentsHandler,
+  speechesHandler,
+  transcriptionsHandler,
+  userSettingsHandler,
+  videosHandler,
+];
+
 db.connect = async () => {
-  if (db.connection) {
-    return;
-  }
-  const dbPath = settings.dbPath();
-  if (!dbPath) {
-    throw new Error("Db path is not ready");
+  // Use a lock to prevent concurrent connections
+  if (db.isConnecting) {
+    throw new Error("Database connection is already in progress");
   }
 
-  const sequelize = new Sequelize({
-    dialect: "sqlite",
-    storage: dbPath,
-    models: [
-      Audio,
-      CacheObject,
-      Chat,
-      ChatAgent,
-      ChatMember,
-      ChatMessage,
-      Conversation,
-      Message,
-      Note,
-      PronunciationAssessment,
-      Recording,
-      Segment,
-      Speech,
-      Transcription,
-      UserSetting,
-      Video,
-    ],
-  });
+  db.isConnecting = true;
 
-  db.connection = sequelize;
-
-  const migrationResolver: Resolver<unknown> = ({
-    name,
-    path: filepath,
-    context,
-  }) => {
-    if (!filepath) {
-      throw new Error(
-        `Can't use default resolver for non-filesystem migrations`
-      );
+  try {
+    if (db.connection) {
+      return;
+    }
+    const dbPath = settings.dbPath();
+    if (!dbPath) {
+      throw new Error("Db path is not ready");
     }
 
-    const loadModule: () => Promise<RunnableMigration<unknown>> = async () => {
-      if (os.platform() === "win32") {
-        return import(`file://${filepath}`) as Promise<
-          RunnableMigration<unknown>
-        >;
-      } else {
-        return import(filepath) as Promise<RunnableMigration<unknown>>;
-      }
-    };
+    const sequelize = new Sequelize({
+      dialect: "sqlite",
+      storage: dbPath,
+      models: [
+        Audio,
+        CacheObject,
+        Chat,
+        ChatAgent,
+        ChatMember,
+        ChatMessage,
+        Conversation,
+        Message,
+        Note,
+        PronunciationAssessment,
+        Recording,
+        Segment,
+        Speech,
+        Transcription,
+        UserSetting,
+        Video,
+      ],
+    });
 
-    const getModule = async () => {
-      return await loadModule();
-    };
-
-    return {
+    const migrationResolver: Resolver<unknown> = ({
       name,
       path: filepath,
-      up: async () => (await getModule()).up({ path: filepath, name, context }),
-      down: async () =>
-        (await getModule()).down?.({ path: filepath, name, context }),
-    };
-  };
-
-  const umzug = new Umzug({
-    migrations: {
-      glob: ["migrations/*.js", { cwd: __dirname }],
-      resolve: migrationResolver,
-    },
-    context: sequelize.getQueryInterface(),
-    storage: new SequelizeStorage({ sequelize }),
-    logger: console,
-  });
-
-  // migrate up to the latest state
-  await umzug.up();
-
-  await sequelize.query("PRAGMA foreign_keys = false;");
-  await sequelize.sync();
-  await sequelize.authenticate();
-
-  // kill the zombie transcribe processes
-  Transcription.findAll({
-    where: {
-      state: "processing",
-    },
-  }).then((transcriptions) => {
-    transcriptions.forEach((transcription) => {
-      if (transcription.result) {
-        transcription.update({ state: "finished" });
-      } else {
-        transcription.update({ state: "pending" });
+      context,
+    }) => {
+      if (!filepath) {
+        throw new Error(
+          `Can't use default resolver for non-filesystem migrations`
+        );
       }
+
+      const loadModule: () => Promise<
+        RunnableMigration<unknown>
+      > = async () => {
+        if (os.platform() === "win32") {
+          return import(`file://${filepath}`) as Promise<
+            RunnableMigration<unknown>
+          >;
+        } else {
+          return import(filepath) as Promise<RunnableMigration<unknown>>;
+        }
+      };
+
+      const getModule = async () => {
+        return await loadModule();
+      };
+
+      return {
+        name,
+        path: filepath,
+        up: async () =>
+          (await getModule()).up({ path: filepath, name, context }),
+        down: async () =>
+          (await getModule()).down?.({ path: filepath, name, context }),
+      };
+    };
+
+    const umzug = new Umzug({
+      migrations: {
+        glob: ["migrations/*.js", { cwd: __dirname }],
+        resolve: migrationResolver,
+      },
+      context: sequelize.getQueryInterface(),
+      storage: new SequelizeStorage({ sequelize }),
+      logger: logger,
     });
-  });
 
-  // migrate settings
-  await UserSetting.migrateFromSettings();
+    try {
+      // migrate up to the latest state
+      await umzug.up();
 
-  // initialize i18n
-  const language = (await UserSetting.get(
-    UserSettingKeyEnum.LANGUAGE
-  )) as string;
-  i18n(language);
+      await sequelize.query("PRAGMA foreign_keys = false;");
+      await sequelize.sync();
+      await sequelize.authenticate();
+    } catch (err) {
+      logger.error(err);
+      await sequelize.close();
+      throw err;
+    }
 
-  // vacuum the database
-  await sequelize.query("VACUUM");
+    // migrate settings
+    await UserSetting.migrateFromSettings();
 
-  // register handlers
-  audiosHandler.register();
-  cacheObjectsHandler.register();
-  chatAgentsHandler.register();
-  chatMembersHandler.register();
-  chatMessagesHandler.register();
-  chatsHandler.register();
-  conversationsHandler.register();
-  messagesHandler.register();
-  notesHandler.register();
-  pronunciationAssessmentsHandler.register();
-  recordingsHandler.register();
-  segmentsHandler.register();
-  speechesHandler.register();
-  transcriptionsHandler.register();
-  userSettingsHandler.register();
-  videosHandler.register();
+    // initialize i18n
+    const language = (await UserSetting.get(
+      UserSettingKeyEnum.LANGUAGE
+    )) as string;
+    i18n(language);
+
+    // vacuum the database
+    await sequelize.query("VACUUM");
+
+    // register handlers
+
+    for (const handler of handlers) {
+      handler.register();
+    }
+
+    db.connection = sequelize;
+    logger.info("Database connection established");
+  } catch (err) {
+    logger.error(err);
+    throw err;
+  } finally {
+    db.isConnecting = false;
+  }
 };
 
 db.disconnect = async () => {
   // unregister handlers
-  audiosHandler.unregister();
-  cacheObjectsHandler.unregister();
-  chatAgentsHandler.unregister();
-  chatMembersHandler.unregister();
-  chatMessagesHandler.unregister();
-  chatsHandler.unregister();
-  conversationsHandler.unregister();
-  messagesHandler.unregister();
-  notesHandler.unregister();
-  pronunciationAssessmentsHandler.unregister();
-  recordingsHandler.unregister();
-  segmentsHandler.unregister();
-  speechesHandler.unregister();
-  transcriptionsHandler.unregister();
-  userSettingsHandler.unregister();
-  videosHandler.unregister();
+  for (const handler of handlers) {
+    handler.unregister();
+  }
 
   await db.connection?.close();
   db.connection = null;