feat: Added rate limiting and database connection retry logic

Author: m-mdy-m

src/bot/index.ts

@@ -17,6 +17,7 @@ export class CopBot {
   private webhookPath = `/bot/${Config.token}`;
   private webhookURL = `${Config.web_hook}${this.webhookPath}`;
   private mode = this.isProduction ? 'webhook' : 'long-polling';
+  private static latestContext: Context | null = null;
   private constructor() {
     this._bot = new Bot<Context>(Config.token);
   }
@@ -27,6 +28,19 @@ export class CopBot {
     }
     return CopBot.instance;
   }
+  public static setContext(ctx: Context): void {
+    logger.info(`Setting new context: at ${new Date().toISOString()}`);
+    this.latestContext = ctx;
+  }
+
+  public static getContext(): Context | null {
+    if (this.latestContext) {
+      logger.info(`Retrieved latest context: at ${new Date().toISOString()}`);
+    } else {
+      logger.warn('Attempted to retrieve context, but no context is set.');
+    }
+    return this.latestContext;
+  }
   // Stop the bot
   async stop(): Promise<void> {
     if (this.healthCheckInterval) {
@@ -82,6 +96,7 @@ export class CopBot {
         app.listen(port, async () => {
           logger.info(`Webhook server running on port ${port}`);
           await this.setupWebhook(this.webhookURL);
+          logger.info(`Bot started in ${this.mode} mode!`);
         });
       } catch (err: any) {
         console.error('Error setting up webhook:', err);
@@ -100,11 +115,26 @@ export class CopBot {
     new GenerateCommand(this._bot).generate();
     this._bot.use(
       limit({
-        onLimitExceeded: (ctx) => ctx.reply('Too many requests! Please slow down.'),
+        onLimitExceeded: async (ctx: Context, next: () => Promise<void>) => {
+          const waitTime = '1000';
+          const message = `⚠️ Rate limit exceeded. Please wait for ${waitTime} ms before trying again.`;
+
+          try {
+            await ctx.reply(message);
+          } catch (error: any) {
+            logger.error(`Failed to send rate limit message: ${error.message}`);
+          }
+          if (next) {
+            await next();
+          }
+        },
       })
     );
     this._bot.on('my_chat_member', (ctx) => this.handleJoinNewChat(ctx));
-    this._bot.on('message', (ctx) => this.handleMessage(ctx));
+    this._bot.on('message', (ctx) => {
+      CopBot.setContext(ctx);
+      this.handleMessage(ctx);
+    });
     this._bot.catch(async (error: BotError<Context>) => {
       if (error.message.includes('timeout')) {
         await error.ctx.reply('The request took too long to process. Please try again later.');

src/database/ConnectionPool.ts

@@ -10,28 +10,37 @@ export class ConnectionPool {
     this._pool = this.initializePool(connectionString);
   }
   async connect(): Promise<boolean> {
+    let client;
     try {
-      await this._pool.connect();
+      client = await this._pool.connect();
       logger.info('Database connection successful');
       return true;
     } catch (error: any) {
       console.error('Database connection error:', error.message);
+
+      // Handle missing database error (PostgreSQL code: 3D000)
       if (error.code === '3D000') {
         console.log(`Database does not exist. Creating database ${Config.database.databaseName}...`);
-        await this.createDatabase();
-        await this.reinitializePool();
+
         try {
-          const client = await this._pool.connect();
+          await this.createDatabase();
+          await this.reinitializePool();
+          client = await this._pool.connect(); // Retry connection
           logger.info('Database connection successful after reinitialization');
           return true;
         } catch (reconnectError: any) {
-          console.error('Reconnection failed:', reconnectError.message);
+          console.error('Reconnection failed after reinitialization:', reconnectError.message);
           return false;
         }
       } else {
         console.error('Unexpected error connecting to the database:', error);
         return false;
       }
+    } finally {
+      // Release client only if it was successfully acquired
+      if (client) {
+        client.release();
+      }
     }
   }
   private async createDatabase(): Promise<void> {
@@ -68,7 +77,9 @@ export class ConnectionPool {
     });
   }
   async reinitializePool() {
-    await this._pool.end(); // Close old connections
+    if (this._pool) {
+      await this._pool.end();
+    }
     const newConnectionString = this.getConnectionString();
     this._pool = this.initializePool(newConnectionString);
     console.warn('Connection pool reinitialized.');

src/database/service/Database.ts

@@ -8,21 +8,16 @@ export class DatabaseService {
     }
     this._client = client;
   }
-
-  /**
-   * Runs a query with parameters and returns the result.
-   */
   async query<T extends QueryResultRow>(sql: string, params: any[] = []): Promise<QueryResult<T>> {
     try {
       return await this._client.query<T>(sql, params);
     } catch (error: any) {
       console.error(`Error executing query: ${sql}`, error);
       throw new Error(`Database query failed: ${error.message}`);
+    } finally {
+      this._client.release();
     }
   }
-  /**
-   * Inserts a new record and returns the inserted row.
-   */
   async insert<T extends QueryResultRow>(tableName: string, data: Record<string, any>, returning: string[] = ['*']): Promise<T> {
     const columns = Object.keys(data).join(', ');
     const values = Object.values(data);
@@ -32,10 +27,6 @@ export class DatabaseService {
     const result = await this.query<T>(sql, values);
     return result.rows[0];
   }
-
-  /**
-   * Updates a record and returns the updated row.
-   */
   async update<T extends QueryResultRow>(tableName: string, data: Record<string, any>, condition: Record<string, any>, returning: string[] = ['*']): Promise<T> {
     const setClauses = Object.keys(data)
       .map((key, i) => `"${key}" = $${i + 1}`)
@@ -49,10 +40,6 @@ export class DatabaseService {
     const result = await this.query<T>(sql, values);
     return result.rows[0];
   }
-
-  /**
-   * Deletes records and optionally returns the deleted rows.
-   */
   async delete<T extends QueryResultRow>(tableName: string, condition: Record<string, any>, returning: string[] = ['*']): Promise<T[]> {
     const whereClauses = Object.keys(condition)
       .map((key, i) => `"${key}" = $${i + 1}`)

src/service/database/ServiceProvider.ts

@@ -6,11 +6,15 @@ import { UserService } from '../../database/models/User';
 import { GroupRuleService } from '../../database/models/GroupRule';
 import { WarningDatabaseService } from '../../database/models/Warning';
 import logger from '../../utils/logger';
+import { CopBot } from '../../bot';
 
 export class ServiceProvider {
   private static instance: ServiceProvider;
   private _clientInstance: Client;
   private _connectionPool!: ConnectionPool;
+
+  private lastRequestTime: number | null = null; // Track the last request time
+  private readonly requestInterval: number = 5000;
   private constructor() {
     this._clientInstance = new Client();
   }
@@ -27,6 +31,28 @@ export class ServiceProvider {
     }
     return ServiceProvider.instance;
   }
+  private async enforceRateLimit(): Promise<void> {
+    const ctx = CopBot.getContext();
+    const now = Date.now();
+    if (this.lastRequestTime) {
+      const elapsed = now - this.lastRequestTime;
+      if (elapsed < this.requestInterval) {
+        const waitTime = this.requestInterval - elapsed;
+
+        if (ctx) {
+          try {
+            await ctx.reply(`⚠️ Rate limit exceeded. Please wait for ${waitTime} ms before making another request.`);
+          } catch (error: any) {
+            logger.error(`Failed to notify user about rate limit: ${error.message}`);
+          }
+        } else {
+          logger.warn('No active context to send a rate limit message.');
+        }
+        await new Promise((resolve) => setTimeout(resolve, waitTime));
+      }
+    }
+    this.lastRequestTime = Date.now();
+  }
   static getInstance() {
     return ServiceProvider.instance;
   }
@@ -41,6 +67,7 @@ export class ServiceProvider {
     await this._connectionPool.close();
   }
   async getPoolClint(): Promise<PoolClient> {
+    await this.enforceRateLimit();
     try {
       const client = await this._connectionPool.getClient();
       client.on('error', (err: any) => {