Skip to content

Commit e533a31

Browse files
Chat2DB-ProChat2DB-Pro
committed
fix exception
1 parent 46e2d52 commit e533a31

5 files changed

Lines changed: 44 additions & 69 deletions

File tree

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/data/source/DataSourceController.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import ai.chat2db.server.domain.api.param.DataSourceUpdateParam;
1313
import ai.chat2db.server.domain.api.service.ConsoleService;
1414
import ai.chat2db.server.domain.api.service.DataSourceService;
15+
import ai.chat2db.server.tools.common.exception.ConnectionException;
1516
import ai.chat2db.spi.model.Database;
1617
import ai.chat2db.spi.ssh.SSHManager;
1718
import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
@@ -99,7 +100,7 @@ public ActionResult sshConnect(@RequestBody SSHTestRequest request) {
99100
session = SSHManager.getSSHSession(sshWebConverter.toInfo(request));
100101
} catch (Exception e) {
101102
log.error("sshConnect error", e);
102-
throw new RuntimeException(e);
103+
throw new ConnectionException("connection.ssh.error",null,e);
103104
} finally {
104105
if (session != null) {
105106
session.disconnect();

‎chat2db-server/chat2db-spi/pom.xml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@
3232
<artifactId>spring-jdbc</artifactId>
3333
</dependency>
3434
<dependency>
35-
<groupId>com.jcraft</groupId>
35+
<groupId>com.github.mwiede</groupId>
3636
<artifactId>jsch</artifactId>
37-
<version>0.1.53</version>
37+
<version>0.2.9</version>
3838
</dependency>
3939
<dependency>
4040
<groupId>com.oracle.ojdbc</groupId>

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,8 @@ private static Connection setConnectInfoThreadLocal(ConnectInfo connectInfo) {
103103
if (session != null) {
104104
url = url.replace(host, "127.0.0.1").replace(port, ssh.getLocalPort());
105105
}
106-
}catch (Exception e){
107-
throw new ConnectionException("connection.ssh.error",null,e);
106+
} catch (Exception e) {
107+
throw new ConnectionException("connection.ssh.error", null, e);
108108
}
109109
try {
110110
DriverConfig config = connectInfo.getDriverConfig();
@@ -133,7 +133,7 @@ private static Connection setConnectInfoThreadLocal(ConnectInfo connectInfo) {
133133
} catch (Exception e) {
134134
}
135135
}
136-
throw new BusinessException("connection.error",null,e1);
136+
throw new BusinessException("connection.error", null, e1);
137137
}
138138
connectInfo.setSession(session);
139139
connectInfo.setConnection(connection);
@@ -167,7 +167,17 @@ public static void removeContext() {
167167
} catch (SQLException e) {
168168
log.error("close connection error", e);
169169
}
170+
170171
CONNECT_INFO_THREAD_LOCAL.remove();
172+
173+
Session session = connectInfo.getSession();
174+
if (session != null && session.isConnected() && connectInfo.getSsh() != null
175+
&& connectInfo.getSsh().isUse()) {
176+
try {
177+
session.delPortForwardingL(Integer.parseInt(connectInfo.getSsh().getLocalPort()));
178+
} catch (JSchException e) {
179+
}
180+
}
171181
}
172182
}
173183

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/ssh/SSHManager.java‎

Lines changed: 22 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,10 @@
11

22
package ai.chat2db.spi.ssh;
33

