開發 IDEA Plugin 引入探針，基於位元組碼插樁獲取執行SQL

作者：小傅哥
部落格：https://bugstack.cn

沉澱、分享、成長，讓自己和他人都能有所收穫！?

一、前言

片面了！

一月三舟，托爾斯泰說：“多麼偉大的作家，也不過就是在書寫自己的片面而已”。何況是我，何況是我們！

雖然我們不書寫文章，但我們寫需求、寫程式碼、寫註釋，當我們遇到了需要被討論的問題點時，往往變成了爭論點。這個好、那個差、你用的都是啥啥啥！

當你把路走窄了，你所能接受到的新的思路、新的想法、新的視野，以及非常重要的收入，也都會隨之減少。只有橫向對比、參考借鑑、查漏補缺，才能讓你的頭腦中會有更多的思路，無論是在寫程式碼上、還是在理財上、還是在生活上。

二、需求目的

你是否有在使用 IntelliJ IDEA 做開發的過程，需要拿到執行 SQL 語句，複製出來做驗證的時候，總是這樣的語句：SELECT * FROM USER WHERE id = ? AND name = ? 又需要自己把 ? 號 替換成入參值呢？

當然這個需求其實並不大，甚至你還可以使用其他方式解決。那麼在本章節會給你提供一個新的思路，可能你幾乎是沒過的方式進行處理。

那麼在這個章節的案例中我們用到基於 IDEA Plugin 開發能力，把位元組碼插樁探針，基於 Javaagent 的能力，注入到程式碼中。再通過增強後的位元組碼，獲取到 com.mysql.jdbc.PreparedStatement -> executeInternal 執行時的物件，從而拿到可以直接測試的 SQL 語句。

三、案例開發

1. 工程結構

guide-idea-plugin-probe
├── .gradle
├── probe-agent
│   ├── src
│   │   └── main
│   │       └── java
│   │           └── cn.bugstack.guide.idea.plugin
│   │               ├── MonitorMethod.java
│   │               └── PreAgent.java
│   └── build.gradle
└── probe-plugin
│   └── src
│   │   └── main
│   │       ├── java
│   │       │    └── cn.bugstack.guide.idea.plugin
│   │       │        └── utils
│   │       │        │    └── PluginUtil.java
│   │       │        └── PerRun.java
│   │       └── resources
│   │           └── META-INF
│   │                └── plugin.xml
│    └── build.gradle
├── build.gradle    
└── gradle.properties

原始碼獲取：#公眾號：bugstack蟲洞棧 回覆：idea 即可下載全部 IDEA 外掛開發原始碼

在此 IDEA 外掛工程中，工程結構分為2塊：

probe-agent：探針模組，用於編譯打包提供位元組碼增強服務，給 probe-plugin 模組使用
probe-plugin：外掛模組，通過 java.programPatcher 載入位元組碼增強包，獲取並列印執行資料庫操作的 SQL 語句。

2. 位元組碼增強獲取 SQL

此處的位元組碼增強方式，採用的 Byte-Buddy 位元組碼框架，它的使用方式更加簡單，在使用的過程中有些像使用 AOP 的攔截方式一樣，獲取到你需要的資訊。

此外在 gradle 打包構建的時候，需要新增 shadowJar 模組，把 Premain-Class 打包進去。這部分程式碼中可以檢視

2.1 探針入口

cn.bugstack.guide.idea.plugin.PreAgent

//JVM 首先嚐試在代理類上呼叫以下方法
public static void premain(String agentArgs, Instrumentation inst) {
    AgentBuilder.Transformer transformer = (builder, typeDescription, classLoader, javaModule) -> {
        return builder
                .method(ElementMatchers.named("executeInternal")) // 攔截任意方法
                .intercept(MethodDelegation.to(MonitorMethod.class)); // 委託
    };
    new AgentBuilder
            .Default()
            .type(ElementMatchers.nameStartsWith("com.mysql.jdbc.PreparedStatement"))
            .transform(transformer)
            .installOn(inst);
}

通過 Byte-buddy 配置，攔截匹配的類和方法，因為這個類和方法下，可以獲取到完整吃執行 SQL 語句。

2.2 攔截 SQL

cn.bugstack.guide.idea.plugin.MonitorMethod

@RuntimeType
public static Object intercept(@This Object obj, @Origin Method method, @SuperCall Callable<?> callable, @AllArguments Object... args) throws Exception {
    try {
        return callable.call();
    } finally {
        String originalSql = (String) BeanUtil.getFieldValue(obj, "originalSql");
        String replaceSql = ReflectUtil.invoke(obj, "asSql");
        System.out.println("資料庫名稱：Mysql");
        System.out.println("執行緒ID：" + Thread.currentThread().getId());
        System.out.println("時間：" + new Date());
        System.out.println("原始SQL：\r\n" + originalSql);
        System.out.println("替換SQL：\r\n" + replaceSql);
    }
}

