基于Drizzle ORM实现通用单用户查询函数的类型问题

需求 #

• 实现通用函数getUser，基于Drizzle-ORM从数据库查询单个用户
• 用户不存在时抛出错误
• 支持通过第一个参数指定要查询的任意列，返回类型需与所选列严格匹配（非空列保持非空，可空列保持可空）

遇到的问题 #

• 初始实现将UserSchema定义为userTable.$inferSelect类型，编写泛型函数时触发错误：Type 'K' cannot be used to index type...
• 调整泛型约束为keyof UserSchema后代码可运行，但使用Partial<UserSchema>作为返回类型时，原本的非空列会被错误标记为可选类型，不符合需求

示例代码 #

import { pgTable, text } from 'drizzle-orm/pg-core';
import { drizzle } from 'drizzle-orm/node-postgres';
import { Client } from 'pg';

const client = new Client({ /* 数据库配置 */ });
await client.connect();
const db = drizzle(client);

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

// 期望调用方式与返回类型
const user1 = await getUser({ name: userTable.name }, "id-123"); // 期望返回 { name: string | null }
const user2 = await getUser({ id: userTable.id }, "id-123");     // 期望返回 { id: string }

解决方案 #

要实现严格匹配所选列的类型，需利用Drizzle的内置类型工具，结合泛型提取所选列的实际类型，而非使用Partial。以下是正确实现：

import type { InferSelectModel } from 'drizzle-orm';
import type { PgColumn } from 'drizzle-orm/pg-core';

type UserSchema = InferSelectModel<typeof userTable>;

// 泛型约束：所选列的键必须是UserSchema的键，值对应表中的列对象
async function getUser<SelectedColumns extends Record<keyof UserSchema, PgColumn>>(
  columns: SelectedColumns,
  userId: string
): Promise<{ [K in keyof SelectedColumns]: UserSchema[K] }> {
  const result = await db
    .select(columns)
    .from(userTable)
    .where(userTable.id.eq(userId))
    .limit(1);

  if (!result[0]) {
    throw new Error('用户不存在');
  }

  return result[0];
}

说明 #

1. 使用InferSelectModel<typeof userTable>准确获取表的查询返回类型，语义清晰且与userTable.$inferSelect等效
2. 泛型SelectedColumns约束为Record<keyof UserSchema, PgColumn>，确保传入的列对象属于userTable的有效列
3. 返回类型使用映射类型{ [K in keyof SelectedColumns]: UserSchema[K] }，严格提取所选列对应的原始类型，完整保留非空/可空的类型特性
4. 查询后检查结果，不存在则抛出错误

此实现下，调用getUser时传入的列会自动推导正确的返回类型，非空列（如id）保持string类型，可空列（如name）保持string | null类型，完全符合需求。

内容的提问来源于stack exchange，提问作者Ahmet yazıcı