チャットデータもdrizzleで扱える

Author: na-trium-144

app/lib/auth.ts

@@ -5,7 +5,9 @@ import { anonymous } from "better-auth/plugins";
 import { migrateChatUser } from "./chatHistory";
 import { getDrizzle } from "./drizzle";
 
-export async function getAuthServer(drizzle: ReturnType<typeof getDrizzle>) {
+export async function getAuthServer(
+  drizzle: Awaited<ReturnType<typeof getDrizzle>>
+) {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   let cloudflareEnv: any;
   try {

app/lib/chatHistory.ts

@@ -1,5 +1,8 @@
 import { headers } from "next/headers";
 import { getAuthServer } from "./auth";
+import { getDrizzle } from "./drizzle";
+import { chat, message } from "@/schema/chat";
+import { and, asc, eq } from "drizzle-orm";
 
 export interface CreateChatMessage {
   role: "user" | "ai" | "error";
@@ -11,66 +14,66 @@ export async function addChat(
   sectionId: string,
   messages: CreateChatMessage[]
 ) {
-  const prisma = await getPrismaClient();
-  const auth = await getAuthServer(prisma);
+  const drizzle = await getDrizzle();
+  const auth = await getAuthServer(drizzle);
   const session = await auth.api.getSession({ headers: await headers() });
   if (!session) {
     throw new Error("Not authenticated");
   }
 
-  return await prisma.chat.create({
-    data: {
+  const [newChat] = await drizzle
+    .insert(chat)
+    .values({
       userId: session.user.id,
       docsId,
       sectionId,
-      messages: {
-        createMany: {
-          data: messages,
-        },
-      },
-    },
-    include: {
-      messages: true,
-    },
-  });
+    })
+    .returning();
+
+  const chatMessages = await drizzle
+    .insert(message)
+    .values(
+      messages.map((msg) => ({
+        chatId: newChat.chatId,
+        role: msg.role,
+        content: msg.content,
+      }))
+    )
+    .returning();
+
+  return {
+    ...newChat,
+    messages: chatMessages,
+  };
 }
 
 export type ChatWithMessages = Awaited<ReturnType<typeof addChat>>;
 
 export async function getChat(docsId: string) {
-  const prisma = await getPrismaClient();
-  const auth = await getAuthServer(prisma);
+  const drizzle = await getDrizzle();
+  const auth = await getAuthServer(drizzle);
   const session = await auth.api.getSession({ headers: await headers() });
   if (!session) {
     return [];
   }
 
-  return await prisma.chat.findMany({
-    where: {
-      userId: session.user.id,
-      docsId,
-    },
-    include: {
+  const chats = await drizzle.query.chat.findMany({
+    where: and(eq(chat.userId, session.user.id), eq(chat.docsId, docsId)),
+    with: {
       messages: {
-        orderBy: {
-          createdAt: "asc",
-        },
+        orderBy: [asc(message.createdAt)],
       },
     },
-    orderBy: {
-      createdAt: "asc",
-    },
+    orderBy: [asc(chat.createdAt)],
   });
+
+  return chats;
 }
 
 export async function migrateChatUser(oldUserId: string, newUserId: string) {
-  const prisma = await getPrismaClient();
-  await prisma.chat.updateMany({
-    where: {
-      userId: oldUserId,
-    },
-    data: {
-      userId: newUserId,
-    },
-  });
+  const drizzle = await getDrizzle();
+  await drizzle
+    .update(chat)
+    .set({ userId: newUserId })
+    .where(eq(chat.userId, oldUserId));
 }

app/lib/drizzle.ts

@@ -2,6 +2,7 @@ import { getCloudflareContext } from "@opennextjs/cloudflare";
 import { drizzle } from "drizzle-orm/node-postgres";
 import { Pool } from "pg";
 import * as authSchema from "../schema/auth";
+import * as chatSchema from "../schema/chat";
 
 export async function getDrizzle() {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
@@ -24,6 +25,7 @@ export async function getDrizzle() {
     client: pool,
     schema: {
       ...authSchema,
+      ...chatSchema,
     },
   });
 }

app/schema/chat.ts

@@ -0,0 +1,29 @@
+import { relations } from "drizzle-orm";
+import { pgTable, text, timestamp, uuid } from "drizzle-orm/pg-core";
+
+export const chat = pgTable("chat", {
+  chatId: uuid("chatId").primaryKey().defaultRandom(),
+  userId: text("userId").notNull(),
+  docsId: text("docsId").notNull(),
+  sectionId: text("sectionId").notNull(),
+  createdAt: timestamp("createdAt").notNull().defaultNow(),
+});
+
+export const message = pgTable("message", {
+  id: uuid("id").primaryKey().defaultRandom(),
+  chatId: uuid("chatId").notNull(),
+  role: text("role").notNull(),
+  content: text("content").notNull(),
+  createdAt: timestamp("createdAt").notNull().defaultNow(),
+});
+
+export const chatRelations = relations(chat, ({ many }) => ({
+  messages: many(message),
+}));
+
+export const messageRelations = relations(message, ({ one }) => ({
+  chat: one(chat, {
+    fields: [message.chatId],
+    references: [chat.chatId],
+  }),
+}));

drizzle/0001_stale_reaper.sql

@@ -0,0 +1,15 @@
+CREATE TABLE "chat" (
+	"chatId" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
+	"userId" text NOT NULL,
+	"docsId" text NOT NULL,
+	"sectionId" text NOT NULL,
+	"createdAt" timestamp DEFAULT now() NOT NULL
+);
+--> statement-breakpoint
+CREATE TABLE "message" (
+	"id" uuid PRIMARY KEY DEFAULT gen_random_uuid() NOT NULL,
+	"chatId" uuid NOT NULL,
+	"role" text NOT NULL,
+	"content" text NOT NULL,
+	"createdAt" timestamp DEFAULT now() NOT NULL
+);