如何在Drizzle ORM基类仓库中实现主键查询的类型安全？

问题背景 #

我正在使用Drizzle ORM为PostgreSQL数据库构建一个泛型抽象基类仓库，可适配所有继承PgTable的表，主键通过构造函数动态传入。现有findOne方法实现如下：

async findOne(id: T['$inferSelect'][ID]): Promise<InferSelectModel<T>> {
  if (!id) {
    throw new Error('ID is required');
  }

  const result = await this.db
    .select()
    .from(this.table)
    .where(eq(sql.raw(`${this.primaryKey.toString()}`), id));

  if (!result.length) {
    throw new Error(`No record found for ID: ${id}`);
  }

  return result[0];
}

核心问题 #

当前使用sql.raw(${this.primaryKey.toString()})的写法完全缺乏类型安全性：无法在编译时校验主键列名是否与表中实际列匹配，只有在运行时才会抛出错误，违背了TypeScript严格类型校验的初衷。

补充背景 #

• 构造函数的primaryKey参数类型为ID extends keyof T['$inferSelect']
• 该仓库用于实现通用CRUD操作（findOne、create、update、delete）
• Drizzle ORM不直接支持无sql.raw的动态键.where()操作

构造函数代码参考：

constructor(
  protected readonly db: DrizzleDB,
  protected readonly table: T,
  protected readonly primaryKey: ID,
) {}

需要解决的问题：

1. 如何让主键的.where()子句具备类型安全性？
2. 在Drizzle ORM仓库类中处理动态主键，是否有更优的实现方法或设计模式？

解决方案 #

1. 实现类型安全的.where()子句 #

核心思路是直接操作Drizzle的列对象，而非字符串键名，让TypeScript在编译时完成校验。

方案一：直接传入列对象（推荐）

修改构造函数的primaryKey参数类型，要求传入表的列对象（而非字符串），从根源上确保类型安全：

import { PgTable, Column, InferSelectModel, eq, DrizzleDB } from 'drizzle-orm/pg-core';

abstract class BaseRepository<T extends PgTable, IDColumn extends Column> {
  constructor(
    protected readonly db: DrizzleDB,
    protected readonly table: T,
    protected readonly primaryKey: IDColumn,
  ) {}

  async findOne(id: IDColumn['_']['data']): Promise<InferSelectModel<T>> {
    if (!id) {
      throw new Error('ID is required');
    }

    // 直接使用列对象，编译时自动校验列的合法性
    const result = await this.db
      .select()
      .from(this.table)
      .where(eq(this.primaryKey, id));

    if (!result.length) {
      throw new Error(`No record found for ID: ${id}`);
    }

    return result[0];
  }
}

使用时直接传入表的列实例（比如users.id），TypeScript会在编译时拦截无效列的传入，彻底避免运行时的列名不匹配问题。

方案二：基于字符串键名推导列对象

如果必须传入字符串键名（比如从配置读取），可以通过类型断言结合表的属性获取列对象，同时保留编译时校验：

abstract class BaseRepository<T extends PgTable, ID extends keyof T['$inferSelect']> {
  protected readonly primaryKeyColumn: T[ID];

  constructor(
    protected readonly db: DrizzleDB,
    protected readonly table: T,
    primaryKey: ID,
  ) {
    // 类型断言确保列存在，编译时校验ID是否为表的有效键
    this.primaryKeyColumn = this.table[primaryKey] as T[ID];
  }

  async findOne(id: T['$inferSelect'][ID]): Promise<InferSelectModel<T>> {
    if (!id) {
      throw new Error('ID is required');
    }

    const result = await this.db
      .select()
      .from(this.table)
      .where(eq(this.primaryKeyColumn, id));

    if (!result.length) {
      throw new Error(`No record found for ID: ${id}`);
    }

    return result[0];
  }
}

2. 更优的实现方法与设计模式 #

方法一：自动提取表的主键列

如果你的表都遵循统一的主键规则（比如主键名为id），可以直接基于表定义自动提取主键，无需手动传入：

abstract class BaseRepository<T extends PgTable & { id: Column }> {
  constructor(
    protected readonly db: DrizzleDB,
    protected readonly table: T,
  ) {}

  async findOne(id: T['id']['_']['data']): Promise<InferSelectModel<T>> {
    const result = await this.db
      .select()
      .from(this.table)
      .where(eq(this.table.id, id));

    if (!result.length) {
      throw new Error(`No record found for ID: ${id}`);
    }

    return result[0];
  }
}

如果主键命名不统一，可以利用Drizzle的内置类型提取表的主键：

import { PgTable, PrimaryKey, InferSelectModel, eq, DrizzleDB } from 'drizzle-orm/pg-core';

// 从表定义中提取主键列的类型
type ExtractPrimaryKey<T extends PgTable> = T['_']['primaryKey'][number];

abstract class BaseRepository<T extends PgTable> {
  constructor(
    protected readonly db: DrizzleDB,
    protected readonly table: T,
  ) {}

  async findOne(id: InferSelectModel<T>[ExtractPrimaryKey<T>['name']]): Promise<InferSelectModel<T>> {
    const primaryKeyName = ExtractPrimaryKey<T>['name'] as keyof T;
    const primaryKey = this.table[primaryKeyName];
    
    const result = await this.db
      .select()
      .from(this.table)
      .where(eq(primaryKey as any, id));

    if (!result.length) {
      throw new Error(`No record found for ID: ${id}`);
    }

    return result[0];
  }
}

方法二：工厂模式简化仓库创建

通过工厂函数自动生成仓库实例，避免手动传入主键的重复操作：

function createRepository<T extends PgTable>(db: DrizzleDB, table: T) {
  return new class extends BaseRepository<T> {
    constructor() {
      super(db, table);
    }
  }();
}

// 使用示例
const userRepository = createRepository(db, users);
await userRepository.findOne(1);

方法三：扩展Drizzle查询构建器

直接为Drizzle的表对象扩展通用CRUD方法，无需单独维护仓库类：

import { PgTable, PrimaryKey, InferSelectModel, eq, DrizzleDB } from 'drizzle-orm/pg-core';

type ExtractPrimaryKey<T extends PgTable> = T['_']['primaryKey'][number];

declare module 'drizzle-orm/pg-core' {
  interface PgTable<TName extends string, TColumns extends Record<string, Column>> {
    findOne(this: PgTable<TName, TColumns>, db: DrizzleDB, id: InferSelectModel<this>[ExtractPrimaryKey<this>['name']]): Promise<InferSelectModel<this>>;
  }
}

PgTable.prototype.findOne = async function(this: PgTable, db: DrizzleDB, id) {
  const primaryKeyName = ExtractPrimaryKey<this>['name'] as keyof this;
  const primaryKey = this[primaryKeyName];
  
  const result = await db.select().from(this).where(eq(primaryKey as any, id));
  if (!result.length) throw new Error(`No record found for ID: ${id}`);
  return result[0];
};

// 使用示例
await users.findOne(db, 1);

内容的提问来源于stack exchange，提问作者Irosha Hewage