4-
import java.util.concurrent.ConcurrentHashMap;
5-
4+
import ai.chat2db.server.tools.common.exception.ConnectionException;
65
import ai.chat2db.spi.model.SSHInfo;
76
import cn.hutool.core.net.NetUtil;
8-
import com.jcraft.jsch.JSch;
7+
import cn.hutool.extra.ssh.JschUtil;
98
import com.jcraft.jsch.Session;
109
import lombok.extern.slf4j.Slf4j;
1110
import org.apache.commons.lang3.StringUtils;
@@ -17,72 +16,35 @@
1716
@Slf4j
1817
public class SSHManager {
1918

20-
private static final ConcurrentHashMap<SSHInfo, Session> SSH_SESSION_MAP = new ConcurrentHashMap();
19+
public static Session getSSHSession(SSHInfo ssh) {
20+
Session session;
21+
try {
22+
byte[] passphrase = StringUtils.isNotBlank(ssh.getPassphrase()) ? StringUtils.getBytes(ssh.getPassphrase(),
23+
"UTF-8") : null;
24+
session = JschUtil.getSession(ssh.getHostName(), Integer.parseInt(ssh.getPort()), ssh.getUserName(),
25+
ssh.getKeyFile(), passphrase);
2126

22-
public static Session getSSHSession(SSHInfo sshInfo) {
23-
Session session = SSH_SESSION_MAP.get(sshInfo);
24-
if (session != null && session.isConnected()) {
25-
return session;
26-
} else {
27-
return createSession(sshInfo);
27+
} catch (Exception e) {
28+
throw new ConnectionException("connection.ssh.error", null, e);
2829
}
29-
}
30-
31-
private static Session createSession(SSHInfo ssh) {
32-
synchronized (ssh) {
33-
Session session = SSH_SESSION_MAP.get(ssh);
34-
if (session != null && session.isConnected()) {
35-
return session;
36-
}
30+
if (StringUtils.isNotBlank(ssh.getRHost()) && StringUtils.isNotBlank(ssh.getRPort())) {
3731
try {
38-
JSch jSch = new JSch();
39-
if (StringUtils.isNotBlank(ssh.getKeyFile()) && StringUtils.isNotBlank(ssh.getPassphrase())) {
40-
jSch.addIdentity(ssh.getKeyFile(), ssh.getPassphrase());
41-
}
42-
session = jSch.getSession(ssh.getUserName(), ssh.getHostName(), Integer.parseInt(ssh.getPort()));
43-
if (StringUtils.isBlank(ssh.getKeyFile()) || StringUtils.isBlank(ssh.getPassphrase())) {
44-
session.setPassword(ssh.getPassword());
45-
}
46-
47-
session.setConfig("StrictHostKeyChecking", "no");
48-
session.connect();
49-
SSH_SESSION_MAP.put(ssh, session);
32+
int localPort = !StringUtils.isBlank(ssh.getLocalPort()) ? Integer.parseInt(ssh.getLocalPort())
33+
: NetUtil.getUsableLocalPort();
34+
ssh.setLocalPort(String.valueOf(localPort));
35+
session.setPortForwardingL(localPort, ssh.getRHost(),
36+
Integer.parseInt(ssh.getRPort()));
5037
} catch (Exception e) {
51-
throw new RuntimeException("create ssh session error", e);
52-
}
53-
54-
if (StringUtils.isNotBlank(ssh.getRHost()) && StringUtils.isNotBlank(ssh.getRPort())) {
55-
try {
56-
int localPort = !StringUtils.isBlank(ssh.getLocalPort()) ? Integer.parseInt(ssh.getLocalPort())
57-
: NetUtil.getUsableLocalPort();
58-
session.setPortForwardingL(localPort, ssh.getRHost(),
59-
Integer.parseInt(ssh.getRPort()));
60-
} catch (Exception e) {
61-
if (session != null && session.isConnected()) {
62-
session.disconnect();
63-
SSH_SESSION_MAP.remove(ssh);
64-
}
65-
throw new RuntimeException(ssh.getLocalPort() + " port is used，please change to another port ", e);
38+
if (session != null && session.isConnected()) {
39+
session.disconnect();
6640
}
41+
throw new ConnectionException("connection.ssh.error", null, e);
6742
}
68-
return session;
6943
}
44+
return session;
7045
}
7146

7247
public static void close() {
73-
SSH_SESSION_MAP.forEach((k, v) -> {
74-
if (v != null && v.isConnected()) {
75-
try {
76-
v.delPortForwardingL(Integer.parseInt(k.getLocalPort()));
77-
} catch (Exception e) {
78-
log.error("delPortForwardingL error", e);
79-
}
80-
try {
81-
v.disconnect();
82-
} catch (Exception e) {
83-
log.error("disconnect error", e);
84-
}
85-
}
86-
});
48+
JschUtil.closeAll();
8749
}
8850
}

‎chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/util/JdbcUtils.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,9 @@
2121
import ai.chat2db.spi.model.SSHInfo;
2222
import ai.chat2db.spi.sql.IDriverManager;
2323
import ai.chat2db.spi.ssh.SSHManager;
24-
import com.jcraft.jsch.JSchException;
2524
import com.jcraft.jsch.Session;
2625
import lombok.extern.slf4j.Slf4j;
26+
import org.apache.commons.lang3.StringUtils;
2727

2828
/**
2929
* jdbc工具类
@@ -235,10 +235,12 @@ public static DataSourceConnect testConnect(String url, String host, String port
235235
}
236236
if (session != null) {
237237
try {
238-
session.delPortForwardingL(Integer.parseInt(ssh.getLocalPort()));
238+
if(StringUtils.isNotBlank(ssh.getLocalPort())) {
239+
session.delPortForwardingL(Integer.parseInt(ssh.getLocalPort()));
240+
}
239241
session.disconnect();
242+
} catch (Exception e) {
240243

241-
} catch (JSchException e) {
242244
}
243245
}
244246
}

0 commit comments

Comments
 (0)