opensource/Chat2DB
1
0
Fork 0
mirror of https://github.com/CodePhiliaX/Chat2DB.git synced 2025-09-25 16:13:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

支持自定义驱动

Browse Source
This commit is contained in:
jipengfei-jpf
2023-07-08 16:22:18 +08:00
parent aff25b2eb9
commit 4f5a58e77b
53 changed files with 643 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java
View File

@ -65,7 +65,12 @@ public class Chat2DBContext {
}
public static Connection getConnection() {
return getConnectInfo().getConnection();
ConnectInfo connectInfo = getConnectInfo();
Connection connection = connectInfo.getConnection();
if (connection == null) {
connection = setConnectInfoThreadLocal(connectInfo);
}
return connection;
}
/**
@ -74,64 +79,67 @@ public class Chat2DBContext {
* @param info
*/
public static void putContext(ConnectInfo info) {
ConnectInfo connectInfo = CONNECT_INFO_THREAD_LOCAL.get();
CONNECT_INFO_THREAD_LOCAL.set(info);
if (connectInfo == null) {
setConnectInfoThreadLocal(info);
if (StringUtils.isNotBlank(info.getDatabaseName())) {
PLUGIN_MAP.get(getConnectInfo().getDbType()).getDBManage().connectDatabase(info.getDatabaseName());
}
}
}
private static void setConnectInfoThreadLocal(ConnectInfo connectInfo) {
Session session = null;
Connection connection = null;
SSHInfo ssh = connectInfo.getSsh();
String url = connectInfo.getUrl();
String host = connectInfo.getHost();
String port = connectInfo.getPort() + "";
try {
ssh.setRHost(host);
ssh.setRPort(port);
session = getSession(ssh);
if (session != null) {
url = url.replace(host, "127.0.0.1").replace(port, ssh.getLocalPort());
}
DriverConfig config = connectInfo.getDriverConfig();
if (config == null) {
config = getDefaultDriverConfig(connectInfo.getDbType());
connectInfo.setDriverConfig(config);
}
connection = IDriverManager.getConnection(url, connectInfo.getUser(), connectInfo.getPassword(),
connectInfo.getDriverConfig(), connectInfo.getExtendMap());
} catch (Exception e1) {
log.error("GetConnect error", e1);
private static Connection setConnectInfoThreadLocal(ConnectInfo connectInfo) {
synchronized (connectInfo) {
Connection connection = connectInfo.getConnection();
if (connection != null) {
try {
connection.close();
} catch (SQLException e) {
log.error("session close error", e);
}
return connection;
}
if (session != null) {
try {
session.delPortForwardingL(Integer.parseInt(ssh.getLocalPort()));
} catch (JSchException e) {
log.error("session delPortForwardingL error", e);
Session session = null;
SSHInfo ssh = connectInfo.getSsh();
String url = connectInfo.getUrl();
String host = connectInfo.getHost();
String port = connectInfo.getPort() + "";
try {
ssh.setRHost(host);
ssh.setRPort(port);
session = getSession(ssh);
if (session != null) {
url = url.replace(host, "127.0.0.1").replace(port, ssh.getLocalPort());
}
try {
session.disconnect();
} catch (Exception e) {
log.error("session disconnect error", e);
DriverConfig config = connectInfo.getDriverConfig();
if (config == null) {
config = getDefaultDriverConfig(connectInfo.getDbType());
connectInfo.setDriverConfig(config);
}
connection = IDriverManager.getConnection(url, connectInfo.getUser(), connectInfo.getPassword(),
connectInfo.getDriverConfig(), connectInfo.getExtendMap());
} catch (Exception e1) {
log.error("GetConnect error", e1);
if (connection != null) {
try {
connection.close();
} catch (SQLException e) {
log.error("session close error", e);
}
}
if (session != null) {
try {
session.delPortForwardingL(Integer.parseInt(ssh.getLocalPort()));
} catch (JSchException e) {
log.error("session delPortForwardingL error", e);
}
try {
session.disconnect();
} catch (Exception e) {
log.error("session disconnect error", e);
}
}
throw new RuntimeException("GetConnect error", e1);
}
throw new RuntimeException("GetConnect error", e1);
connectInfo.setSession(session);
connectInfo.setConnection(connection);
if (StringUtils.isNotBlank(connectInfo.getDatabaseName())) {
PLUGIN_MAP.get(getConnectInfo().getDbType()).getDBManage().connectDatabase(
connectInfo.getDatabaseName());
}
return connection;
}
connectInfo.setSession(session);
connectInfo.setConnection(connection);
}
private static Session getSession(SSHInfo ssh) {

10
chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/IDriverManager.java
View File

@ -71,16 +71,14 @@ public class IDriverManager {
if (url == null) {
throw new SQLException("The url cannot be null", "08001");
}
DriverManager.println("DriverManager.getConnection(\"" + url + "\")");
SQLException reason = null;
DriverEntry driverEntry = DRIVER_ENTRY_MAP.get(driver.getName());
DriverEntry driverEntry = DRIVER_ENTRY_MAP.get(driver.getJdbcDriver());
if (driverEntry == null) {
driverEntry = getJDBCDriver(driver);
}
try {
Connection con = driverEntry.getDriver().connect(url, info);
if (con != null) {
DriverManager.println("getConnection returning " + driverEntry.getDriver().getClass().getName());
return con;
}
} catch (SQLException var7) {
@ -116,13 +114,13 @@ public class IDriverManager {
throws SQLException {
synchronized (driver) {
try {
if (DRIVER_ENTRY_MAP.containsKey(driver.getName())) {
return DRIVER_ENTRY_MAP.get(driver.getName());
if (DRIVER_ENTRY_MAP.containsKey(driver.getJdbcDriver())) {
return DRIVER_ENTRY_MAP.get(driver.getJdbcDriver());
}
ClassLoader cl = getClassLoader(driver);
Driver d = (Driver)cl.loadClass(driver.getJdbcDriverClass()).newInstance();
DriverEntry driverEntry = DriverEntry.builder().driverConfig(driver).driver(d).build();
DRIVER_ENTRY_MAP.put(driver.getName(), driverEntry);
DRIVER_ENTRY_MAP.put(driver.getJdbcDriver(), driverEntry);
return driverEntry;
} catch (Exception e) {
log.error("getJDBCDriver error", e);