feat(drizzle): migrate pgvector example to Drizzle ORM (#923)

storage/postgres-pgvector/.env.example

@@ -1,7 +1,3 @@
 POSTGRES_URL=
-POSTGRES_URL_NON_POOLING=
-POSTGRES_USER=
-POSTGRES_HOST=
-POSTGRES_PASSWORD=
-POSTGRES_DATABASE=
-OPENAI_API_KEY=
+OPENAI_API_KEY=
+OPENAI_ORG_ID=

storage/postgres-pgvector/README.md

@@ -1,7 +1,7 @@
 ---
-name: Vercel Postgres + Prisma + pgvector Next.js Starter
+name: Vercel Postgres + Drizzle ORM + pgvector Next.js Starter
 slug: postgres-pgvector
-description: A Next.js template that uses Vercel Postgres as the database, Prisma as the ORM with pgvector to enable vector similarity search, and OpenAI's text-embedding-ada-002 model for embeddings.
+description: A Next.js template that uses Vercel Postgres as the database, Drizzle ORM as the ORM with pgvector to enable vector similarity search, and OpenAI's text-embedding-ada-002 model for embeddings.
 framework: Next.js
 useCase: Starter
 css: Tailwind
@@ -14,9 +14,9 @@ relatedTemplates:
   - postgres-drizzle
 ---
 
-# Vercel Postgres + Prisma + pgvector Next.js Starter
+# Vercel Postgres + Drizzle ORM + pgvector Next.js Starter
 
-A Next.js template that uses [Vercel Postgres](https://vercel.com/postgres) as the database, [Prisma](https://prisma.io/) as the ORM with [pgvector](https://github.com/pgvector/pgvector-node#prisma) to enable vector similarity search, and OpenAI's [`text-embedding-ada-002`](https://platform.openai.com/docs/guides/embeddings) model for embeddings.
+A Next.js template that uses [Vercel Postgres](https://vercel.com/postgres) as the database, [Drizzle ORM](https://orm.drizzle.team/) as the ORM with [pgvector](https://github.com/pgvector/pgvector-node#drizzle-orm) to enable vector similarity search, and OpenAI's [`text-embedding-ada-002`](https://platform.openai.com/docs/guides/embeddings) model for embeddings.
 
 ## Demo
 

storage/postgres-pgvector/app/actions.tsx

@@ -1,30 +1,33 @@
 'use server'
 
-import prisma from '@/lib/prisma'
+import { db } from '@/drizzle/db'
+import { SelectPokemon, pokemons } from '@/drizzle/schema'
 import { openai } from '@/lib/openai'
-import { type Pokemon } from '@prisma/client'
+import { desc, sql, cosineDistance, gt } from 'drizzle-orm'
 import { embed } from 'ai'
 
 export async function searchPokedex(
   query: string
-): Promise<Array<Pokemon & { similarity: number }>> {
+): Promise<Array<Pick<SelectPokemon, 'id' | 'name'> & { similarity: number }>> {
   try {
     if (query.trim().length === 0) return []
 
     const embedding = await generateEmbedding(query)
     const vectorQuery = `[${embedding.join(',')}]`
-    const pokemon = await prisma.$queryRaw`
-      SELECT
-        id,
-        "name",
-        1 - (embedding <=> ${vectorQuery}::vector) as similarity
-      FROM pokemon
-      where 1 - (embedding <=> ${vectorQuery}::vector) > .5
-      ORDER BY  similarity DESC
-      LIMIT 8;
-    `
 
-    return pokemon as Array<Pokemon & { similarity: number }>
+    const similarity = sql<number>`1 - (${cosineDistance(
+      pokemons.embedding,
+      vectorQuery
+    )})`
+
+    const pokemon = await db
+      .select({ id: pokemons.id, name: pokemons.name, similarity })
+      .from(pokemons)
+      .where(gt(similarity, 0.5))
+      .orderBy((t) => desc(t.similarity))
+      .limit(8)
+
+    return pokemon
   } catch (error) {
     console.error(error)
     throw error

storage/postgres-pgvector/app/layout.tsx

@@ -5,7 +5,7 @@ export const metadata = {
   metadataBase: new URL('https://postgres-pgvector.vercel.app'),
   title: 'Vercel Postgres AI-powered Semantic Search Demo',
   description:
-    'A Next.js app that uses Vercel Postgres with pgvector, Prisma, and OpenAI to power a semantic search.',
+    'A Next.js app that uses Vercel Postgres with pgvector, Drizzle ORM, and OpenAI to power a semantic search.',
 }
 
 const inter = Inter({

storage/postgres-pgvector/app/page.tsx

@@ -43,17 +43,17 @@ export default function Home() {
         </Link>{' '}
         semantic search demo with{' '}
         <Link
-          href="https://github.com/pgvector/pgvector-node#prisma"
+          href="https://github.com/pgvector/pgvector-node#drizzle-orm"
           className="font-medium underline underline-offset-4 hover:text-black transition-colors"
         >
           pgvector
         </Link>
         ,{' '}
         <Link
-          href="https://prisma.io"
+          href="https://orm.drizzle.team"
           className="font-medium underline underline-offset-4 hover:text-black transition-colors"
         >
-          Prisma
+          Drizzle ORM
         </Link>
         , and{' '}
         <Link

storage/postgres-pgvector/components/search.tsx

@@ -7,20 +7,22 @@ import {
   CommandItem,
   CommandList,
 } from '@/components/command'
-import { type Pokemon } from '@prisma/client'
+import { SelectPokemon } from '@/drizzle/schema'
 import { useEffect, useState } from 'react'
 import { useDebounce } from 'use-debounce'
 
 export interface SearchProps {
   searchPokedex: (
     content: string
-  ) => Promise<Array<Pokemon & { similarity: number }>>
+  ) => Promise<
+    Array<Pick<SelectPokemon, 'id' | 'name'> & { similarity: number }>
+  >
 }
 
 export function Search({ searchPokedex }: SearchProps) {
   const [query, setQuery] = useState('')
   const [searchResults, setSearchResults] = useState<
-    Array<Pokemon & { similarity?: number }>
+    Array<Pick<SelectPokemon, 'id' | 'name'> & { similarity?: number }>
   >([])
   const [debouncedQuery] = useDebounce(query, 150)
   useEffect(() => {

storage/postgres-pgvector/drizzle.config.ts

@@ -0,0 +1,14 @@
+import { defineConfig } from 'drizzle-kit'
+
+export default defineConfig({
+  schema: './drizzle/schema.ts',
+  out: './drizzle/migrations',
+  dialect: 'postgresql',
+  dbCredentials: {
+    url: process.env.POSTGRES_URL!,
+  },
+  // prints all statements that will be executed: https://orm.drizzle.team/kit-docs/config-reference#verbose
+  // verbose: true,
+  // always ask for your confirmation to execute statements: https://orm.drizzle.team/kit-docs/config-reference#strict
+  // strict: true,
+})

storage/postgres-pgvector/drizzle/db.ts

@@ -0,0 +1,5 @@
+import { sql } from '@vercel/postgres'
+import { drizzle } from 'drizzle-orm/vercel-postgres'
+import * as schema from './schema'
+
+export const db = drizzle(sql, { schema })

storage/postgres-pgvector/drizzle/schema.ts

@@ -0,0 +1,41 @@
+import {
+  boolean,
+  index,
+  integer,
+  pgTable,
+  text,
+  vector,
+} from 'drizzle-orm/pg-core'
+import { randomUUID } from 'crypto'
+
+export const pokemons = pgTable(
+  'pokemon',
+  {
+    id: text('id')
+      .primaryKey()
+      .notNull()
+      .$defaultFn(() => randomUUID()),
+    number: integer('number').notNull(),
+    name: text('name').notNull(),
+    type1: text('type1').notNull(),
+    type2: text('type2'),
+    total: integer('total').notNull(),
+    hp: integer('hp').notNull(),
+    attack: integer('attack').notNull(),
+    defense: integer('defense').notNull(),
+    spAtk: integer('spAtk').notNull(),
+    spDef: integer('spDef').notNull(),
+    speed: integer('speed').notNull(),
+    generation: integer('generation').notNull(),
+    legendary: boolean('legendary').notNull(),
+    embedding: vector('embedding', { dimensions: 1536 }),
+  },
+  (table) => ({
+    embeddingIndex: index().using(
+      'hnsw',
+      table.embedding.op('vector_cosine_ops')
+    ),
+  })
+)
+
+export type SelectPokemon = typeof pokemons.$inferSelect

storage/postgres-pgvector/prisma/seed.ts → storage/postgres-pgvector/drizzle/seed.ts

@@ -1,7 +1,10 @@
-import prisma from '../lib/prisma'
+import 'dotenv/config'
+import { db } from './db'
+import { pokemons } from './schema'
+import { eq } from 'drizzle-orm'
 import { openai } from '../lib/openai'
-import { embed } from 'ai'
 import pokemon from './pokemon-with-embeddings.json'
+import { embed } from 'ai'
 
 if (!process.env.OPENAI_API_KEY) {
   throw new Error('process.env.OPENAI_API_KEY is not defined. Please set it.')
@@ -13,11 +16,10 @@ if (!process.env.POSTGRES_URL) {
 
 async function main() {
   try {
-    const pika = await prisma.pokemon.findFirst({
-      where: {
-        name: 'Pikachu',
-      },
+    const pika = await db.query.pokemons.findFirst({
+      where: (pokemons, { eq }) => eq(pokemons.name, 'Pikachu'),
     })
+
     if (pika) {
       console.log('Pokédex already seeded!')
       return
@@ -35,16 +37,14 @@ async function main() {
     const { embedding, ...p } = record
 
     // Create the pokemon in the database
-    const pokemon = await prisma.pokemon.create({
-      data: p,
-    })
+    const [pokemon] = await db.insert(pokemons).values(p).returning()
 
-    // Add the embedding
-    await prisma.$executeRaw`
-        UPDATE pokemon
-        SET embedding = ${embedding}::vector
-        WHERE id = ${pokemon.id}
-    `
+    await db
+      .update(pokemons)
+      .set({
+        embedding,
+      })
+      .where(eq(pokemons.id, pokemon.id))
 
     console.log(`Added ${pokemon.number} ${pokemon.name}`)
   }
@@ -58,11 +58,11 @@ async function main() {
 }
 main()
   .then(async () => {
-    await prisma.$disconnect()
+    process.exit(0)
   })
   .catch(async (e) => {
     console.error(e)
-    await prisma.$disconnect()
+
     process.exit(1)
   })
 

storage/postgres-pgvector/package.json

@@ -1,41 +1,43 @@
 {
-  "name": "postgres-prisma-pgvector-openai",
+  "name": "postgres-drizzle-orm-pgvector-openai",
   "repository": "https://github.com/vercel/examples.git",
   "license": "MIT",
   "version": "0.0.0",
   "private": true,
   "scripts": {
-    "dev": "prisma generate && next dev",
-    "build": "prisma generate && prisma db push && prisma db seed && next build",
+    "dev": "next dev",
+    "db:seed": "ts-node --compiler-options '{\"module\":\"CommonJS\"}' drizzle/seed.ts",
+    "build": "drizzle-kit push && pnpm db:seed && next build",
     "start": "next start",
     "lint": "next lint"
   },
-  "prisma": {
-    "seed": "ts-node --compiler-options {\"module\":\"CommonJS\"} prisma/seed.ts"
-  },
   "dependencies": {
     "@ai-sdk/openai": "^0.0.14",
-    "@prisma/client": "^5.14.0",
     "@types/node": "20.12.12",
     "@types/react": "18.3.2",
     "@types/react-dom": "18.3.0",
+    "@vercel/postgres": "^0.8.0",
     "ai": "^3.1.14",
     "autoprefixer": "10.4.19",
     "clsx": "^2.1.1",
     "cmdk": "^1.0.0",
+    "drizzle-orm": "^0.31.0",
     "eslint": "9.3.0",
     "eslint-config-next": "14.2.3",
     "gpt3-tokenizer": "^1.1.5",
     "next": "14.2.3",
     "openai": "^4.47.1",
     "postcss": "8.4.38",
-    "prisma": "^5.14.0",
     "react": "18.3.1",
     "react-dom": "18.3.1",
     "tailwind-merge": "^2.3.0",
     "tailwindcss": "3.4.3",
     "ts-node": "^10.9.2",
     "typescript": "5.4.5",
     "use-debounce": "^10.0.0"
+  },
+  "devDependencies": {
+    "dotenv": "^16.4.5",
+    "drizzle-kit": "^0.22.0"
   }
 }