diff --git a/apps/client/src/App.tsx b/apps/client/src/App.tsx index 66c167d..ee2e3c0 100644 --- a/apps/client/src/App.tsx +++ b/apps/client/src/App.tsx @@ -67,6 +67,7 @@ export default function App() { }> } /> + } /> } /> } /> } /> diff --git a/apps/server/src/app.module.ts b/apps/server/src/app.module.ts index 2079097..fcd040c 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 7bda391..7b48edc 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 3923174..1b049f7 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 { + 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); + } + } + } + } +}