Try pg Pool

Author: arjunkomath

bun.lock

@@ -1,11 +1,13 @@
+import { attachDatabasePool } from "@vercel/functions";
 import { sql } from "drizzle-orm";
 import { drizzle } from "drizzle-orm/node-postgres";
 import { err, ok, type Result } from "neverthrow";
+import { Pool } from "pg";
 import type { Database } from "@/drizzle/types";
 import * as schema from "../../drizzle/schema";
 import { getOwner } from "./useOwner";
 
-const dbInstances = new Map<string, Database>();
+const poolInstances = new Map<string, Pool>();
 
 export function getDatabaseName(ownerId: string): Result<string, string> {
 	if (!ownerId.startsWith("org_") && !ownerId.startsWith("user_")) {
@@ -31,19 +33,22 @@ export async function getDatabaseForOwner(ownerId: string): Promise<Database> {
 		},
 	);
 
-	const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
-	const connectionString = `${process.env.DATABASE_URL}/${databaseName}${sslMode}`;
-
-	if (!dbInstances.has(ownerId)) {
-		dbInstances.set(
-			ownerId,
-			drizzle(connectionString, {
-				schema,
-			}),
-		);
+	if (!poolInstances.has(ownerId)) {
+		const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
+		const connectionString = `${process.env.DATABASE_URL}/${databaseName}${sslMode}`;
+
+		const pool = new Pool({
+			connectionString,
+			min: 1,
+			idleTimeoutMillis: 5000,
+			connectionTimeoutMillis: 5000,
+		});
+
+		attachDatabasePool(pool);
+		poolInstances.set(ownerId, pool);
 	}
 
-	return dbInstances.get(ownerId)!;
+	return drizzle(poolInstances.get(ownerId)!, { schema });
 }
 
 export async function deleteDatabase(ownerId: string) {
@@ -54,22 +59,32 @@ export async function deleteDatabase(ownerId: string) {
 		},
 	);
 
-	const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
+	const pool = poolInstances.get(ownerId);
+	if (pool) {
+		await pool.end();
+		poolInstances.delete(ownerId);
+	}
 
-	const ownerDb = drizzle(`${process.env.DATABASE_URL}/manage${sslMode}`, {
-		schema,
+	const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
+	const managePool = new Pool({
+		connectionString: `${process.env.DATABASE_URL}/manage${sslMode}`,
+		min: 1,
+		idleTimeoutMillis: 5000,
+		connectionTimeoutMillis: 5000,
 	});
 
-	// Terminate all connections to the database before dropping
+	const ownerDb = drizzle(managePool, { schema });
+
 	await ownerDb.execute(sql`
 		SELECT pg_terminate_backend(pid)
 		FROM pg_stat_activity
 		WHERE datname = ${databaseName}
 		  AND pid <> pg_backend_pid()
 	`);
 
-	// Drop the database with FORCE option (PostgreSQL 13+)
 	await ownerDb.execute(
 		sql`DROP DATABASE ${sql.identifier(databaseName)} WITH (FORCE)`,
 	);
+
+	await managePool.end();
 }

lib/utils/useDatabase.ts

@@ -1,11 +1,30 @@
+import { attachDatabasePool } from "@vercel/functions";
 import type { UserJSON } from "@clerk/nextjs/server";
 import { drizzle } from "drizzle-orm/node-postgres";
+import { Pool } from "pg";
 import * as schema from "./drizzle/schema";
 import type { OpsDatabase } from "./drizzle/types";
 
+let opsPool: Pool | null = null;
+let opsDb: OpsDatabase | null = null;
+
 export async function getOpsDatabase(): Promise<OpsDatabase> {
-	const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
-	return drizzle(`${process.env.DATABASE_URL}/manage${sslMode}`, { schema });
+	if (!opsDb) {
+		const sslMode = process.env.DATABASE_SSL === "true" ? "?sslmode=require" : "";
+		const connectionString = `${process.env.DATABASE_URL}/manage${sslMode}`;
+
+		opsPool = new Pool({
+			connectionString,
+			min: 1,
+			idleTimeoutMillis: 5000,
+			connectionTimeoutMillis: 5000,
+		});
+
+		attachDatabasePool(opsPool);
+		opsDb = drizzle(opsPool, { schema });
+	}
+
+	return opsDb;
 }
 
 export async function addUserToOpsDb(userData: UserJSON) {

ops/useOps.ts

@@ -52,6 +52,7 @@
     "@upstash/redis": "^1.35.3",
     "@upstash/search": "^0.1.5",
     "@upstash/workflow": "^0.2.16",
+    "@vercel/functions": "^3.1.4",
     "autoprefixer": "10.4.14",
     "caniuse-lite": "^1.0.30001737",
     "class-variance-authority": "^0.7.1",