Drizzleでクエリ数をカウントする
Next.js
Drizzle
リクエストでどのくらいクエリが発行されているかチェックしたかったため、以下を追加した。
@Injectable({
scope: Scope.REQUEST,
})
export class DbService {
public connectionFromPool: mysql.PoolConnection | null = null;
private readonly queryLogger: QueryCounterLogger;
private drizzleClientForWrite: MySql2Database<typeof schema> | null = null;
private drizzleClientForRead: MySql2Database<typeof schema>;
constructor(private readonly dbConnectionService: DbConnectionService) {
this.queryLogger = new QueryCounterLogger();
this.drizzleClientForRead = drizzle(dbConnectionService.dbConnection, {
mode: 'default',
schema,
logger: this.queryLogger,
});
}
drizzle(): MySql2Database<typeof schema> {
if (!this.drizzleClientForWrite) {
return this.drizzleClientForRead;
}
return this.drizzleClientForWrite;
}
async startTransaction(): Promise<void> {
this.connectionFromPool =
await this.dbConnectionService.dbPool.getConnection();
this.drizzleClientForWrite = drizzle(this.connectionFromPool, {
mode: 'default',
schema,
logger: this.queryLogger,
});
await this.connectionFromPool.beginTransaction();
return;
}
async commitTransaction(): Promise<void> {
if (!this.connectionFromPool) {
throw new Error('connection is not initialized');
}
await this.connectionFromPool.commit();
this.connectionFromPool.release();
}
async rollbackTransaction(): Promise<void> {
if (!this.connectionFromPool) {
throw new Error('connection is not initialized');
}
await this.connectionFromPool.rollback();
this.connectionFromPool.release();
}
async executeInTransaction<T>(callback: () => Promise<T>): Promise<T> {
await this.startTransaction();
try {
const result = await callback();
await this.commitTransaction();
return result;
} catch (e) {
await this.rollbackTransaction();
throw e;
}
}
getQueryCount(): number {
return this.queryLogger.getQueryCount();
}
getQueries(): string[] {
return this.queryLogger.getQueries();
}
}
export class QueryCounterLogger implements Logger {
private queryCount = 0;
private queries: string[] = [];
logQuery(_query: string, _params: unknown[]): void {
this.queryCount++;
this.queries.push(_query);
}
getQueryCount(): number {
return this.queryCount;
}
getQueries(): string[] {
return this.queries;
}
reset(): void {
this.queryCount = 0;
}
}