database connection retry

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);
+        }
+      }
+    }
+  }
+}