攔截方法入參是一種可配置操作，比如 @This Object obj 是為了獲取當前類的執行物件，@Origin Method method 是為了獲取執行方法。
在 finally 塊中，我們可以通過反射拿到當前類的屬性資訊，以及反射拿到執行的 SQL，並做列印輸出。

2.3 編譯打包

在測試和開發 IDEA Plugin 外掛之前，我們需要先進行一個打包操作，這個打包就是把位元組碼增強的程式碼打包整一個 Jar 包。在 build.gradle -> shadowJar

打包編譯後，就可以在 build -> libs 下看到 Jar：probe-agent-1.0-SNAPSHOT-all.jar 這個 Jar 就是用來做位元組碼增強處理的。

2.4 測試驗證

這裡在把寫好的位元組碼增強元件給外掛使用之前，可以做一個測試驗證，避免每次都需要啟動外掛才能做測試。

單元測試

public class ApiTest {

    public static void main(String[] args) throws Exception {

        String URL = "jdbc:mysql://127.0.0.1:3306/itstack?characterEncoding=utf-8";
        String USER = "root";
        String PASSWORD = "123456";
        Class.forName("com.mysql.jdbc.Driver");

        Connection conn = DriverManager.getConnection(URL, USER, PASSWORD);
        String sql="SELECT * FROM USER WHERE id = ? AND name = ?";
        PreparedStatement statement = conn.prepareStatement(sql);
        statement.setLong(1,1L);
        statement.setString(2,"謝飛機");
        ResultSet rs = statement.executeQuery();

        while (rs.next()) {
            System.out.println(rs.getString("name") + " " + rs.getString("address"));
        }

    }

}

VM options：-javaagent:你的路徑\libs\probe-agent-1.0-SNAPSHOT-all.jar
注意在測試執行的時候，你要給 ApiTest 配置 VM options 才能列印攔截 SQL 資訊

測試結果

原始SQL：
SELECT * FROM USER WHERE id = ? AND name = ?
替換SQL：
SELECT * FROM USER WHERE id = 1 AND name = '謝飛機'
謝飛機 北京.大興區.通明湖公園

好啦，這樣我們就可以攔截可以複製執行的 SQL 語句了，接下來我們再做下 IDEA Plugin 的處理。

3. 通過外掛開發引入探針 Jar

接下來我們要把開發好的位元組碼增強 Jar 包，複製到 IDEA Plugin 外掛開發模組中的 libs(可自己建立) 下，之後在 plugin.xml 配置載入 implementation fileTree(dir: 'libs', includes: ['*jar']) 這樣就可以程式中，找到這個 jar 包並配置到程式中。

3.1 複製 jar 到 libs 下

3.2 build.gradle 配置載入

dependencies {
    implementation fileTree(dir: 'libs', includes: ['*jar'])
}

通過 implementation fileTree 引入載入檔案樹的方式，把我們配置好的 Jar 載入到程式執行中。

3.3 程式中引入 javaagent

cn.bugstack.guide.idea.plugin.PerRun

public class PerRun extends JavaProgramPatcher {

    @Override
    public void patchJavaParameters(Executor executor, RunProfile configuration, JavaParameters javaParameters) {

        RunConfiguration runConfiguration = (RunConfiguration) configuration;
        ParametersList vmParametersList = javaParameters.getVMParametersList();
        vmParametersList.addParametersString("-javaagent:" + agentCoreJarPath);
        vmParametersList.addNotEmptyProperty("guide-idea-plugin-probe.projectId", runConfiguration.getProject().getLocationHash());

    }

}

通過繼承 JavaProgramPatcher 類，實現 patchJavaParameters 方法，通過 configuration 屬性來配置我們自己需要被載入的 -javaagent 包。
這樣在通過 IDEA 已經安裝此外掛，執行程式碼的時候，就會執行到這個攔截和列印 SQL 的功能。

3.4 plugin.xml 新增配置

<extensions defaultExtensionNs="com.intellij">
    <!-- Add your extensions here -->
    <java.programPatcher implementation="cn.bugstack.guide.idea.plugin.PerRun"/>
</extensions>

之後你還需要把開發好的載入類，配置到 java.programPatcher 這樣就可以程式執行的時候，被載入到了。

四、測試驗證

準備好一個有資料庫操作的工程，需要的是 JDBC，如果是其他的，你需要自己擴充套件
啟動外掛後，開啟你的工程，執行單元測試，檢視列印區

啟動外掛

如果你是新下載程式碼，那麼可以在 probe-plugin -> Tasks -> intellij -> runIde 中進行執行啟動。