Merge pull request #190 from techulus/main

Reuse DB connection

Author: arjunkomath

lib/utils/useDatabase.ts

@@ -1,26 +1,27 @@
 import { sql } from "drizzle-orm";
 import { drizzle } from "drizzle-orm/node-postgres";
 import { err, ok, type Result } from "neverthrow";
-import { cache } from "react";
 import type { Database } from "@/drizzle/types";
 import * as schema from "../../drizzle/schema";
 import { getOwner } from "./useOwner";
 
+const dbInstances = new Map<string, Database>();
+
 export function getDatabaseName(ownerId: string): Result<string, string> {
 	if (!ownerId.startsWith("org_") && !ownerId.startsWith("user_")) {
 		return err("Invalid owner ID");
 	}
 	return ok(ownerId.toLowerCase().replace(/ /g, "_"));
 }
 
-export const database = cache(async (): Promise<Database> => {
+export async function database(): Promise<Database> {
 	const { ownerId } = await getOwner();
 	if (!ownerId) {
 		throw new Error("Owner ID not found");
 	}
 
 	return getDatabaseForOwner(ownerId);
-});
+}
 
 export async function getDatabaseForOwner(ownerId: string): Promise<Database> {
 	const databaseName = getDatabaseName(ownerId).match(
@@ -31,10 +32,18 @@ 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}`;
 
-	return drizzle(`${process.env.DATABASE_URL}/${databaseName}${sslMode}`, {
-		schema,
-	});
+	if (!dbInstances.has(ownerId)) {
+		dbInstances.set(
+			ownerId,
+			drizzle(connectionString, {
+				schema,
+			}),
+		);
+	}
+
+	return dbInstances.get(ownerId)!;
 }
 
 export async function deleteDatabase(ownerId: string) {

trpc/init.ts

@@ -1,13 +1,12 @@
 import { auth } from "@clerk/nextjs/server";
 import { initTRPC, TRPCError } from "@trpc/server";
-import { cache } from "react";
 import superjson from "superjson";
 import { ZodError } from "zod";
 import { SearchService } from "@/lib/search";
 import { database } from "@/lib/utils/useDatabase";
 import { getOwner, getTimezone } from "@/lib/utils/useOwner";
 
-export const createTRPCContext = cache(async () => {
+export const createTRPCContext = async () => {
 	const timezone = await getTimezone();
 	const { orgSlug, ownerId } = await getOwner();
 	const db = await database();
@@ -18,7 +17,7 @@ export const createTRPCContext = cache(async () => {
 		ownerId,
 		orgSlug,
 	};
-});
+};
 
 export type Context = Awaited<ReturnType<typeof createTRPCContext>>;
 // Avoid exporting the entire t-object