TypeORMのQueryBuilderを拡張

背景

TypeORMでSELECT ～ FOR UPDATEは、QueryBuilderを使って下記のようにすればいいけど、公式を見ても、SELECT ～ FOR UPDATE OF ＜table_name＞ をやる手段が見つからなかった。

hoge.ts

  ...
  .createQueryBuilder()
  .useTransaction(true)
  .setLock("pessimistic_write")
  ...

ということで、そもそもQueryBuilderを拡張するようなことができるのか、試しにやってみた。
(postgresql前提)

QueryBuilderを拡張

TypeORMのSelectQueryBuilderをextendsして、SelectQbというクラスを作成。

selectQb.ts

import { SelectQueryBuilder } from "typeorm";

export class SelectQb<T> extends SelectQueryBuilder<T> {
  private lockTarget = null;

  constructor(qb: SelectQueryBuilder<T>) {
    super(qb);
  }

  public createLockExpression(): string {
    switch (this.expressionMap.lockMode) {
      case "pessimistic_write":
        return " FOR UPDATE OF " + '"' + this.lockTarget + '"';
      default:
        return super.createLockExpression();
    }
  }

  public setLockOf(lockMode: "pessimistic_write", lockTarget: string): this {
    this.expressionMap.lockMode = lockMode;
    this.lockTarget = lockTarget;
    return this;
  }
}

本家のSelectQueryBuilderの実装を見ながら、必要な機能だけ追加。

setLock(lockMode)の代わりに、一緒にテーブル名も指定できるsetLockOf(lockMode, lockTarget)を追加。

SelectQueryBuilderでSQL実行時に呼び出されるcreateLockExpression()を拡張。

確認

下記のようなコードを書いて、SQLを確認。

test.ts

    ...

    const repository = getRepository(User);
    const qb = repository.createQueryBuilder();

    // 拡張
    const selectQb = new SelectQb(qb)
      .useTransaction(true)
      .setLockOf("pessimistic_write", "User");

    // 確認
    console.log(await selectQb.getSql());

    ...

実行結果はこんな感じ。

SELECT "User"."id" AS "User_id", "User"."firstName" AS "User_firstName", "User"."lastName" AS "User_lastName", "User"."age" AS "User_age" FROM "user" "User" FOR UPDATE OF "User"

とりあえず、それっぽいSQLが発行された。

※使っているentityは以前書いたやつ。
https://qiita.com/yusuke-ka/items/195e6bba4f21a659b424