Refactor ShardingSphereDriverUtils (#31705)

apache · Jun 14, 2024 · 646468c · 646468c
1 parent 6c9d2ac
commit 646468c
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 20 deletions.
diff --git a/...main/java/org/apache/shardingsphere/agent/plugin/core/util/ShardingSphereDriverUtils.java b/...main/java/org/apache/shardingsphere/agent/plugin/core/util/ShardingSphereDriverUtils.java
@@ -21,16 +21,12 @@
 import lombok.NoArgsConstructor;
 import org.apache.shardingsphere.driver.ShardingSphereDriver;
 import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
-import org.apache.shardingsphere.driver.jdbc.core.driver.DriverDataSourceCache;
 
-import javax.sql.DataSource;
 import java.sql.Driver;
 import java.sql.DriverManager;
 import java.util.Collections;
 import java.util.Enumeration;
-import java.util.LinkedHashMap;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.Optional;
 
 /**
@@ -45,26 +41,14 @@ public final class ShardingSphereDriverUtils {
      * @return found data source
      */
     public static Map<String, ShardingSphereDataSource> findShardingSphereDataSources() {
-        return findShardingSphereDriver().map(ShardingSphereDriverUtils::findShardingSphereDataSources).orElse(Collections.emptyMap());
-    }
-
-    private static Map<String, ShardingSphereDataSource> findShardingSphereDataSources(final Driver driver) {
-        DriverDataSourceCache dataSourceCache = AgentReflectionUtils.getFieldValue(driver, "dataSourceCache");
-        Map<String, DataSource> dataSourceMap = AgentReflectionUtils.getFieldValue(dataSourceCache, "dataSourceMap");
-        Map<String, ShardingSphereDataSource> result = new LinkedHashMap<>(dataSourceMap.size(), 1F);
-        for (Entry<String, DataSource> entry : dataSourceMap.entrySet()) {
-            if (entry.getValue() instanceof ShardingSphereDataSource) {
-                result.put(entry.getKey(), (ShardingSphereDataSource) entry.getValue());
-            }
-        }
-        return result;
+        return findShardingSphereDriver().map(ShardingSphereDriver::getShardingSphereDataSources).orElse(Collections.emptyMap());
     }
 
     @SuppressWarnings("UseOfJDBCDriverClass")
     private static Optional<ShardingSphereDriver> findShardingSphereDriver() {
-        Enumeration<Driver> driverEnumeration = DriverManager.getDrivers();
-        while (driverEnumeration.hasMoreElements()) {
-            Driver driver = driverEnumeration.nextElement();
+        Enumeration<Driver> drivers = DriverManager.getDrivers();
+        while (drivers.hasMoreElements()) {
+            Driver driver = drivers.nextElement();
             if (driver instanceof ShardingSphereDriver) {
                 return Optional.of((ShardingSphereDriver) driver);
             }

diff --git a/jdbc/src/main/java/org/apache/shardingsphere/driver/ShardingSphereDriver.java b/jdbc/src/main/java/org/apache/shardingsphere/driver/ShardingSphereDriver.java
@@ -18,6 +18,7 @@
 package org.apache.shardingsphere.driver;
 
 import org.apache.shardingsphere.driver.exception.DriverRegisterException;
+import org.apache.shardingsphere.driver.jdbc.core.datasource.ShardingSphereDataSource;
 import org.apache.shardingsphere.driver.jdbc.core.driver.DriverDataSourceCache;
 import org.apache.shardingsphere.infra.annotation.HighFrequencyInvocation;
 
@@ -26,12 +27,16 @@
 import java.sql.DriverManager;
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
+import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Properties;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
 /**
  * ShardingSphere driver.
  */
+@SuppressWarnings("UseOfJDBCDriverClass")
 public final class ShardingSphereDriver implements Driver {
 
     private static final String DRIVER_URL_PREFIX = "jdbc:shardingsphere:";
@@ -50,6 +55,16 @@ public final class ShardingSphereDriver implements Driver {
         }
     }
 
+    /**
+     * Get ShardingSphere data sources.
+     *
+     * @return ShardingSphere data source map
+     */
+    public Map<String, ShardingSphereDataSource> getShardingSphereDataSources() {
+        return dataSourceCache.getDataSourceMap().entrySet().stream()
+                .filter(entry -> entry.getValue() instanceof ShardingSphereDataSource).collect(Collectors.toMap(Entry::getKey, entry -> (ShardingSphereDataSource) entry.getValue()));
+    }
+
     @HighFrequencyInvocation(canBeCached = true)
     @Override
     public Connection connect(final String url, final Properties info) throws SQLException {

diff --git a/...rc/main/java/org/apache/shardingsphere/driver/jdbc/core/driver/DriverDataSourceCache.java b/...rc/main/java/org/apache/shardingsphere/driver/jdbc/core/driver/DriverDataSourceCache.java
@@ -17,6 +17,7 @@
 
 package org.apache.shardingsphere.driver.jdbc.core.driver;
 
+import lombok.Getter;
 import org.apache.shardingsphere.driver.api.yaml.YamlShardingSphereDataSourceFactory;
 import org.apache.shardingsphere.infra.url.core.ShardingSphereURL;
 import org.apache.shardingsphere.infra.url.core.ShardingSphereURLLoadEngine;
@@ -30,6 +31,7 @@
 /**
  * Driver data source cache.
  */
+@Getter
 public final class DriverDataSourceCache {
 
     private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap<>();