From 33456f1bd0522516f5d4b0fea7aad9260168bf4d Mon Sep 17 00:00:00 2001
From: Philipinho <16838612+Philipinho@users.noreply.github.com>
Date: Sat, 27 Apr 2024 18:53:30 +0100
Subject: [PATCH] database connection retry

---
 apps/client/src/App.tsx                       |  1 +
 apps/server/src/app.module.ts                 |  4 +-
 .../src/collaboration/collaboration.module.ts |  8 ++-
 ...kysely-db.module.ts => database.module.ts} | 62 +++++++++++++++++--
 4 files changed, 67 insertions(+), 8 deletions(-)
 rename apps/server/src/kysely/{kysely-db.module.ts => database.module.ts} (56%)

diff --git a/apps/client/src/App.tsx b/apps/client/src/App.tsx
index 66c167dc..ee2e3c01 100644
--- a/apps/client/src/App.tsx
+++ b/apps/client/src/App.tsx
@@ -67,6 +67,7 @@ export default function App() {
 
         <Route path={"/settings"} element={<SettingsLayout />}>
           <Route path={"profile"} element={<AccountSettings />} />
+          <Route path={"profile/preferences"} element={<AccountSettings />} />
           <Route path={"workspace"} element={<WorkspaceSettings />} />
           <Route path={"members"} element={<WorkspaceMembers />} />
           <Route path={"groups"} element={<Groups />} />
diff --git a/apps/server/src/app.module.ts b/apps/server/src/app.module.ts
index 20790972..fcd040cd 100644
--- a/apps/server/src/app.module.ts
+++ b/apps/server/src/app.module.ts
@@ -7,7 +7,7 @@ import { CollaborationModule } from './collaboration/collaboration.module';
 import { WsModule } from './ws/ws.module';
 import { ServeStaticModule } from '@nestjs/serve-static';
 import { join } from 'path';
-import { KyselyDbModule } from './kysely/kysely-db.module';
+import { DatabaseModule } from '@docmost/db/database.module';
 import * as fs from 'fs';
 
 const clientDistPath = join(__dirname, '..', '..', 'client/dist');
@@ -26,7 +26,7 @@ function getServeStaticModule() {
 @Module({
   imports: [
     CoreModule,
-    KyselyDbModule,
+    DatabaseModule,
     EnvironmentModule,
     CollaborationModule,
     WsModule,
diff --git a/apps/server/src/collaboration/collaboration.module.ts b/apps/server/src/collaboration/collaboration.module.ts
index 7bda3918..7b48edc7 100644
--- a/apps/server/src/collaboration/collaboration.module.ts
+++ b/apps/server/src/collaboration/collaboration.module.ts
@@ -39,7 +39,11 @@ export class CollaborationModule implements OnModuleInit, OnModuleDestroy {
   }
 
   onModuleDestroy(): any {
-    this.collaborationGateway.destroy();
-    this.collabWsAdapter.destroy();
+    if (this.collaborationGateway) {
+      this.collaborationGateway.destroy();
+    }
+    if (this.collabWsAdapter) {
+      this.collabWsAdapter.destroy();
+    }
   }
 }
diff --git a/apps/server/src/kysely/kysely-db.module.ts b/apps/server/src/kysely/database.module.ts
similarity index 56%
rename from apps/server/src/kysely/kysely-db.module.ts
rename to apps/server/src/kysely/database.module.ts
index 39231749..1b049f71 100644
--- a/apps/server/src/kysely/kysely-db.module.ts
+++ b/apps/server/src/kysely/database.module.ts
@@ -1,7 +1,13 @@
-import { Global, Module } from '@nestjs/common';
-import { KyselyModule } from 'nestjs-kysely';
+import {
+  Global,
+  Logger,
+  Module,
+  OnApplicationBootstrap,
+  OnModuleDestroy,
+} from '@nestjs/common';
+import { InjectKysely, KyselyModule } from 'nestjs-kysely';
 import { EnvironmentService } from '../integrations/environment/environment.service';
-import { CamelCasePlugin, LogEvent, PostgresDialect } from 'kysely';
+import { CamelCasePlugin, LogEvent, PostgresDialect, sql } from 'kysely';
 import { Pool, types } from 'pg';
 import { GroupRepo } from '@docmost/db/repos/group/group.repo';
 import { WorkspaceRepo } from '@docmost/db/repos/workspace/workspace.repo';
@@ -13,6 +19,8 @@ import { PageRepo } from './repos/page/page.repo';
 import { CommentRepo } from './repos/comment/comment.repo';
 import { PageHistoryRepo } from './repos/page/page-history.repo';
 import { AttachmentRepo } from './repos/attachment/attachment.repo';
+import { KyselyDB } from '@docmost/db/types/kysely.types';
+import * as process from 'node:process';
 
 // https://github.com/brianc/node-postgres/issues/811
 types.setTypeParser(types.builtins.INT8, (val) => Number(val));
@@ -68,4 +76,50 @@ types.setTypeParser(types.builtins.INT8, (val) => Number(val));
     AttachmentRepo,
   ],
 })
-export class KyselyDbModule {}
+export class DatabaseModule implements OnModuleDestroy, OnApplicationBootstrap {
+  private readonly logger = new Logger(DatabaseModule.name);
+
+  constructor(@InjectKysely() private readonly db: KyselyDB) {}
+
+  async onApplicationBootstrap() {
+    await this.establishConnection();
+  }
+
+  async onModuleDestroy(): Promise<void> {
+    if (this.db) {
+      await this.db.destroy();
+    }
+  }
+
+  async establishConnection() {
+    const retryAttempts = 10;
+    const retryDelay = 3000;
+
+    this.logger.log('Establishing database connection');
+    for (let i = 0; i < retryAttempts; i++) {
+      try {
+        await sql`SELECT 1=1`.execute(this.db);
+        this.logger.log('Database connection successful');
+        break;
+      } catch (err) {
+        if (err['errors']) {
+          this.logger.error(err['errors'][0]);
+        } else {
+          this.logger.error(err);
+        }
+
+        if (i < retryAttempts - 1) {
+          this.logger.log(
+            `Retrying [${i + 1}/${retryAttempts}] in ${retryDelay} ms`,
+          );
+          await new Promise((resolve) => setTimeout(resolve, retryDelay));
+        } else {
+          this.logger.error(
+            `Failed to connect to database after ${retryAttempts} attempts. Exiting...`,
+          );
+          process.exit(1);
+        }
+      }
+    }
+  }
+}