Skip to content

fix: 连接池懒加载+有效性校验,驱动资源清理与友好错误提示 #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wl2027 merged 1 commit into from
Jun 5, 2026
Merged

fix: 连接池懒加载+有效性校验,驱动资源清理与友好错误提示 #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/main/java/com/data/pivot/plugin/config/DataPivotInitializer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,9 @@


import com.data.pivot.plugin.entity.custom.DataPivotStrategyInfo;
import com.data.pivot.plugin.tool.DataSourceDriverUtil;
import com.data.pivot.plugin.tool.DatabaseUtil;
import com.data.pivot.plugin.tool.QueryTool;
import com.data.pivot.plugin.context.DataPivotApplication;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.startup.StartupActivity;
Expand All @@ -16,6 +18,8 @@ public void runActivity(Project project) {
initDataPivotDatabaseInfo(project);
initDataPivotSettingInfo(project);
Disposer.register(project, DatabaseUtil::closeConnections);
Disposer.register(project, QueryTool::closeAllConnections);
Disposer.register(project, DataSourceDriverUtil::deregisterAllDrivers);
}

static void initDefaultStrategy(DataPivotApplication application) {
Expand Down
11 changes: 8 additions & 3 deletions src/main/java/com/data/pivot/plugin/tool/DataGripUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,9 +80,12 @@ private static List<String> getIntrospectionObjectNames(LocalDataSource dataSour
}

private static void collectNamingNames(TreePatternNode node, List<String> names) {
if (node == null) {
return;
}
if (node.naming != null && node.naming.names != null) {
for (ObjectName name : node.naming.names) {
if (name.name != null && !name.name.isBlank() && !names.contains(name.name)) {
if (name != null && name.name != null && !name.name.isBlank() && !names.contains(name.name)) {
names.add(name.name);
}
}
Expand All @@ -91,11 +94,13 @@ private static void collectNamingNames(TreePatternNode node, List<String> names)
return;
}
for (TreePatternNode.Group group : node.groups) {
if (group.children == null) {
if (group == null || group.children == null) {
continue;
}
for (TreePatternNode child : group.children) {
collectNamingNames(child, names);
if (child != null) {
collectNamingNames(child, names);
}
}
}
}
Expand Down
66 changes: 53 additions & 13 deletions src/main/java/com/data/pivot/plugin/tool/DataSourceDriverUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
package com.data.pivot.plugin.tool;

import com.data.pivot.plugin.i18n.DataPivotBundle;
import com.intellij.database.dataSource.DatabaseDriver;
import com.intellij.database.dataSource.LocalDataSource;
import com.intellij.database.dataSource.artifacts.DatabaseArtifactContext;
import com.intellij.database.dataSource.artifacts.DatabaseArtifactDefaultContext;
import com.intellij.openapi.vfs.VfsUtilCore;
import com.intellij.util.ui.classpath.SimpleClasspathElement;

import java.io.IOException;
import java.net.URL;
import java.net.URLClassLoader;
import java.nio.file.Path;
Expand All @@ -26,7 +28,7 @@
import java.util.stream.Collectors;

public final class DataSourceDriverUtil {
private static final ConcurrentMap<String, Driver> REGISTERED_DRIVERS = new ConcurrentHashMap<>();
private static final ConcurrentMap<String, RegisteredDriver> REGISTERED_DRIVERS = new ConcurrentHashMap<>();

private DataSourceDriverUtil() {
}
Expand Down Expand Up @@ -64,32 +66,70 @@ public static void ensureDriverRegistered(String dataSourceId, String driverClas

List<String> roots = classRootUrls == null ? Collections.emptyList() : classRootUrls;
String key = driverClassName + "@" + String.join("|", roots);
REGISTERED_DRIVERS.computeIfAbsent(key, ignored -> registerDriver(dataSourceId, driverClassName, roots));
REGISTERED_DRIVERS.computeIfAbsent(key, ignored -> registerDriver(driverClassName, roots));
}

private static Driver registerDriver(String dataSourceId, String driverClassName, List<String> classRootUrls) {
private static RegisteredDriver registerDriver(String driverClassName, List<String> classRootUrls) {
URLClassLoader classLoader = null;
try {
Driver driver = createDriver(driverClassName, classRootUrls);
Driver driver;
if (classRootUrls.isEmpty()) {
driver = (Driver) Class.forName(driverClassName).getDeclaredConstructor().newInstance();
} else {
classLoader = createClassLoader(classRootUrls);
driver = (Driver) Class.forName(driverClassName, true, classLoader).getDeclaredConstructor().newInstance();
}
Driver shim = new DriverShim(driver);
DriverManager.registerDriver(shim);
return shim;
return new RegisteredDriver(shim, classLoader);
} catch (ClassNotFoundException | NoClassDefFoundError e) {
// 驱动 jar 尚未下载到本地(或 jar 内不含该驱动类),引导用户先在 DataGrip 数据源里下载该驱动。
closeQuietly(classLoader);
throw new IllegalStateException(
DataPivotBundle.message("data.pivot.driver.not.downloaded", driverClassName), e);
} catch (Exception e) {
throw new IllegalStateException("Failed to register database driver for " + dataSourceId + ": " + driverClassName, e);
closeQuietly(classLoader);
throw new IllegalStateException(
DataPivotBundle.message("data.pivot.driver.register.fail", driverClassName, String.valueOf(e.getMessage())), e);
}
}

private static Driver createDriver(String driverClassName, List<String> classRootUrls) throws Exception {
if (classRootUrls.isEmpty()) {
return (Driver) Class.forName(driverClassName).getDeclaredConstructor().newInstance();
}

private static URLClassLoader createClassLoader(List<String> classRootUrls) throws Exception {
List<URL> urls = new ArrayList<>();
for (String rootUrl : classRootUrls) {
urls.add(toUrl(rootUrl));
}
return new URLClassLoader(urls.toArray(new URL[0]), DataSourceDriverUtil.class.getClassLoader());
}

/**
* 释放本工具注册过的驱动 shim 及其 {@link URLClassLoader}。应在项目/插件 dispose 时调用,
* 避免 DriverManager 中的 shim 与类加载器长期驻留;下次查询会按需重新注册。
*/
public static void deregisterAllDrivers() {
for (RegisteredDriver registered : REGISTERED_DRIVERS.values()) {
try {
DriverManager.deregisterDriver(registered.shim());
} catch (SQLException ignored) {
// best-effort cleanup on dispose
}
closeQuietly(registered.classLoader());
}
REGISTERED_DRIVERS.clear();
}

private static void closeQuietly(URLClassLoader classLoader) {
if (classLoader == null) {
return;
}
try {
classLoader.close();
} catch (IOException ignored) {
// best-effort cleanup
}
}

URLClassLoader classLoader = new URLClassLoader(urls.toArray(new URL[0]), DataSourceDriverUtil.class.getClassLoader());
return (Driver) Class.forName(driverClassName, true, classLoader).getDeclaredConstructor().newInstance();
private record RegisteredDriver(Driver shim, URLClassLoader classLoader) {
}

private static URL toUrl(String rootUrl) throws Exception {
Expand Down
154 changes: 127 additions & 27 deletions src/main/java/com/data/pivot/plugin/tool/QueryTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
* 查询工具类，复用 IntelliJ IDEA Database Tools 中已配置的数据源驱动。
Expand All @@ -31,7 +32,8 @@ public class QueryTool {
private static final int LIMIT = 20;
private static final int QUERY_TIMEOUT = 5;
private static final int MAX_CONNECTIONS = 10;
private static final ConcurrentMap<String, BlockingQueue<Connection>> connectionPools = new ConcurrentHashMap<>();
private static final int VALIDATION_TIMEOUT = 2;
private static final ConcurrentMap<String, ConnectionPool> connectionPools = new ConcurrentHashMap<>();

private static Connection getConnection(DatabaseQueryConfig config) throws InterruptedException, SQLException {
if (DBType.MONGO.equals(config.getDbType())) {
Expand All @@ -45,22 +47,12 @@ private static Connection getConnection(DatabaseQueryConfig config) throws Inter
);

String poolKey = config.getDataSourceId() + "@" + config.getUrl();
BlockingQueue<Connection> pool = getOrCreateConnectionPool(poolKey, config);
return pool.poll(QUERY_TIMEOUT, TimeUnit.SECONDS);
ConnectionPool pool = getOrCreateConnectionPool(poolKey, config);
return pool.acquire();
}

private static BlockingQueue<Connection> getOrCreateConnectionPool(String poolKey, DatabaseQueryConfig config) {
return connectionPools.computeIfAbsent(poolKey, id -> {
BlockingQueue<Connection> pool = new LinkedBlockingQueue<>(MAX_CONNECTIONS);
for (int i = 0; i < MAX_CONNECTIONS; i++) {
try {
pool.offer(createConnection(config));
} catch (SQLException e) {
throw new RuntimeException("Failed to create a new connection", e);
}
}
return pool;
});
private static ConnectionPool getOrCreateConnectionPool(String poolKey, DatabaseQueryConfig config) {
return connectionPools.computeIfAbsent(poolKey, id -> new ConnectionPool(config));
}

private static Connection createConnection(DatabaseQueryConfig config) throws SQLException {
Expand Down Expand Up @@ -97,10 +89,15 @@ static String addUrlParameters(String url, String parameters) {
}

private static void releaseConnection(DatabaseQueryConfig config, Connection connection) {
if (connection == null) {
return;
}
String poolKey = config.getDataSourceId() + "@" + config.getUrl();
BlockingQueue<Connection> pool = connectionPools.get(poolKey);
ConnectionPool pool = connectionPools.get(poolKey);
if (pool != null) {
pool.offer(connection);
pool.release(connection);
} else {
ConnectionPool.closeQuietly(connection);
}
}

Expand Down Expand Up @@ -197,18 +194,121 @@ private static List<Map<String, Object>> executeQuery(Connection connection, Str
}

public static void closeAllConnections() {
connectionPools.values().forEach(pool -> {
while (!pool.isEmpty()) {
try {
Connection connection = pool.poll();
if (connection != null) {
connection.close();
connectionPools.values().forEach(ConnectionPool::closeAll);
connectionPools.clear();
}

/**
* 懒加载连接池:连接按需创建(上限 {@link #MAX_CONNECTIONS}),取用时通过 {@link Connection#isValid(int)}
* 校验有效性,失效连接会被丢弃并按需重建。相比一次性预建连接,既避免首次查询因单个连接创建失败而整体抛错,
* 也避免长时间空闲后使用到已被数据库侧断开的 stale 连接。
*/
private static final class ConnectionPool {
private final DatabaseQueryConfig config;
private final BlockingQueue<Connection> idle = new LinkedBlockingQueue<>(MAX_CONNECTIONS);
private final AtomicInteger total = new AtomicInteger(0);

private ConnectionPool(DatabaseQueryConfig config) {
this.config = config;
}

/**
* 取出一个可用连接;池内无空闲且未达上限时新建,达到上限则等待归还,超时返回 {@code null}。
*/
private Connection acquire() throws SQLException, InterruptedException {
long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(QUERY_TIMEOUT);
while (true) {
Connection pooled = idle.poll();
if (pooled != null) {
if (isUsable(pooled)) {
return pooled;
}
discard(pooled);
continue;
}
if (tryReserveSlot()) {
try {
return createConnection(config);
} catch (SQLException e) {
total.decrementAndGet();
throw e;
}
} catch (SQLException ignored) {
// Ignore close failures while disposing pooled connections.
}
long remainingNanos = deadlineNanos - System.nanoTime();
if (remainingNanos <= 0) {
return null;
}
pooled = idle.poll(remainingNanos, TimeUnit.NANOSECONDS);
if (pooled == null) {
return null;
}
if (isUsable(pooled)) {
return pooled;
}
discard(pooled);
}
});
connectionPools.clear();
}

/**
* 归还连接:已关闭或池已满则直接释放,否则放回空闲队列。
*/
private void release(Connection connection) {
if (connection == null) {
return;
}
if (isClosedQuietly(connection) || !idle.offer(connection)) {
discard(connection);
}
}

private boolean tryReserveSlot() {
int current;
do {
current = total.get();
if (current >= MAX_CONNECTIONS) {
return false;
}
} while (!total.compareAndSet(current, current + 1));
return true;
}

private void discard(Connection connection) {
closeQuietly(connection);
total.decrementAndGet();
}

private void closeAll() {
Connection connection;
while ((connection = idle.poll()) != null) {
closeQuietly(connection);
}
total.set(0);
}

private static boolean isUsable(Connection connection) {
try {
return connection != null && !connection.isClosed() && connection.isValid(VALIDATION_TIMEOUT);
} catch (SQLException e) {
return false;
}
}

private static boolean isClosedQuietly(Connection connection) {
try {
return connection.isClosed();
} catch (SQLException e) {
return true;
}
}

private static void closeQuietly(Connection connection) {
try {
if (connection != null) {
connection.close();
}
} catch (SQLException ignored) {
// Ignore close failures while disposing pooled connections.
}
}
}
}
2 changes: 2 additions & 0 deletions src/main/resources/messages/DataPivotBundle.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ data.pivot.notice.setting.null.action = \u6253\u5F00 data-pivot \u914D\u7F6E\u98
data.pivot.notice.connection.null = {0} \u6570\u636E\u5E93\u8FDE\u63A5\u5F02\u5E38,\u8BF7\u68C0\u67E5 DataGrip \u8BE5\u6570\u636E\u6E90\u8FDE\u63A5\u60C5\u51B5\u540E\u5237\u65B0
data.pivot.notice.connection.driver.null = {0} \u6570\u636E\u5E93\u9A71\u52A8\u4E0D\u5B58\u5728
data.pivot.notice.rom.data.error = \u52A0\u8F7D{0}\u6570\u636E\u5931\u8D25,\u5F02\u5E38\u4FE1\u606F{1}
data.pivot.driver.not.downloaded = {0} \u9A71\u52A8\u5C1A\u672A\u4E0B\u8F7D,\u8BF7\u5148\u5728 DataGrip \u6570\u636E\u6E90\u8BBE\u7F6E\u4E2D\u4E0B\u8F7D\u8BE5\u6570\u636E\u5E93\u9A71\u52A8\u540E\u91CD\u8BD5
data.pivot.driver.register.fail = {0} \u9A71\u52A8\u6CE8\u518C\u5931\u8D25:{1}

data.pivot.dialog.setting.repeat = {0} \u5DF2\u5B58\u5728\u5BF9\u5E94\u7684 mapping ,\u4E0D\u53EF\u4EE5\u91CD\u590D
data.pivot.dialog.setting.module = module
Expand Down
2 changes: 2 additions & 0 deletions src/main/resources/messages/DataPivotBundle_en_US.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ data.pivot.notice.setting.null.action = open the data-pivot configuration page
data.pivot.notice.connection.null = {0} The database connection is abnormal, check the DataGrip connection status of the data source and refresh it
data.pivot.notice.connection.driver.null = {0} The database driver does not exist
data.pivot.notice.rom.data.error = Loading {0} data failed, exception message: {1}
data.pivot.driver.not.downloaded = The database driver "{0}" has not been downloaded yet. Open the data source settings in DataGrip / Database Tools, download the driver, then try again.
data.pivot.driver.register.fail = Failed to register the database driver "{0}": {1}

data.pivot.dialog.setting.repeat = {0} The corresponding mapping already exists and cannot be repeated
data.pivot.dialog.setting.module = module
Expand Down
2 changes: 2 additions & 0 deletions src/main/resources/messages/DataPivotBundle_zh_CN.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ data.pivot.notice.setting.null.action = \u6253\u5F00 data-pivot \u914D\u7F6E\u98
data.pivot.notice.connection.null = {0} \u6570\u636E\u5E93\u8FDE\u63A5\u5F02\u5E38,\u8BF7\u68C0\u67E5 DataGrip \u8BE5\u6570\u636E\u6E90\u8FDE\u63A5\u60C5\u51B5\u540E\u5237\u65B0
data.pivot.notice.connection.driver.null = {0} \u6570\u636E\u5E93\u9A71\u52A8\u4E0D\u5B58\u5728
data.pivot.notice.rom.data.error = \u52A0\u8F7D{0}\u6570\u636E\u5931\u8D25,\u5F02\u5E38\u4FE1\u606F{1}
data.pivot.driver.not.downloaded = {0} \u9A71\u52A8\u5C1A\u672A\u4E0B\u8F7D,\u8BF7\u5148\u5728 DataGrip \u6570\u636E\u6E90\u8BBE\u7F6E\u4E2D\u4E0B\u8F7D\u8BE5\u6570\u636E\u5E93\u9A71\u52A8\u540E\u91CD\u8BD5
data.pivot.driver.register.fail = {0} \u9A71\u52A8\u6CE8\u518C\u5931\u8D25:{1}

data.pivot.dialog.setting.repeat = {0} \u5DF2\u5B58\u5728\u5BF9\u5E94\u7684 mapping ,\u4E0D\u53EF\u4EE5\u91CD\u590D
data.pivot.dialog.setting.module = module
Expand Down
Loading