chore: review and improve MongoDB connection handling for serverless runtimes

Author: TanCodeX

lib/mongodb.test.ts

@@ -2,12 +2,27 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import mongoose from 'mongoose';
 import dbConnect from './mongodb';
 
+const { mockMongooseConnection } = vi.hoisted(() => ({
+  mockMongooseConnection: {
+    readyState: 0,
+  },
+}));
+
 vi.mock('mongoose', () => ({
   default: {
     connect: vi.fn(),
+    disconnect: vi.fn(),
+    connection: mockMongooseConnection,
   },
 }));
 
+const setConnectedMongoose = (resolvedValue: typeof mongoose) => {
+  vi.mocked(mongoose.connect).mockImplementation(async () => {
+    mockMongooseConnection.readyState = 1;
+    return resolvedValue;
+  });
+};
+
 describe('dbConnect', () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -17,10 +32,15 @@ describe('dbConnect', () => {
       global.mongoose.conn = null;
       global.mongoose.promise = null;
     }
+
+    delete process.env.NEXT_RUNTIME;
+    delete process.env.MONGODB_URI;
+    mockMongooseConnection.readyState = 0;
   });
 
   afterEach(() => {
     delete process.env.MONGODB_URI;
+    delete process.env.NEXT_RUNTIME;
   });
 
   it('throws an error if MONGODB_URI is not defined', async () => {
@@ -34,12 +54,16 @@ describe('dbConnect', () => {
   it('connects to mongoose and caches the connection', async () => {
     process.env.MONGODB_URI = 'mongodb://localhost:27017/test';
     const mockMongoose = { connection: 'mock' };
-    vi.mocked(mongoose.connect).mockResolvedValue(mockMongoose as unknown as typeof mongoose);
+    setConnectedMongoose(mockMongoose as unknown as typeof mongoose);
 
     const conn1 = await dbConnect();
     expect(mongoose.connect).toHaveBeenCalledTimes(1);
     expect(mongoose.connect).toHaveBeenCalledWith('mongodb://localhost:27017/test', {
       bufferCommands: false,
+      maxPoolSize: 10,
+      minPoolSize: 0,
+      maxIdleTimeMS: 30000,
+      serverSelectionTimeoutMS: 5000,
     });
     expect(conn1).toBe(mockMongoose);
 
@@ -64,12 +88,16 @@ describe('dbConnect', () => {
     process.env.MONGODB_URI = specificUri;
 
     const mockMongoose = { connection: 'mock' };
-    vi.mocked(mongoose.connect).mockResolvedValue(mockMongoose as unknown as typeof mongoose);
+    setConnectedMongoose(mockMongoose as unknown as typeof mongoose);
 
     await dbConnect();
 
     expect(mongoose.connect).toHaveBeenCalledWith(specificUri, {
       bufferCommands: false,
+      maxPoolSize: 10,
+      minPoolSize: 0,
+      maxIdleTimeMS: 30000,
+      serverSelectionTimeoutMS: 5000,
     });
   });
 
@@ -85,4 +113,29 @@ describe('dbConnect', () => {
     // The promise should be cleared so it can try again
     expect(global.mongoose.promise).toBeNull();
   });
+
+  it('throws when called from the Edge runtime', async () => {
+    vi.stubEnv('NEXT_RUNTIME', 'edge');
+    process.env.MONGODB_URI = 'mongodb://localhost:27017/test';
+
+    await expect(dbConnect()).rejects.toThrow(
+      'MongoDB is not supported in the Edge runtime. Use the Node.js runtime.'
+    );
+
+    expect(mongoose.connect).not.toHaveBeenCalled();
+  });
+
+  it('clears a stale cached connection before reconnecting', async () => {
+    process.env.MONGODB_URI = 'mongodb://localhost:27017/test';
+    global.mongoose.conn = {} as typeof mongoose;
+
+    const mockMongoose = { connection: 'mock' };
+    setConnectedMongoose(mockMongoose as unknown as typeof mongoose);
+
+    const conn = await dbConnect();
+
+    expect(mongoose.connect).toHaveBeenCalledTimes(1);
+    expect(global.mongoose.conn).toBe(mockMongoose);
+    expect(conn).toBe(mockMongoose);
+  });
 });

lib/mongodb.ts

@@ -1,6 +1,7 @@
 import mongoose from 'mongoose';
 
 declare global {
+  // Cached across hot reloads and repeated serverless invocations in the same process.
   var mongoose: {
     conn: typeof import('mongoose') | null;
     promise: Promise<typeof import('mongoose')> | null;
@@ -14,10 +15,19 @@ if (!cached) {
 }
 
 async function dbConnect() {
-  if (cached.conn) {
+  if (process.env.NEXT_RUNTIME === 'edge') {
+    throw new Error('MongoDB is not supported in the Edge runtime. Use the Node.js runtime.');
+  }
+
+  if (cached.conn && mongoose.connection.readyState === 1) {
     return cached.conn;
   }
 
+  if (cached.conn && mongoose.connection.readyState !== 1) {
+    cached.conn = null;
+    cached.promise = null;
+  }
+
   if (!cached.promise) {
     const MONGODB_URI = process.env.MONGODB_URI;
 
@@ -27,6 +37,10 @@ async function dbConnect() {
 
     const opts = {
       bufferCommands: false,
+      maxPoolSize: 10,
+      minPoolSize: 0,
+      maxIdleTimeMS: 30000,
+      serverSelectionTimeoutMS: 5000,
     };
 
     cached.promise = mongoose.connect(MONGODB_URI, opts).then((mongoose) => {