MyBatis経由で実行するSQLにリクエストID入れてみる。ダーティハックだけど。

MyBatis経由で実行される全てのSQLに対して SQLコメントを追加したい気持ちになったので簡単にどうやるのかを書いておく。

stackoverflowにも質問を書いたけどググってたら思いついてしまった。（ほぼコピペだけど）

ちなみにリフレクションでこじ開けるので使用するときには注意が必要だ。

最初に経緯を書いておく。どういうコメントを追加したいかというと以下のようなコメントである。

SELECT * FROM users; /* TraceID: foo-bar-baz */

TraceID/RequestID/CorrelationID、なんでも良いんだけど SQLからトランザクションを追いかけたい。そのために、実行されるSQL全てにSQLコメントを追加したくなった。

この手法はもともと isucon9の優勝者のブログに書かれていてなるほど、これは便利そうだ、となった覚えがある。

ユースケースに関しては紹介したブログにも書かれているがスロークエリログ単体でトランザクションを特定できるのが良い点だと思っている。

リクエストIDはアプリケーション側でログ出してる場合は出てると思うんだけど SQLに書かれていないと、スロークエリログ単体で見た場合に分からない。また、ある程度類推はできるとは思うが、スロークエリだったクエリを投げたのは、どのトランザクションか特定するのは困難になる。

というわけでこの記事で紹介するコードを書くきっかけになった次第である。

で、どう実装するかというと、こうだ。

package com.wreulicke.github.mybatis;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Invocation;
import org.slf4j.MDC;
import org.springframework.jdbc.core.SqlProvider;
import org.springframework.stereotype.Component;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Field;
import java.lang.reflect.Type;
import java.util.Properties;

@Component
public class RequestIdInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        return null;
    }

    @Override
    public Object plugin(Object target) {
        String requestid = MDC.get("requestid");
        if (requestid == null) {
            return target;
        }

        try {
            Type[] types = target.getClass().getGenericInterfaces();
            for(Type type : types) {
                if(type == StatementHandler.class) {
                    StatementHandler sh = (StatementHandler) target;
                    Field field = ReflectionUtils.findField(sh.getClass(), "delegate");
                    ReflectionUtils.makeAccessible(field);
                    Object delegate = ReflectionUtils.getField(field, target);

                    Field boundSqlField = ReflectionUtils.findField(delegate.getClass(), "boundSql");
                    ReflectionUtils.makeAccessible(boundSqlField);
                    BoundSql boundSql = (BoundSql) ReflectionUtils.getField(boundSqlField, delegate);
                    String sql = boundSql.getSql();

                    Field sqlField = ReflectionUtils.findField(boundSql.getClass(), "sql");
                    ReflectionUtils.makeAccessible(sqlField);

                    String newSql = sql + " /* RequestID:" + requestid + "*/";

                    ReflectionUtils.setField(sqlField, boundSql, newSql);
                }
            }
            return target;
        } catch(Exception e) {
            throw new AssertionError(e.getMessage(), e);
        }
    }

    @Override
    public void setProperties(Properties properties) {

    }
}

SQL持ってるフィールドからSQL引っ張ってきて書き換えて上書きするだけである。

終わり。

もうちょっとMyBatis側のAPIでどうにかしたいけどこのぐらいしか思いつかなかった。

ちなみに筆者は mybatis-spring-boot-starterを使っていて Componentに登録するだけで、Interceptorが登録されるようになっているのでそのほかの環境で使う場合は注意してほしい。

参考にしたのは以下のコードだ。 github.com ほとんどコピペである。