From 46049d02ce92b5a91970c45d6a369e755a5d5a76 Mon Sep 17 00:00:00 2001
From: Michael Pellegrini <466696+mpellegrini@users.noreply.github.com>
Date: Tue, 9 Apr 2024 22:34:25 -0400
Subject: [PATCH] feat(packages/drizzledb-pg): drizzle relational queries

* setup drizzle client to accept relations by passing
  in schemas
* definie relations in drizzle schema files
* add new findBy query using relational queries
---
 packages/drizzledb-pg/src/client/pg-client.ts |  4 ++-
 .../drizzledb-pg/src/schema/auth-schema.ts    | 34 +++++++++++++++++++
 packages/drizzledb-pg/src/schema/index.ts     | 25 +-------------
 packages/lucia-auth/src/repository/index.ts   | 21 +++++++++---
 4 files changed, 55 insertions(+), 29 deletions(-)
 create mode 100644 packages/drizzledb-pg/src/schema/auth-schema.ts

diff --git a/packages/drizzledb-pg/src/client/pg-client.ts b/packages/drizzledb-pg/src/client/pg-client.ts
index 77d1333..df9bcd8 100644
--- a/packages/drizzledb-pg/src/client/pg-client.ts
+++ b/packages/drizzledb-pg/src/client/pg-client.ts
@@ -2,6 +2,8 @@ import { drizzle } from 'drizzle-orm/node-postgres'
 import pg from 'pg'
 import { parseEnv, z } from 'znv'
 
+import * as schema from '../schema/index.js'
+
 import 'dotenv/config'
 
 const { DB_CONNECTION_STRING, DB_MAX_CONNECTIONS } = parseEnv(process.env, {
@@ -34,4 +36,4 @@ connection.on('connect', () => console.log(`Connected`))
 connection.on('acquire', () => console.log('acquired a connection from the pool'))
 connection.on('release', () => console.log('released the connection to the pool'))
 
-export const db = drizzle(connection)
+export const db = drizzle(connection, { schema })
diff --git a/packages/drizzledb-pg/src/schema/auth-schema.ts b/packages/drizzledb-pg/src/schema/auth-schema.ts
new file mode 100644
index 0000000..fe18051
--- /dev/null
+++ b/packages/drizzledb-pg/src/schema/auth-schema.ts
@@ -0,0 +1,34 @@
+import { relations } from 'drizzle-orm'
+import { pgSchema, text, timestamp, uuid } from 'drizzle-orm/pg-core'
+
+import { citext } from './custom-types.js'
+
+export const authSchema = pgSchema('auth')
+
+export const users = authSchema.table('users', {
+  id: uuid('id').primaryKey().defaultRandom(),
+  username: citext('username').notNull().unique(),
+  hashedPassword: text('hashed_password').notNull(),
+})
+
+export const sessions = authSchema.table('sessions', {
+  id: text('id').primaryKey(),
+  userId: uuid('user_id')
+    .notNull()
+    .references(() => users.id),
+  expiresAt: timestamp('expires_at', {
+    withTimezone: true,
+    mode: 'date',
+  }).notNull(),
+})
+
+export const usersRelations = relations(users, ({ many }) => ({
+  sessions: many(sessions),
+}))
+
+export const sessionsRelations = relations(sessions, ({ one }) => ({
+  user: one(users, {
+    fields: [sessions.userId],
+    references: [users.id],
+  }),
+}))
diff --git a/packages/drizzledb-pg/src/schema/index.ts b/packages/drizzledb-pg/src/schema/index.ts
index a50b79b..5782596 100644
--- a/packages/drizzledb-pg/src/schema/index.ts
+++ b/packages/drizzledb-pg/src/schema/index.ts
@@ -1,24 +1 @@
-import { pgSchema, text, timestamp, uuid } from 'drizzle-orm/pg-core'
-
-import { citext } from './custom-types.js'
-
-export { count } from 'drizzle-orm'
-
-export const authSchema = pgSchema('auth')
-
-export const users = authSchema.table('users', {
-  id: uuid('id').primaryKey().defaultRandom(),
-  username: citext('username').notNull().unique(),
-  hashedPassword: text('hashed_password').notNull(),
-})
-
-export const sessions = authSchema.table('sessions', {
-  id: text('id').primaryKey(),
-  userId: uuid('user_id')
-    .notNull()
-    .references(() => users.id),
-  expiresAt: timestamp('expires_at', {
-    withTimezone: true,
-    mode: 'date',
-  }).notNull(),
-})
+export * from './auth-schema.js'
diff --git a/packages/lucia-auth/src/repository/index.ts b/packages/lucia-auth/src/repository/index.ts
index e4a618b..c8ca875 100644
--- a/packages/lucia-auth/src/repository/index.ts
+++ b/packages/lucia-auth/src/repository/index.ts
@@ -1,19 +1,32 @@
 import { Argon2id } from 'oslo/password'
 
 import { db } from '@packages/drizzledb-pg/client'
-import * as schema from '@packages/drizzledb-pg/schema'
 import { users } from '@packages/drizzledb-pg/schema'
 
-export type NewUser = Omit<typeof schema.users.$inferInsert, 'id' | 'hashedPassword'> & {
+export type NewUser = Omit<typeof users.$inferInsert, 'id' | 'hashedPassword'> & {
   password: string
 }
 
-export type User = typeof schema.users.$inferSelect
+export type User = typeof users.$inferSelect
+
+export const findByUsername = async (username: string): Promise<User | undefined> => {
+  return db.query.users.findFirst({
+    columns: { id: true, username: true, hashedPassword: true },
+    where: (col, { eq }) => eq(col.username, username),
+  })
+}
 
 export const createUser = async (user: NewUser): Promise<string | undefined> => {
   const hashedPassword = await new Argon2id().hash(user.password)
+
+  await db.query.users.findMany({
+    with: {
+      sessions: true,
+    },
+  })
+
   return db
-    .insert(schema.users)
+    .insert(users)
     .values({ username: user.username, hashedPassword })
     .returning({ id: users.id })
     .then((res) => res[0]?.id)