Generic function for drizzle orm select method

Question

I would like to have a generic function to select a single user from the database using drizzle-orm, and throw error if user does not exist. Here is what I want to achieve:

const userTable = pgTable("user", { id: text("id").primaryKey(), name: text("name") });

const user1 = await getUser({ name: userTable.name }, "id-123"); // user1: { name: string | null }
const user2 = await getUser({ id: userTable.id }, "id-123");     // user2: { id: string }

I want any column to be optionally selectable through the first parameter provided, and it should be reflected on the return type.

Here is how I tried to write the generic getUser function:

type UserSchema = typeof userTable.$inferSelect;

async function getUser<T extends { [K in keyof Partial<UserSchema>]: (typeof userTable)[K] }>(
    cols: T,
    id: string,
  ): Promise<{ [K in keyof T]: UserSchema[K] }> {
    const user = (await this.db.select(cols).from(userTable).where(eq(userTable.id, id)))[0];
    if (!user) {
      throw new UserNotFoundError(id);
    }

    return user;
  }

But it says:

Type 'K' cannot be used to index type '{ id: string; name: string | null; }'.

UPDATE:

Here is where I am at right now:

async function getUser<T extends keyof UserSchema>(
    cols: { [K in T]: (typeof userTable)[K] },
    id: string,
  ): Promise<{ [K in T]: UserSchema[K] }> {
    const user = (await this.db.select(cols).from(userTable).where(eq(userTable.id, id)))[0];
    if (!user) {
      throw new UserNotFoundError(id);
    }

    return user;
  }

This does not work as it is, but when I change the return type to Promise<{ [K in T]: Partial<UserSchema>[K] }>, it works, but this isn't what I want. Because this causes all the keys passed to the function to be returned with a possibility of undefined even though the column is not nullable.

Answer 1

I recently answered a very similar question in the github discussion. You got the main idea right (making the selection generic), you just need to translate transform it into the result with some type casting and the help of some of the drizzle own types:

import { pgTable, varchar, text, PgSelectBase } from "drizzle-orm/pg-core"
import { drizzle} from "drizzle-orm/node-postgres"
import { eq, SQL, sql } from "drizzle-orm"

export const users = pgTable('employees', {
  id: text('id'),
  firstName: varchar('first_name', { length: 256 }).notNull(),
  lastName: varchar('last_name', { length: 256 }),
  email: varchar('email', { length: 256 }).unique(),
});

const db = drizzle({} as any)

type UserSelectFields = typeof users['_']['columns']

const getUser = async <P extends Partial<UserSelectFields>>(
  projection: P,
  id: string
) => {
  const rows = await db
    .select(projection)
    .from(users)
    .where(eq(users.id, id))
    .limit(1) as Awaited<PgSelectBase<any, P, "single", any>>;
  
  if (!rows.length) {
    throw new Error()
  }
  
  return rows[0];
}

const myUser = await getUser({ id: users.id, firstName: users.firstName}, "id")

Check this typescript playground to see it in action.