mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-20 00:17:44 +00:00
fix(agents): share runtimes across connection sessions
This commit is contained in:
@@ -106,7 +106,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
Connection conn = requireConnected();
|
||||
uncheckedVoid(() -> beforeQueryExecution(conn, options.getTimeoutSecs()));
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
@@ -122,7 +122,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
Connection conn = requireConnected();
|
||||
uncheckedVoid(() -> beforeQueryExecution(conn, options.getTimeoutSecs()));
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return JdbcExecutor.current().executePage(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
@@ -134,19 +134,19 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult fetchQueryPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize);
|
||||
return JdbcExecutor.current().fetchPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean closeQuerySession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeQuerySession(sessionId);
|
||||
return JdbcExecutor.current().closeQuerySession(sessionId);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
Connection conn = requireConnected();
|
||||
uncheckedVoid(() -> beforeQueryExecution(conn, options.getTimeoutSecs()));
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
@@ -158,12 +158,12 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult fetchTableReadPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchTableReadPage(sessionId, pageSize);
|
||||
return JdbcExecutor.current().fetchTableReadPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean closeTableReadSession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeTableReadSession(sessionId);
|
||||
return JdbcExecutor.current().closeTableReadSession(sessionId);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
package com.dbx.agent;
|
||||
|
||||
final class AgentExecutionContext {
|
||||
private static final ThreadLocal<JdbcExecutor> JDBC_EXECUTOR = new ThreadLocal<>();
|
||||
|
||||
private AgentExecutionContext() {
|
||||
}
|
||||
|
||||
static JdbcExecutor jdbcExecutor() {
|
||||
JdbcExecutor executor = JDBC_EXECUTOR.get();
|
||||
return executor == null ? JdbcExecutor.INSTANCE : executor;
|
||||
}
|
||||
|
||||
static <T> T withJdbcExecutor(JdbcExecutor executor, DatabaseAgent.ThrowingSupplier<T> supplier) throws Exception {
|
||||
JdbcExecutor previous = JDBC_EXECUTOR.get();
|
||||
JDBC_EXECUTOR.set(executor);
|
||||
try {
|
||||
return supplier.get();
|
||||
} finally {
|
||||
if (previous == null) {
|
||||
JDBC_EXECUTOR.remove();
|
||||
} else {
|
||||
JDBC_EXECUTOR.set(previous);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -6,9 +6,14 @@ import java.util.List;
|
||||
|
||||
public final class AgentProtocol {
|
||||
public static final int PROTOCOL_VERSION = 1;
|
||||
public static final int MULTI_SESSION_PROTOCOL_VERSION = 2;
|
||||
|
||||
public static final String METHOD_HANDSHAKE = "handshake";
|
||||
public static final String METHOD_CONNECT = "connect";
|
||||
public static final String METHOD_OPEN_SESSION = "open_session";
|
||||
public static final String METHOD_CLOSE_SESSION = "close_session";
|
||||
public static final String METHOD_VALIDATE_SESSION = "validate_session";
|
||||
public static final String METHOD_CANCEL_SESSION = "cancel_session";
|
||||
public static final String METHOD_TEST_CONNECTION = "test_connection";
|
||||
public static final String METHOD_VALIDATE_CONNECTION = "validate_connection";
|
||||
public static final String METHOD_LIST_DATABASES = "list_databases";
|
||||
@@ -67,6 +72,7 @@ public final class AgentProtocol {
|
||||
public static final String CAPABILITY_TRANSACTION = "transaction";
|
||||
public static final String CAPABILITY_DDL = "ddl";
|
||||
public static final String CAPABILITY_KV = "kv";
|
||||
public static final String CAPABILITY_MULTI_SESSION = "multi_session";
|
||||
|
||||
public static final List<String> CAPABILITIES = Collections.unmodifiableList(Arrays.asList(
|
||||
CAPABILITY_CONNECT,
|
||||
@@ -120,6 +126,20 @@ public final class AgentProtocol {
|
||||
METHOD_SHUTDOWN
|
||||
));
|
||||
|
||||
public static final List<String> MULTI_SESSION_METHODS;
|
||||
|
||||
static {
|
||||
List<String> methods = new java.util.ArrayList<>(COMMON_METHODS);
|
||||
int insertAt = methods.indexOf(METHOD_CONNECT) + 1;
|
||||
methods.addAll(insertAt, Arrays.asList(
|
||||
METHOD_OPEN_SESSION,
|
||||
METHOD_CLOSE_SESSION,
|
||||
METHOD_VALIDATE_SESSION,
|
||||
METHOD_CANCEL_SESSION
|
||||
));
|
||||
MULTI_SESSION_METHODS = Collections.unmodifiableList(methods);
|
||||
}
|
||||
|
||||
public static final List<String> MONGO_LEGACY_METHODS = Collections.unmodifiableList(Arrays.asList(
|
||||
MONGO_METHOD_LIST_DATABASES,
|
||||
MONGO_METHOD_LIST_COLLECTIONS,
|
||||
@@ -151,6 +171,12 @@ public final class AgentProtocol {
|
||||
return new HandshakeResult(PROTOCOL_VERSION, PROTOCOL_VERSION, CAPABILITIES);
|
||||
}
|
||||
|
||||
public static HandshakeResult multiSessionHandshakeResult() {
|
||||
List<String> capabilities = new java.util.ArrayList<>(CAPABILITIES);
|
||||
capabilities.add(CAPABILITY_MULTI_SESSION);
|
||||
return new HandshakeResult(MULTI_SESSION_PROTOCOL_VERSION, MULTI_SESSION_PROTOCOL_VERSION, capabilities);
|
||||
}
|
||||
|
||||
public static final class HandshakeResult {
|
||||
private final int protocolVersion;
|
||||
private final int agentProtocolVersion;
|
||||
|
||||
@@ -47,46 +47,46 @@ public abstract class BaseDatabaseAgent implements DatabaseAgent {
|
||||
@Override
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
Connection conn = requireConnected();
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return JdbcExecutor.current().executePage(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
JdbcExecutor.INSTANCE::defaultResultValue
|
||||
JdbcExecutor.current()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult fetchQueryPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize);
|
||||
return JdbcExecutor.current().fetchPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean closeQuerySession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeQuerySession(sessionId);
|
||||
return JdbcExecutor.current().closeQuerySession(sessionId);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
JdbcExecutor.INSTANCE::defaultResultValue
|
||||
JdbcExecutor.current()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult fetchTableReadPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchTableReadPage(sessionId, pageSize);
|
||||
return JdbcExecutor.current().fetchTableReadPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean closeTableReadSession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeTableReadSession(sessionId);
|
||||
return JdbcExecutor.current().closeTableReadSession(sessionId);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
@@ -102,22 +102,22 @@ public interface DatabaseAgent {
|
||||
if (conn == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return AgentExecutionContext.jdbcExecutor().executePage(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
JdbcExecutor.INSTANCE::defaultResultValue
|
||||
AgentExecutionContext.jdbcExecutor()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
default QueryPageResult fetchQueryPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchPage(sessionId, pageSize);
|
||||
return AgentExecutionContext.jdbcExecutor().fetchPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
default boolean closeQuerySession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeQuerySession(sessionId);
|
||||
return AgentExecutionContext.jdbcExecutor().closeQuerySession(sessionId);
|
||||
}
|
||||
|
||||
default QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
@@ -125,22 +125,22 @@ public interface DatabaseAgent {
|
||||
if (conn == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return AgentExecutionContext.jdbcExecutor().startTableRead(
|
||||
conn,
|
||||
sql,
|
||||
schema,
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
JdbcExecutor.INSTANCE::defaultResultValue
|
||||
AgentExecutionContext.jdbcExecutor()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
default QueryPageResult fetchTableReadPage(String sessionId, int pageSize) {
|
||||
return JdbcExecutor.INSTANCE.fetchTableReadPage(sessionId, pageSize);
|
||||
return AgentExecutionContext.jdbcExecutor().fetchTableReadPage(sessionId, pageSize);
|
||||
}
|
||||
|
||||
default boolean closeTableReadSession(String sessionId) {
|
||||
return JdbcExecutor.INSTANCE.closeTableReadSession(sessionId);
|
||||
return AgentExecutionContext.jdbcExecutor().closeTableReadSession(sessionId);
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -24,8 +24,13 @@ public final class JdbcExecutor {
|
||||
|
||||
private final ConcurrentHashMap<String, QuerySession> sessions = new ConcurrentHashMap<>();
|
||||
private final ConcurrentHashMap<String, QuerySession> tableReadSessions = new ConcurrentHashMap<>();
|
||||
private final java.util.Set<Statement> activeStatements = ConcurrentHashMap.newKeySet();
|
||||
|
||||
private JdbcExecutor() {
|
||||
public JdbcExecutor() {
|
||||
}
|
||||
|
||||
public static JdbcExecutor current() {
|
||||
return AgentExecutionContext.jdbcExecutor();
|
||||
}
|
||||
|
||||
public QueryResult execute(Connection conn, String sql, String schema, Function<String, String> setSchemaSql) {
|
||||
@@ -72,6 +77,8 @@ public final class JdbcExecutor {
|
||||
applySchema(conn, schema, setSchemaSql);
|
||||
|
||||
try (Statement stmt = conn.createStatement()) {
|
||||
activeStatements.add(stmt);
|
||||
try {
|
||||
int effectiveMaxRows = Math.max(maxRows, 1);
|
||||
stmt.setMaxRows(effectiveMaxRows + 1);
|
||||
applyQueryTimeout(stmt, timeoutSecs);
|
||||
@@ -96,6 +103,9 @@ public final class JdbcExecutor {
|
||||
elapsed,
|
||||
false
|
||||
);
|
||||
} finally {
|
||||
activeStatements.remove(stmt);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -199,6 +209,7 @@ public final class JdbcExecutor {
|
||||
applySchema(conn, schema, setSchemaSql);
|
||||
|
||||
Statement stmt = conn.createStatement();
|
||||
activeStatements.add(stmt);
|
||||
try {
|
||||
applyQueryTimeout(stmt, options.getTimeoutSecs());
|
||||
if (options.getFetchSize() != null && options.getFetchSize() > 0) {
|
||||
@@ -210,6 +221,7 @@ public final class JdbcExecutor {
|
||||
long elapsed = System.currentTimeMillis() - start;
|
||||
if (!hasResultSet) {
|
||||
int updateCount = stmt.getUpdateCount();
|
||||
activeStatements.remove(stmt);
|
||||
stmt.close();
|
||||
return new QueryPageResult(
|
||||
Collections.emptyList(),
|
||||
@@ -249,6 +261,7 @@ public final class JdbcExecutor {
|
||||
targetSessions.put(sessionId, session);
|
||||
return readSessionPage(targetSessions, session, options.getPageSize(), elapsed);
|
||||
} catch (Exception e) {
|
||||
activeStatements.remove(stmt);
|
||||
try {
|
||||
stmt.close();
|
||||
} catch (Exception ignored) {
|
||||
@@ -284,6 +297,20 @@ public final class JdbcExecutor {
|
||||
closeAllSessions(tableReadSessions);
|
||||
}
|
||||
|
||||
public void cancelActiveStatements() {
|
||||
for (Statement statement : activeStatements) {
|
||||
try {
|
||||
statement.cancel();
|
||||
} catch (SQLException ignored) {
|
||||
}
|
||||
}
|
||||
// Paged result sets can remain idle between fetch requests, so no
|
||||
// Statement is executing when cancellation arrives. Close those
|
||||
// session-owned cursors as part of the same cancellation boundary.
|
||||
closeAllQuerySessions();
|
||||
closeAllTableReadSessions();
|
||||
}
|
||||
|
||||
public int expireIdleQuerySessions() {
|
||||
return expireIdleQuerySessions(System.currentTimeMillis(), QUERY_SESSION_IDLE_TIMEOUT_MILLIS);
|
||||
}
|
||||
@@ -504,6 +531,7 @@ public final class JdbcExecutor {
|
||||
return false;
|
||||
}
|
||||
synchronized (session) {
|
||||
activeStatements.remove(session.statement);
|
||||
try {
|
||||
session.resultSet.close();
|
||||
} catch (Exception ignored) {
|
||||
|
||||
@@ -23,6 +23,7 @@ public final class JsonRpcServer {
|
||||
private static final long CONNECTION_VALIDATION_INTERVAL_MILLIS = 5_000L;
|
||||
|
||||
private final DatabaseAgent agent;
|
||||
private final JdbcExecutor jdbcExecutor;
|
||||
private final Gson gson = new GsonBuilder()
|
||||
// JDBC DECIMAL/NUMERIC values can exceed JavaScript Number precision after JSON-RPC parsing.
|
||||
.registerTypeAdapter(
|
||||
@@ -39,6 +40,7 @@ public final class JsonRpcServer {
|
||||
|
||||
public JsonRpcServer(DatabaseAgent agent) {
|
||||
this.agent = agent;
|
||||
this.jdbcExecutor = new JdbcExecutor();
|
||||
}
|
||||
|
||||
public void run() {
|
||||
@@ -84,6 +86,14 @@ public final class JsonRpcServer {
|
||||
}
|
||||
}
|
||||
|
||||
Object dispatchForRuntime(String method, JsonObject params) throws Exception {
|
||||
return AgentExecutionContext.withJdbcExecutor(jdbcExecutor, () -> dispatch(method, params));
|
||||
}
|
||||
|
||||
void cancelActiveStatements() {
|
||||
jdbcExecutor.cancelActiveStatements();
|
||||
}
|
||||
|
||||
private Object dispatch(String method, JsonObject params) throws Exception {
|
||||
if (AgentProtocol.METHOD_HANDSHAKE.equals(method)) {
|
||||
return AgentProtocol.handshakeResult();
|
||||
@@ -243,18 +253,17 @@ public final class JsonRpcServer {
|
||||
return agent.executeBatch(statements, stringOrNull(params, "schema"));
|
||||
}
|
||||
if (AgentProtocol.METHOD_DISCONNECT.equals(method)) {
|
||||
JdbcExecutor.INSTANCE.closeAllQuerySessions();
|
||||
JdbcExecutor.INSTANCE.closeAllTableReadSessions();
|
||||
jdbcExecutor.closeAllQuerySessions();
|
||||
jdbcExecutor.closeAllTableReadSessions();
|
||||
agent.disconnect();
|
||||
lastConnectParams = null;
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
|
||||
JdbcExecutor.INSTANCE.closeAllQuerySessions();
|
||||
JdbcExecutor.INSTANCE.closeAllTableReadSessions();
|
||||
jdbcExecutor.closeAllQuerySessions();
|
||||
jdbcExecutor.closeAllTableReadSessions();
|
||||
agent.disconnect();
|
||||
lastConnectParams = null;
|
||||
System.exit(0);
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
throw new IllegalArgumentException("Unknown method: " + method);
|
||||
@@ -300,8 +309,8 @@ public final class JsonRpcServer {
|
||||
return;
|
||||
}
|
||||
|
||||
JdbcExecutor.INSTANCE.closeAllQuerySessions();
|
||||
JdbcExecutor.INSTANCE.closeAllTableReadSessions();
|
||||
jdbcExecutor.closeAllQuerySessions();
|
||||
jdbcExecutor.closeAllTableReadSessions();
|
||||
try {
|
||||
agent.disconnect();
|
||||
} catch (Exception ignored) {
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
package com.dbx.agent;
|
||||
|
||||
import com.google.gson.Gson;
|
||||
import com.google.gson.JsonElement;
|
||||
import com.google.gson.JsonObject;
|
||||
import com.google.gson.JsonParser;
|
||||
|
||||
import java.io.BufferedReader;
|
||||
import java.io.InputStreamReader;
|
||||
import java.util.Collections;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
public final class MultiSessionJsonRpcServer {
|
||||
private static final String LEGACY_SESSION_ID = "__legacy__";
|
||||
private static final int MAX_SESSIONS = 256;
|
||||
|
||||
private final Supplier<? extends DatabaseAgent> agentFactory;
|
||||
private final Map<String, Session> sessions = new ConcurrentHashMap<>();
|
||||
private final ExecutorService requests = Executors.newCachedThreadPool();
|
||||
private final Gson gson = new Gson();
|
||||
private final Object outputLock = new Object();
|
||||
|
||||
public MultiSessionJsonRpcServer(Supplier<? extends DatabaseAgent> agentFactory) {
|
||||
this.agentFactory = agentFactory;
|
||||
}
|
||||
|
||||
public void run() {
|
||||
synchronized (outputLock) {
|
||||
System.out.println("{\"ready\":true}");
|
||||
System.out.flush();
|
||||
}
|
||||
try (BufferedReader reader = new BufferedReader(new InputStreamReader(System.in))) {
|
||||
String line;
|
||||
while ((line = reader.readLine()) != null) {
|
||||
JsonObject request = JsonParser.parseString(line).getAsJsonObject();
|
||||
String method = request.get("method").getAsString();
|
||||
if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
|
||||
writeResponse(handleRequest(request));
|
||||
closeAllSessions();
|
||||
requests.shutdown();
|
||||
return;
|
||||
}
|
||||
requests.submit(() -> writeResponse(handleRequest(request)));
|
||||
}
|
||||
} catch (Exception e) {
|
||||
closeAllSessions();
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
String handleRequest(String line) {
|
||||
return gson.toJson(handleRequest(JsonParser.parseString(line).getAsJsonObject()));
|
||||
}
|
||||
|
||||
private JsonObject handleRequest(JsonObject request) {
|
||||
JsonElement id = request.get("id");
|
||||
String method = request.get("method").getAsString();
|
||||
JsonObject params = request.has("params") && request.get("params").isJsonObject()
|
||||
? request.getAsJsonObject("params")
|
||||
: new JsonObject();
|
||||
JsonObject response = new JsonObject();
|
||||
response.addProperty("jsonrpc", "2.0");
|
||||
response.add("id", id);
|
||||
try {
|
||||
Object result;
|
||||
if (AgentProtocol.METHOD_HANDSHAKE.equals(method)) {
|
||||
result = AgentProtocol.multiSessionHandshakeResult();
|
||||
} else if (AgentProtocol.METHOD_OPEN_SESSION.equals(method)) {
|
||||
result = openSession(requiredSessionId(params), params);
|
||||
} else if (AgentProtocol.METHOD_CLOSE_SESSION.equals(method)) {
|
||||
result = closeSession(requiredSessionId(params));
|
||||
} else if (AgentProtocol.METHOD_VALIDATE_SESSION.equals(method)) {
|
||||
result = session(requiredSessionId(params)).handle("validate_connection", params);
|
||||
} else if (AgentProtocol.METHOD_CANCEL_SESSION.equals(method)) {
|
||||
session(requiredSessionId(params)).cancel();
|
||||
result = Collections.singletonMap("ok", true);
|
||||
} else if (AgentProtocol.METHOD_TEST_CONNECTION.equals(method)) {
|
||||
result = new JsonRpcServer(agentFactory.get()).dispatchForRuntime(method, params);
|
||||
} else if (AgentProtocol.METHOD_CONNECT.equals(method)) {
|
||||
closeSession(LEGACY_SESSION_ID);
|
||||
result = openSession(LEGACY_SESSION_ID, params);
|
||||
} else if (AgentProtocol.METHOD_DISCONNECT.equals(method)) {
|
||||
result = closeSession(LEGACY_SESSION_ID);
|
||||
} else if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
|
||||
result = Collections.singletonMap("ok", true);
|
||||
} else {
|
||||
String sessionId = params.has("agentSessionId") ? params.get("agentSessionId").getAsString() : LEGACY_SESSION_ID;
|
||||
result = session(sessionId).handle(method, params);
|
||||
}
|
||||
response.add("result", gson.toJsonTree(result));
|
||||
} catch (Throwable error) {
|
||||
JsonObject rpcError = new JsonObject();
|
||||
rpcError.addProperty("code", -1);
|
||||
rpcError.addProperty("message", error.getMessage() == null ? error.toString() : error.getMessage());
|
||||
response.add("error", rpcError);
|
||||
}
|
||||
return response;
|
||||
}
|
||||
|
||||
private Object openSession(String sessionId, JsonObject params) throws Exception {
|
||||
if (sessions.size() >= MAX_SESSIONS && !sessions.containsKey(sessionId)) {
|
||||
throw new IllegalStateException("Agent session limit reached: " + MAX_SESSIONS);
|
||||
}
|
||||
Session session = new Session(new JsonRpcServer(agentFactory.get()));
|
||||
Session existing = sessions.putIfAbsent(sessionId, session);
|
||||
if (existing != null) {
|
||||
throw new IllegalStateException("Agent session already exists: " + sessionId);
|
||||
}
|
||||
try {
|
||||
return session.handle(AgentProtocol.METHOD_CONNECT, params);
|
||||
} catch (Exception error) {
|
||||
sessions.remove(sessionId, session);
|
||||
session.close();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private Object closeSession(String sessionId) {
|
||||
Session session = sessions.remove(sessionId);
|
||||
if (session != null) {
|
||||
session.close();
|
||||
}
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
||||
private Session session(String sessionId) {
|
||||
Session session = sessions.get(sessionId);
|
||||
if (session == null) {
|
||||
throw new IllegalStateException("Agent session not found: " + sessionId);
|
||||
}
|
||||
return session;
|
||||
}
|
||||
|
||||
private void closeAllSessions() {
|
||||
for (String sessionId : sessions.keySet()) {
|
||||
closeSession(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
private static String requiredSessionId(JsonObject params) {
|
||||
if (!params.has("agentSessionId") || params.get("agentSessionId").getAsString().trim().isEmpty()) {
|
||||
throw new IllegalArgumentException("agentSessionId is required");
|
||||
}
|
||||
return params.get("agentSessionId").getAsString();
|
||||
}
|
||||
|
||||
private void writeResponse(JsonObject response) {
|
||||
synchronized (outputLock) {
|
||||
System.out.println(gson.toJson(response));
|
||||
System.out.flush();
|
||||
}
|
||||
}
|
||||
|
||||
private static final class Session {
|
||||
private final JsonRpcServer server;
|
||||
|
||||
private Session(JsonRpcServer server) {
|
||||
this.server = server;
|
||||
}
|
||||
|
||||
private synchronized Object handle(String method, JsonObject params) throws Exception {
|
||||
return server.dispatchForRuntime(method, params);
|
||||
}
|
||||
|
||||
private synchronized void close() {
|
||||
try {
|
||||
server.dispatchForRuntime(AgentProtocol.METHOD_DISCONNECT, new JsonObject());
|
||||
} catch (Exception ignored) {
|
||||
}
|
||||
}
|
||||
|
||||
private void cancel() {
|
||||
server.cancelActiveStatements();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -30,7 +30,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -44,7 +44,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return JdbcExecutor.current().executePage(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -56,7 +56,7 @@ public abstract class PostgresLikeAgent extends AbstractJdbcAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
{
|
||||
"protocolVersion": 2,
|
||||
"handshakeMethod": "handshake",
|
||||
"handshakeResponseFields": [
|
||||
"protocolVersion",
|
||||
"agentProtocolVersion",
|
||||
"capabilities"
|
||||
],
|
||||
"allCapabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"kv",
|
||||
"multi_session"
|
||||
],
|
||||
"capabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"multi_session"
|
||||
],
|
||||
"defaultSqlCapabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"multi_session"
|
||||
],
|
||||
"commonMethods": [
|
||||
"handshake",
|
||||
"connect",
|
||||
"open_session",
|
||||
"close_session",
|
||||
"validate_session",
|
||||
"cancel_session",
|
||||
"test_connection",
|
||||
"validate_connection",
|
||||
"list_databases",
|
||||
"list_schemas",
|
||||
"list_tables",
|
||||
"list_objects",
|
||||
"list_data_types",
|
||||
"completion_assistant_search_v1",
|
||||
"get_object_source",
|
||||
"get_table_ddl",
|
||||
"get_columns",
|
||||
"list_indexes",
|
||||
"list_foreign_keys",
|
||||
"list_triggers",
|
||||
"execute_query",
|
||||
"execute_query_page",
|
||||
"fetch_query_page",
|
||||
"close_query_session",
|
||||
"start_table_read",
|
||||
"fetch_table_read_page",
|
||||
"close_table_read_session",
|
||||
"get_explain_info",
|
||||
"execute_batch",
|
||||
"execute_transaction",
|
||||
"disconnect",
|
||||
"shutdown"
|
||||
],
|
||||
"mongoLegacyMethods": [
|
||||
"list_databases",
|
||||
"list_collections",
|
||||
"find_documents",
|
||||
"find_documents_extended_json",
|
||||
"count_documents",
|
||||
"server_version",
|
||||
"create_index",
|
||||
"drop_indexes",
|
||||
"drop_collection",
|
||||
"insert_document",
|
||||
"update_document",
|
||||
"update_documents",
|
||||
"delete_document",
|
||||
"delete_documents"
|
||||
],
|
||||
"kvMethods": [
|
||||
"kv_list_prefix",
|
||||
"kv_get",
|
||||
"kv_put",
|
||||
"kv_delete"
|
||||
],
|
||||
"sessionField": "agentSessionId",
|
||||
"cursorSessionField": "sessionId"
|
||||
}
|
||||
@@ -18,6 +18,9 @@ import java.util.Arrays;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertFalse;
|
||||
@@ -37,7 +40,7 @@ class CommonJavaCompatibilityTest {
|
||||
|
||||
@Test
|
||||
void agentProtocolMatchesContractResource() {
|
||||
JsonObject contract = protocolContract();
|
||||
JsonObject contract = protocolContract("/agent-protocol-v1.json");
|
||||
|
||||
assertEquals(AgentProtocol.PROTOCOL_VERSION, contract.get("protocolVersion").getAsInt());
|
||||
assertEquals(AgentProtocol.METHOD_HANDSHAKE, contract.get("handshakeMethod").getAsString());
|
||||
@@ -53,6 +56,14 @@ class CommonJavaCompatibilityTest {
|
||||
assertEquals(AgentProtocol.KV_METHODS, strings(contract.getAsJsonArray("kvMethods")));
|
||||
}
|
||||
|
||||
@Test
|
||||
void multiSessionAgentProtocolMatchesV2ContractResource() {
|
||||
JsonObject contract = protocolContract("/agent-protocol-v2.json");
|
||||
|
||||
assertEquals(AgentProtocol.MULTI_SESSION_PROTOCOL_VERSION, contract.get("protocolVersion").getAsInt());
|
||||
assertEquals(AgentProtocol.MULTI_SESSION_METHODS, strings(contract.getAsJsonArray("commonMethods")));
|
||||
}
|
||||
|
||||
@Test
|
||||
void jsonRpcServerExposesProtocolHandshake() {
|
||||
JsonRpcServer server = new JsonRpcServer(new MinimalAgent());
|
||||
@@ -72,6 +83,99 @@ class CommonJavaCompatibilityTest {
|
||||
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), "metadata"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void multiSessionServerCreatesAndClosesIndependentAgents() {
|
||||
java.util.List<TrackingAgent> created = new java.util.ArrayList<>();
|
||||
MultiSessionJsonRpcServer server = new MultiSessionJsonRpcServer(() -> {
|
||||
TrackingAgent agent = new TrackingAgent();
|
||||
created.add(agent);
|
||||
return agent;
|
||||
});
|
||||
|
||||
JsonObject handshake = JsonParser.parseString(server.handleRequest(
|
||||
"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"handshake\",\"params\":{}}"
|
||||
)).getAsJsonObject().getAsJsonObject("result");
|
||||
assertEquals(2, handshake.get("protocolVersion").getAsInt());
|
||||
assertTrue(containsCapability(handshake.getAsJsonArray("capabilities"), "multi_session"));
|
||||
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"open_session\",\"params\":{\"agentSessionId\":\"a\"}}");
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"open_session\",\"params\":{\"agentSessionId\":\"b\"}}");
|
||||
assertEquals(2, created.size());
|
||||
assertEquals(1, created.get(0).connectCount);
|
||||
assertEquals(1, created.get(1).connectCount);
|
||||
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"close_session\",\"params\":{\"agentSessionId\":\"a\"}}");
|
||||
assertEquals(1, created.get(0).disconnectCount);
|
||||
assertEquals(0, created.get(1).disconnectCount);
|
||||
}
|
||||
|
||||
@Test
|
||||
void multiSessionCancellationOnlyCancelsTargetSession() throws Exception {
|
||||
List<CancelTrackingAgent> created = Collections.synchronizedList(new ArrayList<>());
|
||||
MultiSessionJsonRpcServer server = new MultiSessionJsonRpcServer(() -> {
|
||||
CancelTrackingAgent agent = new CancelTrackingAgent();
|
||||
created.add(agent);
|
||||
return agent;
|
||||
});
|
||||
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"open_session\",\"params\":{\"agentSessionId\":\"a\"}}");
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"open_session\",\"params\":{\"agentSessionId\":\"b\"}}");
|
||||
|
||||
Thread queryA = new Thread(() -> server.handleRequest(
|
||||
"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"execute_query\",\"params\":{\"agentSessionId\":\"a\",\"sql\":\"SELECT 1\"}}"
|
||||
));
|
||||
Thread queryB = new Thread(() -> server.handleRequest(
|
||||
"{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"execute_query\",\"params\":{\"agentSessionId\":\"b\",\"sql\":\"SELECT 1\"}}"
|
||||
));
|
||||
queryA.start();
|
||||
queryB.start();
|
||||
assertTrue(created.get(0).statementStarted.await(5, TimeUnit.SECONDS));
|
||||
assertTrue(created.get(1).statementStarted.await(5, TimeUnit.SECONDS));
|
||||
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":5,\"method\":\"cancel_session\",\"params\":{\"agentSessionId\":\"a\"}}");
|
||||
assertTrue(created.get(0).canceled.get());
|
||||
assertFalse(created.get(1).canceled.get());
|
||||
|
||||
created.get(1).release.countDown();
|
||||
queryA.join(5_000L);
|
||||
queryB.join(5_000L);
|
||||
assertFalse(queryA.isAlive());
|
||||
assertFalse(queryB.isAlive());
|
||||
}
|
||||
|
||||
@Test
|
||||
void cancelActiveStatementsClosesOnlyThatExecutorPagedSessions() {
|
||||
JdbcExecutor target = new JdbcExecutor();
|
||||
JdbcExecutor other = new JdbcExecutor();
|
||||
List<String> targetCalls = new ArrayList<>();
|
||||
List<String> otherCalls = new ArrayList<>();
|
||||
|
||||
QueryPageResult targetPage = target.executePage(
|
||||
pagedConnection(targetCalls),
|
||||
"SELECT 1",
|
||||
null,
|
||||
ignored -> null,
|
||||
new QueryPageOptions(1, null, 10, 0),
|
||||
target::defaultResultValue
|
||||
);
|
||||
QueryPageResult otherPage = other.executePage(
|
||||
pagedConnection(otherCalls),
|
||||
"SELECT 1",
|
||||
null,
|
||||
ignored -> null,
|
||||
new QueryPageOptions(1, null, 10, 0),
|
||||
other::defaultResultValue
|
||||
);
|
||||
|
||||
target.cancelActiveStatements();
|
||||
|
||||
assertThrows(IllegalArgumentException.class, () -> target.fetchPage(targetPage.getSession_id(), 1));
|
||||
assertNotNull(other.fetchPage(otherPage.getSession_id(), 1));
|
||||
assertTrue(targetCalls.contains("resultSet.close"));
|
||||
assertTrue(targetCalls.contains("statement.close"));
|
||||
assertFalse(otherCalls.contains("resultSet.close"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void jsonRpcServerSerializesArbitraryPrecisionNumbersAsStrings() {
|
||||
JsonRpcServer server = new JsonRpcServer(new PreciseNumberAgent());
|
||||
@@ -464,6 +568,69 @@ class CommonJavaCompatibilityTest {
|
||||
}
|
||||
}
|
||||
|
||||
private static final class TrackingAgent extends MinimalAgent {
|
||||
private int connectCount;
|
||||
private int disconnectCount;
|
||||
|
||||
@Override
|
||||
public void connect(ConnectParams params) {
|
||||
connectCount += 1;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void disconnect() {
|
||||
disconnectCount += 1;
|
||||
}
|
||||
}
|
||||
|
||||
private static final class CancelTrackingAgent extends MinimalAgent {
|
||||
private final CountDownLatch statementStarted = new CountDownLatch(1);
|
||||
private final CountDownLatch release = new CountDownLatch(1);
|
||||
private final AtomicBoolean canceled = new AtomicBoolean();
|
||||
private final Connection connection = proxy(Connection.class, (method, args) -> {
|
||||
if ("createStatement".equals(method.getName())) {
|
||||
return proxy(java.sql.Statement.class, (statementMethod, statementArgs) -> {
|
||||
if ("execute".equals(statementMethod.getName())) {
|
||||
statementStarted.countDown();
|
||||
try {
|
||||
release.await(5, TimeUnit.SECONDS);
|
||||
} catch (InterruptedException interrupted) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new RuntimeException(interrupted);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
if ("cancel".equals(statementMethod.getName())) {
|
||||
canceled.set(true);
|
||||
release.countDown();
|
||||
return null;
|
||||
}
|
||||
return defaultValue(statementMethod.getReturnType());
|
||||
});
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.current().execute(
|
||||
connection,
|
||||
sql,
|
||||
schema,
|
||||
ignored -> null,
|
||||
options.getMaxRows(),
|
||||
options.getFetchSize(),
|
||||
options.getTimeoutSecs(),
|
||||
JdbcExecutor.current()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
@Override
|
||||
public Connection getConnection() {
|
||||
return connection;
|
||||
}
|
||||
}
|
||||
|
||||
private static final class TransactionAgent extends MinimalAgent {
|
||||
private final Connection connection;
|
||||
|
||||
@@ -710,6 +877,41 @@ class CommonJavaCompatibilityTest {
|
||||
});
|
||||
}
|
||||
|
||||
private static Connection pagedConnection(List<String> calls) {
|
||||
java.sql.ResultSetMetaData metadata = proxy(java.sql.ResultSetMetaData.class, (method, args) -> {
|
||||
if ("getColumnCount".equals(method.getName())) return 1;
|
||||
if ("getColumnLabel".equals(method.getName())) return "value";
|
||||
if ("getColumnType".equals(method.getName())) return java.sql.Types.INTEGER;
|
||||
if ("getColumnTypeName".equals(method.getName())) return "integer";
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
int[] row = {0};
|
||||
java.sql.ResultSet resultSet = proxy(java.sql.ResultSet.class, (method, args) -> {
|
||||
if ("getMetaData".equals(method.getName())) return metadata;
|
||||
if ("next".equals(method.getName())) return ++row[0] <= 3;
|
||||
if ("getInt".equals(method.getName())) return row[0];
|
||||
if ("wasNull".equals(method.getName())) return false;
|
||||
if ("close".equals(method.getName())) {
|
||||
calls.add("resultSet.close");
|
||||
return null;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
java.sql.Statement statement = proxy(java.sql.Statement.class, (method, args) -> {
|
||||
if ("execute".equals(method.getName())) return true;
|
||||
if ("getResultSet".equals(method.getName())) return resultSet;
|
||||
if ("close".equals(method.getName())) {
|
||||
calls.add("statement.close");
|
||||
return null;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
return proxy(Connection.class, (method, args) -> {
|
||||
if ("createStatement".equals(method.getName())) return statement;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static <T> T proxy(Class<T> type, MethodHandler handler) {
|
||||
InvocationHandler invocationHandler = new InvocationHandler() {
|
||||
@Override
|
||||
@@ -742,10 +944,10 @@ class CommonJavaCompatibilityTest {
|
||||
return false;
|
||||
}
|
||||
|
||||
private static JsonObject protocolContract() {
|
||||
InputStream stream = CommonJavaCompatibilityTest.class.getResourceAsStream("/agent-protocol-v1.json");
|
||||
private static JsonObject protocolContract(String resourcePath) {
|
||||
InputStream stream = CommonJavaCompatibilityTest.class.getResourceAsStream(resourcePath);
|
||||
if (stream == null) {
|
||||
throw new AssertionError("agent-protocol-v1.json resource missing");
|
||||
throw new AssertionError(resourcePath + " resource missing");
|
||||
}
|
||||
return JsonParser.parseReader(new InputStreamReader(stream, StandardCharsets.UTF_8)).getAsJsonObject();
|
||||
}
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
# Agent Protocol v2: Multi-session runtimes
|
||||
|
||||
Protocol v2 allows one Agent process to serve multiple isolated database sessions. The handshake advertises `protocolVersion: 2` and the `multi_session` capability. DBX falls back to the v1 one-process-per-pool lifecycle when that capability is absent.
|
||||
|
||||
## Session lifecycle
|
||||
|
||||
- `open_session` creates one logical database session. Parameters contain the normal connection fields plus `agentSessionId`.
|
||||
- Every connection-scoped RPC contains `agentSessionId`.
|
||||
- `validate_session` validates and, where supported, reconnects only that session.
|
||||
- `cancel_session` cancels active statements and cursor fetches for only that session; other sessions in the runtime continue normally.
|
||||
- `close_session` closes the session connection, query cursors, and table-read cursors without affecting other sessions.
|
||||
- `shutdown` closes all sessions and terminates the runtime.
|
||||
|
||||
`agentSessionId` identifies a logical database connection. Existing `sessionId` fields remain pagination cursor identifiers and must not be used as logical connection identifiers.
|
||||
|
||||
## Concurrency
|
||||
|
||||
Requests for different sessions may execute concurrently. Requests for the same session are serialized because connection state, transactions, schema changes, and driver connections are not generally safe for concurrent use. JSON-RPC responses may be returned out of order and are correlated by request `id`.
|
||||
|
||||
## Runtime compatibility
|
||||
|
||||
Runtime reuse keys include the Agent driver key, executable or JAR path, launch arguments, working directory, JRE selection, JVM options, classpath-affecting options, and native executable version boundary. Host, account, schema, and credentials belong to sessions and are not part of the runtime key.
|
||||
|
||||
Etcd and ZooKeeper retain the legacy path because they use the key-value Agent protocol rather than the SQL session contract. Older Agent binaries and JARs also remain on the legacy path.
|
||||
|
||||
## Resource limits and recovery
|
||||
|
||||
A runtime accepts at most 256 logical sessions. Closing the final session starts a 30-second grace period before the process exits, preventing rapid tab open/close cycles from repeatedly starting a runtime. Process EOF fails all pending requests; the failed runtime is removed from reuse and recreated on demand. Connection validation and reconnect operate on a single logical session.
|
||||
|
||||
## Driver author guidance
|
||||
|
||||
Use `MultiSessionJsonRpcServer(YourAgent::new)` for Java SQL Agents so each logical session receives a new `DatabaseAgent` and JDBC connection. Do not store connection, statement, cursor, transaction, or schema state in static mutable fields. Use the session execution context for paged query resources. Native Agents must provide equivalent per-session state and synchronized stdout writes.
|
||||
|
||||
The Xugu native Agent keeps one database connection per logical session plus one shared control connection per database endpoint. Because `go-xugu-driver` does not interrupt network reads through `context.Context`, cancellation records the server-side session ID and calls `DBMS_DBA.KILL_SESSION_TRANS` through the shared control connection.
|
||||
@@ -8,7 +8,7 @@ import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.QueryResult;
|
||||
import com.dbx.agent.TableInfo;
|
||||
@@ -217,7 +217,7 @@ public final class AccessAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -225,7 +225,7 @@ public final class AccessAgent extends BaseDatabaseAgent {
|
||||
options.getMaxRows(),
|
||||
options.getFetchSize(),
|
||||
options.getTimeoutSecs(),
|
||||
JdbcExecutor.INSTANCE::defaultResultValue
|
||||
JdbcExecutor.current()::defaultResultValue
|
||||
);
|
||||
}
|
||||
|
||||
@@ -367,6 +367,6 @@ public final class AccessAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new AccessAgent()).run();
|
||||
new MultiSessionJsonRpcServer(AccessAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.TableInfo;
|
||||
@@ -229,6 +229,6 @@ public final class BigQueryAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new BigQueryAgent()).run();
|
||||
new MultiSessionJsonRpcServer(BigQueryAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TriggerInfo;
|
||||
import java.sql.PreparedStatement;
|
||||
@@ -189,6 +189,6 @@ public final class CassandraAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new CassandraAgent()).run();
|
||||
new MultiSessionJsonRpcServer(CassandraAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
@@ -828,7 +828,7 @@ public final class DamengAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -842,7 +842,7 @@ public final class DamengAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return JdbcExecutor.current().executePage(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -854,7 +854,7 @@ public final class DamengAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -1322,6 +1322,6 @@ public final class DamengAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new DamengAgent()).run();
|
||||
new MultiSessionJsonRpcServer(DamengAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.databend;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
@@ -294,6 +294,6 @@ public final class DatabendAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new DatabendAgent()).run();
|
||||
new MultiSessionJsonRpcServer(DatabendAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -3,7 +3,7 @@ package com.dbx.agent.databricks;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class DatabricksAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile DATABRICKS_PROFILE = new DatabricksProfile();
|
||||
@@ -13,7 +13,7 @@ public final class DatabricksAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new DatabricksAgent()).run();
|
||||
new MultiSessionJsonRpcServer(DatabricksAgent::new).run();
|
||||
}
|
||||
|
||||
private static final class DatabricksProfile extends JdbcAgentProfile {
|
||||
|
||||
@@ -9,7 +9,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -387,7 +387,7 @@ public final class Db2Agent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -510,6 +510,6 @@ public final class Db2Agent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new Db2Agent()).run();
|
||||
new MultiSessionJsonRpcServer(Db2Agent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.exasol;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class ExasolAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile EXASOL_PROFILE = new JdbcAgentProfile(
|
||||
@@ -17,6 +17,6 @@ public final class ExasolAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new ExasolAgent()).run();
|
||||
new MultiSessionJsonRpcServer(ExasolAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.firebird;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class FirebirdAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile FIREBIRD_PROFILE = new JdbcAgentProfile(
|
||||
@@ -17,6 +17,6 @@ public final class FirebirdAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new FirebirdAgent()).run();
|
||||
new MultiSessionJsonRpcServer(FirebirdAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
@@ -340,6 +340,6 @@ public final class Gbase8aAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new Gbase8aAgent()).run();
|
||||
new MultiSessionJsonRpcServer(Gbase8aAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
@@ -376,7 +376,7 @@ public final class Gbase8sAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new Gbase8sAgent()).run();
|
||||
new MultiSessionJsonRpcServer(Gbase8sAgent::new).run();
|
||||
}
|
||||
|
||||
private static int port(ConnectParams params) {
|
||||
|
||||
@@ -9,7 +9,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -406,7 +406,7 @@ public final class GoldendbAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -550,6 +550,6 @@ public final class GoldendbAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new GoldendbAgent()).run();
|
||||
new MultiSessionJsonRpcServer(GoldendbAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
package com.dbx.agent.h2legacy;
|
||||
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.h2.H2Agent;
|
||||
|
||||
public final class H2LegacyAgent extends H2Agent {
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new H2LegacyAgent()).run();
|
||||
new MultiSessionJsonRpcServer(H2LegacyAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -414,7 +414,7 @@ public class H2Agent extends AbstractJdbcAgent {
|
||||
|
||||
@Override
|
||||
protected Object resultValue(ResultSet rs, int index, int sqlType) {
|
||||
return unchecked(() -> JdbcExecutor.INSTANCE.defaultResultValue(rs, index, sqlType));
|
||||
return unchecked(() -> JdbcExecutor.current().defaultResultValue(rs, index, sqlType));
|
||||
}
|
||||
|
||||
@Override
|
||||
@@ -511,6 +511,6 @@ public class H2Agent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new H2Agent()).run();
|
||||
new MultiSessionJsonRpcServer(H2Agent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
package com.dbx.agent.highgo;
|
||||
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.PostgresLikeAgent;
|
||||
import com.dbx.agent.PostgresLikeAgentProfile;
|
||||
|
||||
@@ -15,6 +15,6 @@ public final class HighgoAgent extends PostgresLikeAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new HighgoAgent()).run();
|
||||
new MultiSessionJsonRpcServer(HighgoAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TriggerInfo;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
@@ -269,6 +269,6 @@ public final class HiveAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new HiveAgent()).run();
|
||||
new MultiSessionJsonRpcServer(HiveAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -459,7 +459,7 @@ public final class InformixAgent extends BaseDatabaseAgent {
|
||||
case "ROLLBACK WORK" -> "ROLLBACK";
|
||||
default -> sql;
|
||||
};
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
normalizedSql,
|
||||
schema,
|
||||
@@ -551,6 +551,6 @@ public final class InformixAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new InformixAgent()).run();
|
||||
new MultiSessionJsonRpcServer(InformixAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ package com.dbx.agent.iotdb;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.Statement;
|
||||
import java.util.ArrayList;
|
||||
@@ -41,6 +41,6 @@ public final class IoTDBAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new IoTDBAgent()).run();
|
||||
new MultiSessionJsonRpcServer(IoTDBAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,7 @@ package com.dbx.agent.iris;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.StandardJdbcMetadata;
|
||||
|
||||
import java.sql.Connection;
|
||||
@@ -154,6 +154,6 @@ public final class IrisAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new IrisAgent()).run();
|
||||
new MultiSessionJsonRpcServer(IrisAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -834,6 +834,6 @@ public final class KingbaseAgent extends PostgresLikeAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new KingbaseAgent()).run();
|
||||
new MultiSessionJsonRpcServer(KingbaseAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TriggerInfo;
|
||||
import java.sql.ResultSet;
|
||||
@@ -139,6 +139,6 @@ public final class KylinAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new KylinAgent()).run();
|
||||
new MultiSessionJsonRpcServer(KylinAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -42,6 +42,9 @@ import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Set;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import javax.net.ssl.KeyManager;
|
||||
import javax.net.ssl.KeyManagerFactory;
|
||||
import javax.net.ssl.SSLContext;
|
||||
@@ -59,7 +62,10 @@ public final class MongoAgent {
|
||||
private static final JsonWriterSettings EXTENDED_JSON_SETTINGS = JsonWriterSettings.builder()
|
||||
.outputMode(JsonMode.RELAXED)
|
||||
.build();
|
||||
private static MongoClient client;
|
||||
private static final String LEGACY_SESSION_ID = "__legacy__";
|
||||
private static final int MAX_SESSIONS = 256;
|
||||
private static final ThreadLocal<MongoClient> CURRENT_CLIENT = new ThreadLocal<>();
|
||||
private static MongoClient legacyClient;
|
||||
|
||||
private MongoAgent() {
|
||||
}
|
||||
@@ -110,7 +116,7 @@ public final class MongoAgent {
|
||||
return builder;
|
||||
}
|
||||
|
||||
private static Object connect(JsonObject params) {
|
||||
private static MongoClient openClient(JsonObject params) {
|
||||
JsonObject connObj = params.has("connection") && params.get("connection").isJsonObject()
|
||||
? params.getAsJsonObject("connection")
|
||||
: params;
|
||||
@@ -118,11 +124,19 @@ public final class MongoAgent {
|
||||
|
||||
MongoClientSettings.Builder builder = configureBuilder(connObj);
|
||||
|
||||
if (client != null) {
|
||||
MongoClient client = MongoClients.create(builder.build());
|
||||
try {
|
||||
client.getDatabase(database).runCommand(new Document("ping", 1));
|
||||
} catch (RuntimeException error) {
|
||||
client.close();
|
||||
throw error;
|
||||
}
|
||||
client = MongoClients.create(builder.build());
|
||||
client.getDatabase(database).runCommand(new Document("ping", 1));
|
||||
return client;
|
||||
}
|
||||
|
||||
private static Object connect(JsonObject params) {
|
||||
closeLegacyClient();
|
||||
legacyClient = openClient(params);
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
||||
@@ -880,10 +894,7 @@ public final class MongoAgent {
|
||||
case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENT -> deleteDocument(params);
|
||||
case AgentProtocol.MONGO_METHOD_DELETE_DOCUMENTS -> deleteDocuments(params);
|
||||
case AgentProtocol.METHOD_DISCONNECT, AgentProtocol.METHOD_SHUTDOWN -> {
|
||||
if (client != null) {
|
||||
client.close();
|
||||
client = null;
|
||||
}
|
||||
closeLegacyClient();
|
||||
if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
|
||||
System.exit(0);
|
||||
}
|
||||
@@ -894,12 +905,23 @@ public final class MongoAgent {
|
||||
}
|
||||
|
||||
private static MongoClient requireClient() {
|
||||
MongoClient client = CURRENT_CLIENT.get();
|
||||
if (client == null) {
|
||||
client = legacyClient;
|
||||
}
|
||||
if (client == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
return client;
|
||||
}
|
||||
|
||||
private static void closeLegacyClient() {
|
||||
if (legacyClient != null) {
|
||||
legacyClient.close();
|
||||
legacyClient = null;
|
||||
}
|
||||
}
|
||||
|
||||
private static String coalesce(String value) {
|
||||
return value == null ? "" : value;
|
||||
}
|
||||
@@ -936,16 +958,149 @@ public final class MongoAgent {
|
||||
public static void main(String[] args) throws Exception {
|
||||
System.out.println("{\"ready\":true}");
|
||||
System.out.flush();
|
||||
new RuntimeServer().run();
|
||||
}
|
||||
|
||||
BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
|
||||
while (true) {
|
||||
String line = reader.readLine();
|
||||
if (line == null) {
|
||||
break;
|
||||
private static final class RuntimeServer {
|
||||
private final Map<String, Session> sessions = new ConcurrentHashMap<>();
|
||||
private final ExecutorService requests = Executors.newCachedThreadPool();
|
||||
private final Object outputLock = new Object();
|
||||
|
||||
private void run() throws Exception {
|
||||
BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));
|
||||
String line;
|
||||
while ((line = reader.readLine()) != null) {
|
||||
String request = line;
|
||||
requests.submit(() -> writeResponse(handleRuntimeRequest(request)));
|
||||
}
|
||||
closeAllSessions();
|
||||
requests.shutdownNow();
|
||||
}
|
||||
|
||||
System.out.println(handleRequest(line));
|
||||
System.out.flush();
|
||||
private String handleRuntimeRequest(String line) {
|
||||
JsonObject req = JsonParser.parseString(line).getAsJsonObject();
|
||||
JsonElement id = req.get("id");
|
||||
String method = req.get("method").getAsString();
|
||||
JsonObject params = req.has("params") && req.get("params").isJsonObject()
|
||||
? req.getAsJsonObject("params")
|
||||
: new JsonObject();
|
||||
JsonObject response = new JsonObject();
|
||||
response.addProperty("jsonrpc", "2.0");
|
||||
response.add("id", id);
|
||||
try {
|
||||
Object result;
|
||||
if (AgentProtocol.METHOD_HANDSHAKE.equals(method)) {
|
||||
result = AgentProtocol.multiSessionHandshakeResult();
|
||||
} else if (AgentProtocol.METHOD_OPEN_SESSION.equals(method)) {
|
||||
result = openSession(requiredSessionId(params), params);
|
||||
} else if (AgentProtocol.METHOD_CLOSE_SESSION.equals(method)) {
|
||||
result = closeSession(requiredSessionId(params));
|
||||
} else if (AgentProtocol.METHOD_VALIDATE_SESSION.equals(method)) {
|
||||
result = session(requiredSessionId(params)).validate(params);
|
||||
} else if (AgentProtocol.METHOD_CANCEL_SESSION.equals(method)) {
|
||||
// The legacy synchronous MongoDB driver has no safe per-operation cancel API.
|
||||
result = Collections.singletonMap("ok", true);
|
||||
} else if (AgentProtocol.METHOD_CONNECT.equals(method)) {
|
||||
closeSession(LEGACY_SESSION_ID);
|
||||
result = openSession(LEGACY_SESSION_ID, params);
|
||||
} else if (AgentProtocol.METHOD_DISCONNECT.equals(method)) {
|
||||
result = closeSession(LEGACY_SESSION_ID);
|
||||
} else if (AgentProtocol.METHOD_SHUTDOWN.equals(method)) {
|
||||
closeAllSessions();
|
||||
result = Collections.singletonMap("ok", true);
|
||||
} else {
|
||||
String sessionId = params.has("agentSessionId")
|
||||
? params.get("agentSessionId").getAsString()
|
||||
: LEGACY_SESSION_ID;
|
||||
result = session(sessionId).handle(method, params);
|
||||
}
|
||||
response.add("result", GSON.toJsonTree(result));
|
||||
} catch (Exception error) {
|
||||
JsonObject rpcError = new JsonObject();
|
||||
rpcError.addProperty("code", -1);
|
||||
rpcError.addProperty("message", error.getMessage() == null ? "Unknown error" : error.getMessage());
|
||||
response.add("error", rpcError);
|
||||
}
|
||||
return GSON.toJson(response);
|
||||
}
|
||||
|
||||
private Object openSession(String sessionId, JsonObject params) {
|
||||
if (sessions.size() >= MAX_SESSIONS && !sessions.containsKey(sessionId)) {
|
||||
throw new IllegalStateException("Agent session limit reached: " + MAX_SESSIONS);
|
||||
}
|
||||
Session created = new Session(openClient(params));
|
||||
Session existing = sessions.putIfAbsent(sessionId, created);
|
||||
if (existing != null) {
|
||||
created.close();
|
||||
throw new IllegalStateException("Agent session already exists: " + sessionId);
|
||||
}
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
||||
private Object closeSession(String sessionId) {
|
||||
Session removed = sessions.remove(sessionId);
|
||||
if (removed != null) {
|
||||
removed.close();
|
||||
}
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
||||
private Session session(String sessionId) {
|
||||
Session session = sessions.get(sessionId);
|
||||
if (session == null) {
|
||||
throw new IllegalStateException("Agent session not found: " + sessionId);
|
||||
}
|
||||
return session;
|
||||
}
|
||||
|
||||
private void closeAllSessions() {
|
||||
for (String sessionId : sessions.keySet()) {
|
||||
closeSession(sessionId);
|
||||
}
|
||||
}
|
||||
|
||||
private static String requiredSessionId(JsonObject params) {
|
||||
if (!params.has("agentSessionId") || params.get("agentSessionId").getAsString().trim().isEmpty()) {
|
||||
throw new IllegalArgumentException("agentSessionId is required");
|
||||
}
|
||||
return params.get("agentSessionId").getAsString();
|
||||
}
|
||||
|
||||
private void writeResponse(String response) {
|
||||
synchronized (outputLock) {
|
||||
System.out.println(response);
|
||||
System.out.flush();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private static final class Session {
|
||||
private final MongoClient client;
|
||||
|
||||
private Session(MongoClient client) {
|
||||
this.client = client;
|
||||
}
|
||||
|
||||
private synchronized Object handle(String method, JsonObject params) {
|
||||
CURRENT_CLIENT.set(client);
|
||||
try {
|
||||
return dispatch(method, params);
|
||||
} finally {
|
||||
CURRENT_CLIENT.remove();
|
||||
}
|
||||
}
|
||||
|
||||
private Object validate(JsonObject params) {
|
||||
JsonObject connection = params.has("connection") && params.get("connection").isJsonObject()
|
||||
? params.getAsJsonObject("connection")
|
||||
: params;
|
||||
String database = defaultString(stringOrNull(connection, "database"), "admin");
|
||||
client.getDatabase(database).runCommand(new Document("ping", 1));
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
||||
private synchronized void close() {
|
||||
client.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -98,6 +98,16 @@ class MongoAgentTest {
|
||||
assertTrue(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_METADATA));
|
||||
}
|
||||
|
||||
@Test
|
||||
void legacyJsonRpcHandshakeRemainsProtocolV1() {
|
||||
JsonObject result = JsonParser.parseString(MongoAgent.handleRequest(
|
||||
"{\"jsonrpc\":\"2.0\",\"id\":71,\"method\":\"handshake\",\"params\":{}}"
|
||||
)).getAsJsonObject().getAsJsonObject("result");
|
||||
|
||||
assertEquals(1, result.get("protocolVersion").getAsInt());
|
||||
assertFalse(containsCapability(result.getAsJsonArray("capabilities"), AgentProtocol.CAPABILITY_MULTI_SESSION));
|
||||
}
|
||||
|
||||
@Test
|
||||
void listIndexesMethodIsRecognizedOverJsonRpc() {
|
||||
String response = MongoAgent.handleRequest(
|
||||
|
||||
@@ -8,7 +8,7 @@ import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.QueryResult;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TransactionExecutor;
|
||||
@@ -279,7 +279,7 @@ public class Neo4jAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -345,6 +345,6 @@ public class Neo4jAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new Neo4jAgent()).run();
|
||||
new MultiSessionJsonRpcServer(Neo4jAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -9,7 +9,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.TableInfo;
|
||||
@@ -618,6 +618,6 @@ public final class OceanBaseOracleAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new OceanBaseOracleAgent()).run();
|
||||
new MultiSessionJsonRpcServer(OceanBaseOracleAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
_ "github.com/sijms/go-ora/v2"
|
||||
@@ -25,8 +26,11 @@ import (
|
||||
)
|
||||
|
||||
const protocolVersion = 1
|
||||
const multiSessionProtocolVersion = 2
|
||||
const defaultMaxRows = 1000
|
||||
const oracleCharsetZHS32GB18030 = 854
|
||||
const legacyAgentSessionID = "__legacy__"
|
||||
const maxAgentSessions = 256
|
||||
|
||||
var (
|
||||
oraclePlSQLBlockStartRegexp = regexp.MustCompile(`(?is)^\s*(?:DECLARE|BEGIN|CREATE\s+(?:OR\s+REPLACE\s+)?(?:(?:EDITIONABLE|NONEDITIONABLE)\s+)?(?:FUNCTION|PROCEDURE|TRIGGER|PACKAGE(?:\s+BODY)?|TYPE(?:\s+BODY)?))\b`)
|
||||
@@ -317,11 +321,28 @@ type server struct {
|
||||
tableReadSessions map[string]*querySession
|
||||
nextSessionID int64
|
||||
nextTableReadSessionID int64
|
||||
activeCancelMu sync.Mutex
|
||||
activeCancel context.CancelFunc
|
||||
activeRows map[*sql.Rows]context.CancelFunc
|
||||
}
|
||||
|
||||
type agentSession struct {
|
||||
server *server
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
type runtimeServer struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*agentSession
|
||||
}
|
||||
|
||||
func main() {
|
||||
s := newServer()
|
||||
runtime := newRuntimeServer()
|
||||
encoder := json.NewEncoder(os.Stdout)
|
||||
var encoderMu sync.Mutex
|
||||
var requests sync.WaitGroup
|
||||
shutdown := make(chan struct{})
|
||||
var shutdownOnce sync.Once
|
||||
fmt.Fprintln(os.Stdout, `{"ready":true}`)
|
||||
|
||||
scanner := bufio.NewScanner(os.Stdin)
|
||||
@@ -331,22 +352,230 @@ func main() {
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
resp, shutdown := s.handleLine(line)
|
||||
if err := encoder.Encode(resp); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
var requestEnvelope request
|
||||
if json.Unmarshal([]byte(line), &requestEnvelope) == nil && requestEnvelope.Method == "shutdown" {
|
||||
requests.Wait()
|
||||
resp, _ := runtime.handleLine(line)
|
||||
encoderMu.Lock()
|
||||
err := encoder.Encode(resp)
|
||||
encoderMu.Unlock()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if shutdown {
|
||||
requests.Add(1)
|
||||
go func() {
|
||||
defer requests.Done()
|
||||
resp, shouldShutdown := runtime.handleLine(line)
|
||||
encoderMu.Lock()
|
||||
err := encoder.Encode(resp)
|
||||
encoderMu.Unlock()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
shutdownOnce.Do(func() { close(shutdown) })
|
||||
return
|
||||
}
|
||||
if shouldShutdown {
|
||||
shutdownOnce.Do(func() { close(shutdown) })
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case <-shutdown:
|
||||
requests.Wait()
|
||||
return
|
||||
default:
|
||||
}
|
||||
}
|
||||
requests.Wait()
|
||||
if err := scanner.Err(); err != nil && !errors.Is(err, io.EOF) {
|
||||
fmt.Fprintf(os.Stderr, "failed to read stdin: %v\n", err)
|
||||
}
|
||||
}
|
||||
|
||||
func newRuntimeServer() *runtimeServer {
|
||||
return &runtimeServer{sessions: map[string]*agentSession{}}
|
||||
}
|
||||
|
||||
func (r *runtimeServer) handleLine(line string) (response, bool) {
|
||||
var req request
|
||||
if err := json.Unmarshal([]byte(line), &req); err != nil {
|
||||
return errorResponse(nil, err), false
|
||||
}
|
||||
if len(req.ID) == 0 {
|
||||
req.ID = json.RawMessage("1")
|
||||
}
|
||||
result, shutdown, err := r.dispatch(req.Method, req.Params)
|
||||
if err != nil {
|
||||
return errorResponse(req.ID, err), false
|
||||
}
|
||||
return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown
|
||||
}
|
||||
|
||||
func (r *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
|
||||
switch method {
|
||||
case "handshake":
|
||||
return map[string]any{
|
||||
"protocolVersion": multiSessionProtocolVersion,
|
||||
"agentProtocolVersion": multiSessionProtocolVersion,
|
||||
"capabilities": []string{"connect", "test_connection", "metadata", "query", "ddl", "multi_session"},
|
||||
}, false, nil
|
||||
case "open_session":
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
if agentSessionID == "" {
|
||||
return nil, false, errors.New("agentSessionId is required")
|
||||
}
|
||||
var connectParams connectParams
|
||||
if err := decodeParams(params, &connectParams); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, r.openSession(agentSessionID, connectParams)
|
||||
case "close_session":
|
||||
return map[string]bool{"ok": true}, false, r.closeSession(stringParam(params, "agentSessionId"))
|
||||
case "validate_session":
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
session, err := r.session(agentSessionID)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
if _, _, err := session.server.dispatch("validate_connection", params); err == nil {
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
}
|
||||
// Reconnect only this logical session. Other sessions in the runtime keep
|
||||
// their connections, transactions, cursors, and in-flight requests.
|
||||
if err := session.server.connect(session.server.params); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
case "cancel_session":
|
||||
session, err := r.session(stringParam(params, "agentSessionId"))
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
session.server.cancelActiveQuery()
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
case "test_connection":
|
||||
return newServer().dispatch(method, params)
|
||||
case "shutdown":
|
||||
return map[string]bool{"ok": true}, true, r.closeAllSessions()
|
||||
case "connect":
|
||||
var connectParams connectParams
|
||||
if err := decodeParams(params, &connectParams); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, r.replaceSession(legacyAgentSessionID, connectParams)
|
||||
case "disconnect":
|
||||
return map[string]bool{"ok": true}, false, r.closeSession(legacyAgentSessionID)
|
||||
default:
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
if agentSessionID == "" {
|
||||
agentSessionID = legacyAgentSessionID
|
||||
}
|
||||
return r.withSession(agentSessionID, method, params)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *runtimeServer) withSession(agentSessionID, method string, params map[string]json.RawMessage) (any, bool, error) {
|
||||
session, err := r.session(agentSessionID)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
// Oracle connection state, transactions, and cursors are session-scoped;
|
||||
// serialize one session while allowing separate sessions to run in parallel.
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
return session.server.dispatch(method, params)
|
||||
}
|
||||
|
||||
func (r *runtimeServer) openSession(agentSessionID string, params connectParams) error {
|
||||
r.mu.Lock()
|
||||
if _, exists := r.sessions[agentSessionID]; exists {
|
||||
r.mu.Unlock()
|
||||
return fmt.Errorf("agent session already exists: %s", agentSessionID)
|
||||
}
|
||||
if len(r.sessions) >= maxAgentSessions {
|
||||
r.mu.Unlock()
|
||||
return fmt.Errorf("agent session limit reached: %d", maxAgentSessions)
|
||||
}
|
||||
session := &agentSession{server: newServer()}
|
||||
r.sessions[agentSessionID] = session
|
||||
r.mu.Unlock()
|
||||
|
||||
// Reserve the id under the registry lock, then connect outside it so unrelated
|
||||
// sessions can establish database connections concurrently.
|
||||
session.mu.Lock()
|
||||
err := session.server.connect(params)
|
||||
session.mu.Unlock()
|
||||
if err != nil {
|
||||
r.mu.Lock()
|
||||
if r.sessions[agentSessionID] == session {
|
||||
delete(r.sessions, agentSessionID)
|
||||
}
|
||||
r.mu.Unlock()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *runtimeServer) replaceSession(agentSessionID string, params connectParams) error {
|
||||
_ = r.closeSession(agentSessionID)
|
||||
return r.openSession(agentSessionID, params)
|
||||
}
|
||||
|
||||
func (r *runtimeServer) session(agentSessionID string) (*agentSession, error) {
|
||||
if agentSessionID == "" {
|
||||
return nil, errors.New("agentSessionId is required")
|
||||
}
|
||||
r.mu.RLock()
|
||||
session := r.sessions[agentSessionID]
|
||||
r.mu.RUnlock()
|
||||
if session == nil {
|
||||
return nil, fmt.Errorf("agent session not found: %s", agentSessionID)
|
||||
}
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func (r *runtimeServer) closeSession(agentSessionID string) error {
|
||||
if agentSessionID == "" {
|
||||
return errors.New("agentSessionId is required")
|
||||
}
|
||||
r.mu.Lock()
|
||||
session := r.sessions[agentSessionID]
|
||||
delete(r.sessions, agentSessionID)
|
||||
r.mu.Unlock()
|
||||
if session == nil {
|
||||
return nil
|
||||
}
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
return session.server.disconnect()
|
||||
}
|
||||
|
||||
func (r *runtimeServer) closeAllSessions() error {
|
||||
r.mu.Lock()
|
||||
sessions := r.sessions
|
||||
r.sessions = map[string]*agentSession{}
|
||||
r.mu.Unlock()
|
||||
var firstErr error
|
||||
for _, session := range sessions {
|
||||
session.mu.Lock()
|
||||
err := session.server.disconnect()
|
||||
session.mu.Unlock()
|
||||
if firstErr == nil && err != nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func newServer() *server {
|
||||
return &server{sessions: map[string]*querySession{}, tableReadSessions: map[string]*querySession{}}
|
||||
return &server{
|
||||
sessions: map[string]*querySession{},
|
||||
tableReadSessions: map[string]*querySession{},
|
||||
activeRows: map[*sql.Rows]context.CancelFunc{},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *server) handleLine(line string) (response, bool) {
|
||||
@@ -389,6 +618,11 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any
|
||||
}
|
||||
defer db.Close()
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
case "validate_connection":
|
||||
if s.db == nil {
|
||||
return nil, false, errors.New("not connected")
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, pingDB(s.db, 5*time.Second)
|
||||
case "list_databases":
|
||||
result, err := s.listDatabases()
|
||||
return result, false, err
|
||||
@@ -750,7 +984,7 @@ func (s *server) listDatabases() ([]databaseInfo, error) {
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []databaseInfo
|
||||
for rows.Next() {
|
||||
var name string
|
||||
@@ -810,7 +1044,7 @@ func (s *server) listDatabasesFiltered(visibleSchemas []string) ([]databaseInfo,
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []databaseInfo
|
||||
for rows.Next() {
|
||||
var name string
|
||||
@@ -1054,7 +1288,7 @@ func (s *server) listTables(schema string, constraints metadataListConstraints)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []tableInfo
|
||||
for rows.Next() {
|
||||
var item tableInfo
|
||||
@@ -1082,7 +1316,7 @@ func (s *server) listObjects(schema string, constraints metadataListConstraints)
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []objectInfo
|
||||
for rows.Next() {
|
||||
var item objectInfo
|
||||
@@ -1124,7 +1358,7 @@ ORDER BY c.COLUMN_ID`, []any{schema, table})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []columnInfo
|
||||
for rows.Next() {
|
||||
var item columnInfo
|
||||
@@ -1165,7 +1399,7 @@ ORDER BY COLUMN_ID`, []any{schema, table})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []oracleColumnMeta
|
||||
for rows.Next() {
|
||||
var item oracleColumnMeta
|
||||
@@ -1201,7 +1435,7 @@ ORDER BY i.INDEX_NAME, ic.COLUMN_POSITION`, []any{schema, table})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
|
||||
byName := map[string]*indexInfo{}
|
||||
order := []string{}
|
||||
@@ -1259,7 +1493,7 @@ ORDER BY ac.CONSTRAINT_NAME, acc.POSITION`, []any{schema, table})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []foreignKeyInfo
|
||||
for rows.Next() {
|
||||
var item foreignKeyInfo
|
||||
@@ -1285,7 +1519,7 @@ ORDER BY TRIGGER_NAME`, []any{schema, table})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []triggerInfo
|
||||
for rows.Next() {
|
||||
var item triggerInfo
|
||||
@@ -1329,7 +1563,7 @@ ORDER BY LINE`, []any{schema, strings.ToUpper(name), upperType})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var builder strings.Builder
|
||||
for rows.Next() {
|
||||
var line string
|
||||
@@ -1748,7 +1982,7 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes
|
||||
}
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
columnTypes := columnTypeNames(rows)
|
||||
@@ -1760,14 +1994,14 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes
|
||||
result, err := readQuerySessionPage(session, pageSize)
|
||||
result.ExecutionTimeMS = time.Since(start).Milliseconds()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
if result.HasMore {
|
||||
sessionID := s.storeQuerySession(session)
|
||||
result.SessionID = &sessionID
|
||||
} else {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -1814,7 +2048,7 @@ func (s *server) startTableRead(opts queryOptions, pageSize int) (queryPageResul
|
||||
}
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
columnTypes := columnTypeNames(rows)
|
||||
@@ -1826,14 +2060,14 @@ func (s *server) startTableRead(opts queryOptions, pageSize int) (queryPageResul
|
||||
result, err := readQuerySessionPage(session, pageSize)
|
||||
result.ExecutionTimeMS = time.Since(start).Milliseconds()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
if result.HasMore {
|
||||
sessionID := s.storeTableReadSession(session)
|
||||
result.SessionID = &sessionID
|
||||
} else {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -1868,7 +2102,7 @@ func (s *server) closeQuerySession(sessionID string) bool {
|
||||
if session == nil {
|
||||
return false
|
||||
}
|
||||
session.rows.Close()
|
||||
s.closeRows(session.rows)
|
||||
delete(s.sessions, sessionID)
|
||||
return true
|
||||
}
|
||||
@@ -1878,7 +2112,7 @@ func (s *server) closeTableReadSession(sessionID string) bool {
|
||||
if session == nil {
|
||||
return false
|
||||
}
|
||||
session.rows.Close()
|
||||
s.closeRows(session.rows)
|
||||
delete(s.tableReadSessions, sessionID)
|
||||
return true
|
||||
}
|
||||
@@ -1964,7 +2198,7 @@ func (s *server) executeSelect(sqlText string, maxRows int) (queryResult, error)
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
@@ -2029,7 +2263,7 @@ func (s *server) queryRowsWithXMLTypeRewriteIfNeeded(sqlText string) (*sql.Rows,
|
||||
}
|
||||
rewritten, err := s.rewriteXMLTypeSelectSQL(sqlText)
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return nil, err
|
||||
}
|
||||
if rewritten == sqlText {
|
||||
@@ -2037,7 +2271,7 @@ func (s *server) queryRowsWithXMLTypeRewriteIfNeeded(sqlText string) (*sql.Rows,
|
||||
}
|
||||
// Only pay the ALL_TAB_COLUMNS rewrite cost when the result metadata shows
|
||||
// XMLTYPE. Ordinary Oracle queries should not run dictionary probes first.
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return s.queryRows(rewritten, nil)
|
||||
}
|
||||
|
||||
@@ -2671,10 +2905,51 @@ func (s *server) queryRows(sqlText string, args []any) (*sql.Rows, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(args) == 0 {
|
||||
return db.Query(sqlText)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s.activeCancelMu.Lock()
|
||||
s.activeCancel = cancel
|
||||
s.activeCancelMu.Unlock()
|
||||
rows, queryErr := db.QueryContext(ctx, sqlText, args...)
|
||||
s.activeCancelMu.Lock()
|
||||
s.activeCancel = nil
|
||||
if queryErr != nil {
|
||||
cancel()
|
||||
} else {
|
||||
// Keep the context alive for paged reads; database/sql may continue
|
||||
// fetching from the driver until Rows is closed or exhausted.
|
||||
s.activeRows[rows] = cancel
|
||||
}
|
||||
return db.Query(sqlText, args...)
|
||||
s.activeCancelMu.Unlock()
|
||||
return rows, queryErr
|
||||
}
|
||||
|
||||
func (s *server) cancelActiveQuery() {
|
||||
s.activeCancelMu.Lock()
|
||||
cancels := make([]context.CancelFunc, 0, len(s.activeRows)+1)
|
||||
if s.activeCancel != nil {
|
||||
cancels = append(cancels, s.activeCancel)
|
||||
}
|
||||
for _, cancel := range s.activeRows {
|
||||
cancels = append(cancels, cancel)
|
||||
}
|
||||
s.activeCancelMu.Unlock()
|
||||
for _, cancel := range cancels {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *server) closeRows(rows *sql.Rows) error {
|
||||
if rows == nil {
|
||||
return nil
|
||||
}
|
||||
s.activeCancelMu.Lock()
|
||||
cancel := s.activeRows[rows]
|
||||
delete(s.activeRows, rows)
|
||||
s.activeCancelMu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
return rows.Close()
|
||||
}
|
||||
|
||||
func decodeParams(params map[string]json.RawMessage, target any) error {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -49,6 +50,87 @@ func TestHandshakeResponse(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeHandshakeAdvertisesMultiSessionProtocol(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":7,"method":"handshake","params":{"appVersion":"dev"}}`)
|
||||
if shutdown || resp.Error != nil {
|
||||
t.Fatalf("unexpected handshake response: shutdown=%v error=%v", shutdown, resp.Error)
|
||||
}
|
||||
data, err := json.Marshal(resp.Result)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var result struct {
|
||||
ProtocolVersion int `json:"protocolVersion"`
|
||||
AgentProtocolVersion int `json:"agentProtocolVersion"`
|
||||
Capabilities []string `json:"capabilities"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.ProtocolVersion != 2 || result.AgentProtocolVersion != 2 {
|
||||
t.Fatalf("unexpected protocol versions: %+v", result)
|
||||
}
|
||||
if !contains(result.Capabilities, "multi_session") {
|
||||
t.Fatalf("expected multi_session capability, got %v", result.Capabilities)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeMissingAgentSessionDoesNotUseQueryCursorSessionID(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":8,"method":"fetch_query_page","params":{"sessionId":"cursor-1","pageSize":10}}`)
|
||||
if shutdown {
|
||||
t.Fatal("fetch_query_page should not shut down the runtime")
|
||||
}
|
||||
if resp.Error == nil || !strings.Contains(resp.Error.Message, legacyAgentSessionID) {
|
||||
t.Fatalf("expected missing legacy agent session error, got %#v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeCloseOneSessionKeepsOtherSessionRegistered(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
runtime.sessions["a"] = &agentSession{server: newServer()}
|
||||
runtime.sessions["b"] = &agentSession{server: newServer()}
|
||||
|
||||
if err := runtime.closeSession("a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := runtime.session("a"); err == nil {
|
||||
t.Fatal("closed session should be removed")
|
||||
}
|
||||
if _, err := runtime.session("b"); err != nil {
|
||||
t.Fatalf("other session should remain registered: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeCancelSessionOnlyCancelsTargetSession(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
serverA := newServer()
|
||||
serverB := newServer()
|
||||
ctxA, cancelA := context.WithCancel(context.Background())
|
||||
ctxB, cancelB := context.WithCancel(context.Background())
|
||||
serverA.activeCancel = cancelA
|
||||
serverB.activeCancel = cancelB
|
||||
runtime.sessions["a"] = &agentSession{server: serverA}
|
||||
runtime.sessions["b"] = &agentSession{server: serverB}
|
||||
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":9,"method":"cancel_session","params":{"agentSessionId":"a"}}`)
|
||||
if shutdown || resp.Error != nil {
|
||||
t.Fatalf("unexpected cancel response: shutdown=%v error=%v", shutdown, resp.Error)
|
||||
}
|
||||
select {
|
||||
case <-ctxA.Done():
|
||||
default:
|
||||
t.Fatal("target session was not canceled")
|
||||
}
|
||||
select {
|
||||
case <-ctxB.Done():
|
||||
t.Fatal("canceling session a should not cancel session b")
|
||||
default:
|
||||
}
|
||||
cancelB()
|
||||
}
|
||||
|
||||
func TestCloseMissingQuerySessionReturnsFalse(t *testing.T) {
|
||||
s := newServer()
|
||||
resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":8,"method":"close_query_session","params":{"sessionId":"missing"}}`)
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.oscar;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class OscarAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile OSCAR_PROFILE = new JdbcAgentProfile(
|
||||
@@ -16,6 +16,6 @@ public final class OscarAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new OscarAgent()).run();
|
||||
new MultiSessionJsonRpcServer(OscarAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.saphana;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
|
||||
@@ -21,6 +21,6 @@ public final class SapHanaAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new SapHanaAgent()).run();
|
||||
new MultiSessionJsonRpcServer(SapHanaAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -376,6 +376,6 @@ public final class SnowflakeAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new SnowflakeAgent()).run();
|
||||
new MultiSessionJsonRpcServer(SnowflakeAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TriggerInfo;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
@@ -336,6 +336,6 @@ public final class SparkAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new SparkAgent()).run();
|
||||
new MultiSessionJsonRpcServer(SparkAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
+2
-2
@@ -3,7 +3,7 @@ package com.dbx.agent.sqlserverlegacy;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
import java.security.Security;
|
||||
import java.sql.Connection;
|
||||
@@ -259,6 +259,6 @@ public final class SqlServerLegacyAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) throws Exception {
|
||||
new JsonRpcServer(new SqlServerLegacyAgent()).run();
|
||||
new MultiSessionJsonRpcServer(SqlServerLegacyAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -407,7 +407,7 @@ public final class SundbAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -525,6 +525,6 @@ public final class SundbAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new SundbAgent()).run();
|
||||
new MultiSessionJsonRpcServer(SundbAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
@@ -191,7 +191,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
return JdbcExecutor.INSTANCE.execute(
|
||||
return JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -205,7 +205,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.executePage(
|
||||
return JdbcExecutor.current().executePage(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -217,7 +217,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.INSTANCE.startTableRead(
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
requireConnected(),
|
||||
sql,
|
||||
schema,
|
||||
@@ -497,7 +497,7 @@ public final class TDengineAgent extends BaseDatabaseAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new TDengineAgent()).run();
|
||||
new MultiSessionJsonRpcServer(TDengineAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.teradata;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class TeradataAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile TERADATA_PROFILE = new JdbcAgentProfile(
|
||||
@@ -16,6 +16,6 @@ public final class TeradataAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new TeradataAgent()).run();
|
||||
new MultiSessionJsonRpcServer(TeradataAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,7 +7,7 @@ import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcIdentifiers;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.MetadataSqlSupport;
|
||||
import com.dbx.agent.TableInfo;
|
||||
@@ -340,6 +340,6 @@ public final class TrinoAgent extends AbstractJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new TrinoAgent()).run();
|
||||
new MultiSessionJsonRpcServer(TrinoAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
package com.dbx.agent.vastbase;
|
||||
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
import com.dbx.agent.PostgresLikeAgent;
|
||||
import com.dbx.agent.PostgresLikeAgentProfile;
|
||||
@@ -16,7 +16,7 @@ public final class VastbaseAgent extends PostgresLikeAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new VastbaseAgent()).run();
|
||||
new MultiSessionJsonRpcServer(VastbaseAgent::new).run();
|
||||
}
|
||||
|
||||
@Override
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.vertica;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class VerticaAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile VERTICA_PROFILE = new JdbcAgentProfile(
|
||||
@@ -16,6 +16,6 @@ public final class VerticaAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new VerticaAgent()).run();
|
||||
new MultiSessionJsonRpcServer(VerticaAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
+513
-44
@@ -3,6 +3,7 @@ package main
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -13,14 +14,18 @@ import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
_ "gitee.com/XuguDB/go-xugu-driver"
|
||||
)
|
||||
|
||||
const protocolVersion = 1
|
||||
const multiSessionProtocolVersion = 2
|
||||
const defaultMaxRows = 1000
|
||||
const defaultXuguPort = 5138
|
||||
const legacyAgentSessionID = "__legacy__"
|
||||
const maxAgentSessions = 256
|
||||
const xuguListDatabasesSQL = `
|
||||
SELECT DB_NAME
|
||||
FROM ALL_DATABASES
|
||||
@@ -259,15 +264,43 @@ type triggerInfo struct {
|
||||
}
|
||||
|
||||
type server struct {
|
||||
db *sql.DB
|
||||
params connectParams
|
||||
sessions map[string]*querySession
|
||||
nextSessionID int64
|
||||
db *sql.DB
|
||||
cancelDB *sql.DB
|
||||
ownsCancelDB bool
|
||||
params connectParams
|
||||
nodeID int
|
||||
databaseSessionID int64
|
||||
sessions map[string]*querySession
|
||||
nextSessionID int64
|
||||
activeCancelMu sync.Mutex
|
||||
activeCancel context.CancelFunc
|
||||
activeRows map[*sql.Rows]context.CancelFunc
|
||||
}
|
||||
|
||||
type agentSession struct {
|
||||
server *server
|
||||
controlKey string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
type sharedControl struct {
|
||||
db *sql.DB
|
||||
refs int
|
||||
}
|
||||
|
||||
type runtimeServer struct {
|
||||
mu sync.RWMutex
|
||||
sessions map[string]*agentSession
|
||||
connectMu sync.Mutex
|
||||
controlMu sync.Mutex
|
||||
controls map[string]*sharedControl
|
||||
}
|
||||
|
||||
func main() {
|
||||
s := newServer()
|
||||
runtime := newRuntimeServer()
|
||||
encoder := json.NewEncoder(os.Stdout)
|
||||
var encoderMu sync.Mutex
|
||||
var requests sync.WaitGroup
|
||||
fmt.Fprintln(os.Stdout, `{"ready":true}`)
|
||||
|
||||
scanner := bufio.NewScanner(os.Stdin)
|
||||
@@ -277,22 +310,260 @@ func main() {
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
resp, shutdown := s.handleLine(line)
|
||||
if err := encoder.Encode(resp); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
return
|
||||
}
|
||||
if shutdown {
|
||||
var requestEnvelope request
|
||||
if json.Unmarshal([]byte(line), &requestEnvelope) == nil && requestEnvelope.Method == "shutdown" {
|
||||
requests.Wait()
|
||||
resp, _ := runtime.handleLine(line)
|
||||
encoderMu.Lock()
|
||||
err := encoder.Encode(resp)
|
||||
encoderMu.Unlock()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
requests.Add(1)
|
||||
go func(line string) {
|
||||
defer requests.Done()
|
||||
resp, _ := runtime.handleLine(line)
|
||||
encoderMu.Lock()
|
||||
err := encoder.Encode(resp)
|
||||
encoderMu.Unlock()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "failed to write response: %v\n", err)
|
||||
}
|
||||
}(line)
|
||||
}
|
||||
requests.Wait()
|
||||
if err := scanner.Err(); err != nil && !errors.Is(err, io.EOF) {
|
||||
fmt.Fprintf(os.Stderr, "failed to read stdin: %v\n", err)
|
||||
}
|
||||
}
|
||||
|
||||
func newServer() *server {
|
||||
return &server{sessions: map[string]*querySession{}}
|
||||
return &server{sessions: map[string]*querySession{}, activeRows: map[*sql.Rows]context.CancelFunc{}}
|
||||
}
|
||||
|
||||
func newRuntimeServer() *runtimeServer {
|
||||
return &runtimeServer{sessions: map[string]*agentSession{}, controls: map[string]*sharedControl{}}
|
||||
}
|
||||
|
||||
func (r *runtimeServer) handleLine(line string) (response, bool) {
|
||||
var req request
|
||||
if err := json.Unmarshal([]byte(line), &req); err != nil {
|
||||
return errorResponse(nil, err), false
|
||||
}
|
||||
if len(req.ID) == 0 {
|
||||
req.ID = json.RawMessage("1")
|
||||
}
|
||||
result, shutdown, err := r.dispatch(req.Method, req.Params)
|
||||
if err != nil {
|
||||
return errorResponse(req.ID, err), false
|
||||
}
|
||||
return response{JSONRPC: "2.0", ID: req.ID, Result: result}, shutdown
|
||||
}
|
||||
|
||||
func (r *runtimeServer) dispatch(method string, params map[string]json.RawMessage) (any, bool, error) {
|
||||
switch method {
|
||||
case "handshake":
|
||||
return map[string]any{
|
||||
"protocolVersion": multiSessionProtocolVersion,
|
||||
"agentProtocolVersion": multiSessionProtocolVersion,
|
||||
"capabilities": []string{"connect", "test_connection", "metadata", "query", "ddl", "multi_session"},
|
||||
}, false, nil
|
||||
case "open_session":
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
if agentSessionID == "" {
|
||||
return nil, false, errors.New("agentSessionId is required")
|
||||
}
|
||||
var cp connectParams
|
||||
if err := decodeParams(params, &cp); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, r.openSession(agentSessionID, cp)
|
||||
case "close_session":
|
||||
return map[string]bool{"ok": true}, false, r.closeSession(stringParam(params, "agentSessionId"))
|
||||
case "validate_session":
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
session, err := r.session(agentSessionID)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
if err := session.server.validateConnection(); err == nil {
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
}
|
||||
r.connectMu.Lock()
|
||||
err = session.server.connectWithControl(session.server.params, session.server.cancelDB, false)
|
||||
r.connectMu.Unlock()
|
||||
return map[string]bool{"ok": true}, false, err
|
||||
case "cancel_session":
|
||||
session, err := r.session(stringParam(params, "agentSessionId"))
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
session.server.cancelActiveQuery()
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
case "test_connection":
|
||||
return newServer().dispatch(method, params)
|
||||
case "shutdown":
|
||||
return map[string]bool{"ok": true}, true, r.closeAllSessions()
|
||||
case "connect":
|
||||
var cp connectParams
|
||||
if err := decodeParams(params, &cp); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, r.replaceSession(legacyAgentSessionID, cp)
|
||||
case "disconnect":
|
||||
return map[string]bool{"ok": true}, false, r.closeSession(legacyAgentSessionID)
|
||||
default:
|
||||
agentSessionID := stringParam(params, "agentSessionId")
|
||||
if agentSessionID == "" {
|
||||
agentSessionID = legacyAgentSessionID
|
||||
}
|
||||
return r.withSession(agentSessionID, method, params)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *runtimeServer) withSession(agentSessionID, method string, params map[string]json.RawMessage) (any, bool, error) {
|
||||
session, err := r.session(agentSessionID)
|
||||
if err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
// Database, schema, transaction, and cursor state are connection-scoped.
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
return session.server.dispatch(method, params)
|
||||
}
|
||||
|
||||
func (r *runtimeServer) openSession(agentSessionID string, params connectParams) error {
|
||||
r.mu.Lock()
|
||||
if _, exists := r.sessions[agentSessionID]; exists {
|
||||
r.mu.Unlock()
|
||||
return fmt.Errorf("agent session already exists: %s", agentSessionID)
|
||||
}
|
||||
if len(r.sessions) >= maxAgentSessions {
|
||||
r.mu.Unlock()
|
||||
return fmt.Errorf("agent session limit reached: %d", maxAgentSessions)
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
server := newServer()
|
||||
params.URLParams = appendURLParam(params.URLParams, "APP_NAME", xuguSessionAppName(agentSessionID))
|
||||
r.connectMu.Lock()
|
||||
controlKey, controlDB, err := r.acquireControl(params)
|
||||
if err == nil {
|
||||
err = server.connectWithControl(params, controlDB, false)
|
||||
}
|
||||
r.connectMu.Unlock()
|
||||
if err != nil {
|
||||
if controlKey != "" {
|
||||
r.releaseControl(controlKey)
|
||||
}
|
||||
return err
|
||||
}
|
||||
session := &agentSession{server: server, controlKey: controlKey}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if _, exists := r.sessions[agentSessionID]; exists {
|
||||
_ = server.disconnect()
|
||||
r.releaseControl(controlKey)
|
||||
return fmt.Errorf("agent session already exists: %s", agentSessionID)
|
||||
}
|
||||
if len(r.sessions) >= maxAgentSessions {
|
||||
_ = server.disconnect()
|
||||
r.releaseControl(controlKey)
|
||||
return fmt.Errorf("agent session limit reached: %d", maxAgentSessions)
|
||||
}
|
||||
r.sessions[agentSessionID] = session
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *runtimeServer) replaceSession(agentSessionID string, params connectParams) error {
|
||||
_ = r.closeSession(agentSessionID)
|
||||
return r.openSession(agentSessionID, params)
|
||||
}
|
||||
|
||||
func (r *runtimeServer) session(agentSessionID string) (*agentSession, error) {
|
||||
r.mu.RLock()
|
||||
session := r.sessions[agentSessionID]
|
||||
r.mu.RUnlock()
|
||||
if session == nil {
|
||||
return nil, fmt.Errorf("agent session not found: %s", agentSessionID)
|
||||
}
|
||||
return session, nil
|
||||
}
|
||||
|
||||
func (r *runtimeServer) closeSession(agentSessionID string) error {
|
||||
r.mu.Lock()
|
||||
session := r.sessions[agentSessionID]
|
||||
delete(r.sessions, agentSessionID)
|
||||
r.mu.Unlock()
|
||||
if session == nil {
|
||||
return nil
|
||||
}
|
||||
session.mu.Lock()
|
||||
defer session.mu.Unlock()
|
||||
err := session.server.disconnect()
|
||||
r.releaseControl(session.controlKey)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *runtimeServer) acquireControl(params connectParams) (string, *sql.DB, error) {
|
||||
r.controlMu.Lock()
|
||||
defer r.controlMu.Unlock()
|
||||
cancelParams := xuguControlParams(params)
|
||||
key := buildDSN(cancelParams)
|
||||
if control := r.controls[key]; control != nil {
|
||||
control.refs++
|
||||
return key, control.db, nil
|
||||
}
|
||||
db, err := openDB(cancelParams)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
db.Close()
|
||||
return "", nil, err
|
||||
}
|
||||
r.controls[key] = &sharedControl{db: db, refs: 1}
|
||||
return key, db, nil
|
||||
}
|
||||
|
||||
func (r *runtimeServer) releaseControl(key string) {
|
||||
if key == "" {
|
||||
return
|
||||
}
|
||||
r.controlMu.Lock()
|
||||
defer r.controlMu.Unlock()
|
||||
control := r.controls[key]
|
||||
if control == nil {
|
||||
return
|
||||
}
|
||||
control.refs--
|
||||
if control.refs <= 0 {
|
||||
_ = control.db.Close()
|
||||
delete(r.controls, key)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *runtimeServer) closeAllSessions() error {
|
||||
r.mu.Lock()
|
||||
ids := make([]string, 0, len(r.sessions))
|
||||
for id := range r.sessions {
|
||||
ids = append(ids, id)
|
||||
}
|
||||
r.mu.Unlock()
|
||||
var firstErr error
|
||||
for _, id := range ids {
|
||||
if err := r.closeSession(id); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
func (s *server) handleLine(line string) (response, bool) {
|
||||
@@ -340,6 +611,8 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any
|
||||
return nil, false, err
|
||||
}
|
||||
return map[string]bool{"ok": true}, false, nil
|
||||
case "validate_connection":
|
||||
return map[string]bool{"ok": true}, false, s.validateConnection()
|
||||
case "list_databases":
|
||||
result, err := s.listDatabases()
|
||||
return result, false, err
|
||||
@@ -409,6 +682,18 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any
|
||||
return result, false, err
|
||||
case "close_query_session":
|
||||
return s.closeQuerySession(stringParam(params, "sessionId")), false, nil
|
||||
case "start_table_read":
|
||||
var opts queryOptions
|
||||
if err := decodeParams(params, &opts); err != nil {
|
||||
return nil, false, err
|
||||
}
|
||||
result, err := s.executeQueryPage(opts, intParam(params, "pageSize"))
|
||||
return result, false, err
|
||||
case "fetch_table_read_page":
|
||||
result, err := s.fetchQueryPage(stringParam(params, "sessionId"), intParam(params, "pageSize"))
|
||||
return result, false, err
|
||||
case "close_table_read_session":
|
||||
return s.closeQuerySession(stringParam(params, "sessionId")), false, nil
|
||||
case "list_indexes":
|
||||
if err := s.useDatabase(stringParam(params, "database")); err != nil {
|
||||
return nil, false, err
|
||||
@@ -451,44 +736,162 @@ func (s *server) dispatch(method string, params map[string]json.RawMessage) (any
|
||||
}
|
||||
|
||||
func (s *server) connect(params connectParams) error {
|
||||
_ = s.disconnect()
|
||||
db, err := openDB(params)
|
||||
cancelParams := xuguControlParams(params)
|
||||
cancelDB, err := openDB(cancelParams)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
if err := cancelDB.PingContext(ctx); err != nil {
|
||||
cancelDB.Close()
|
||||
return err
|
||||
}
|
||||
if err := s.connectWithControl(params, cancelDB, true); err != nil {
|
||||
cancelDB.Close()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *server) connectWithControl(params connectParams, cancelDB *sql.DB, ownsCancelDB bool) error {
|
||||
_ = s.disconnect()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
|
||||
defer cancel()
|
||||
before, err := xuguDatabaseSessions(cancelDB)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
db, err := openDB(params)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := db.PingContext(ctx); err != nil {
|
||||
db.Close()
|
||||
return err
|
||||
}
|
||||
after, err := xuguDatabaseSessions(cancelDB)
|
||||
if err != nil {
|
||||
db.Close()
|
||||
return err
|
||||
}
|
||||
databaseSession, err := newXuguDatabaseSession(before, after)
|
||||
if err != nil {
|
||||
db.Close()
|
||||
return err
|
||||
}
|
||||
s.db = db
|
||||
s.cancelDB = cancelDB
|
||||
s.ownsCancelDB = ownsCancelDB
|
||||
s.params = params
|
||||
s.nodeID = databaseSession.nodeID
|
||||
s.databaseSessionID = databaseSession.sessionID
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *server) disconnect() error {
|
||||
s.cancelActiveQuery()
|
||||
s.closeAllQuerySessions()
|
||||
if s.ownsCancelDB && s.cancelDB != nil {
|
||||
_ = s.cancelDB.Close()
|
||||
}
|
||||
if s.db == nil {
|
||||
return nil
|
||||
}
|
||||
err := s.db.Close()
|
||||
s.db = nil
|
||||
s.cancelDB = nil
|
||||
s.ownsCancelDB = false
|
||||
s.nodeID = 0
|
||||
s.databaseSessionID = 0
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *server) validateConnection() error {
|
||||
if s.db == nil {
|
||||
return errors.New("not connected")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
return s.db.PingContext(ctx)
|
||||
}
|
||||
|
||||
func openDB(params connectParams) (*sql.DB, error) {
|
||||
dsn := buildDSN(params)
|
||||
db, err := sql.Open("xugu", dsn)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
db.SetMaxOpenConns(4)
|
||||
// One logical Agent session must map to exactly one server-side session so
|
||||
// schema, transaction, cursor, and cancellation state stay deterministic.
|
||||
db.SetMaxOpenConns(1)
|
||||
db.SetMaxIdleConns(1)
|
||||
db.SetConnMaxLifetime(30 * time.Minute)
|
||||
return db, nil
|
||||
}
|
||||
|
||||
func appendURLParam(raw, key, value string) string {
|
||||
parts := make([]string, 0, 2)
|
||||
for _, part := range strings.Split(raw, ";") {
|
||||
part = strings.TrimSpace(part)
|
||||
if part != "" && !strings.HasPrefix(strings.ToUpper(part), strings.ToUpper(key)+"=") {
|
||||
parts = append(parts, part)
|
||||
}
|
||||
}
|
||||
parts = append(parts, key+"="+value)
|
||||
return strings.Join(parts, ";")
|
||||
}
|
||||
|
||||
func xuguSessionAppName(agentSessionID string) string {
|
||||
digest := sha256.Sum256([]byte(agentSessionID))
|
||||
return fmt.Sprintf("DBX_%x", digest[:8])
|
||||
}
|
||||
|
||||
func xuguControlParams(params connectParams) connectParams {
|
||||
params.Database = "SYSTEM"
|
||||
params.ConnectionString = ""
|
||||
params.URLParams = appendURLParam(params.URLParams, "APP_NAME", "DBX_CONTROL")
|
||||
return params
|
||||
}
|
||||
|
||||
type xuguDatabaseSession struct {
|
||||
nodeID int
|
||||
sessionID int64
|
||||
}
|
||||
|
||||
func xuguDatabaseSessions(db *sql.DB) (map[xuguDatabaseSession]struct{}, error) {
|
||||
rows, err := db.Query("SELECT NODEID, SESSION_ID FROM SYS_SESSIONS")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
result := map[xuguDatabaseSession]struct{}{}
|
||||
for rows.Next() {
|
||||
var session xuguDatabaseSession
|
||||
if err := rows.Scan(&session.nodeID, &session.sessionID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
result[session] = struct{}{}
|
||||
}
|
||||
return result, rows.Err()
|
||||
}
|
||||
|
||||
func newXuguDatabaseSession(
|
||||
before map[xuguDatabaseSession]struct{},
|
||||
after map[xuguDatabaseSession]struct{},
|
||||
) (xuguDatabaseSession, error) {
|
||||
var candidates []xuguDatabaseSession
|
||||
for session := range after {
|
||||
if _, existed := before[session]; !existed {
|
||||
candidates = append(candidates, session)
|
||||
}
|
||||
}
|
||||
if len(candidates) != 1 {
|
||||
return xuguDatabaseSession{}, fmt.Errorf("failed to identify Xugu server session: found %d new sessions", len(candidates))
|
||||
}
|
||||
return candidates[0], nil
|
||||
}
|
||||
|
||||
func buildDSN(params connectParams) string {
|
||||
connectionString := strings.TrimSpace(params.ConnectionString)
|
||||
if looksLikeXuguDSN(connectionString) {
|
||||
@@ -684,7 +1087,7 @@ func (s *server) listDatabases() ([]databaseInfo, error) {
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []databaseInfo
|
||||
for rows.Next() {
|
||||
var name string
|
||||
@@ -759,7 +1162,7 @@ func (s *server) listSchemas() ([]string, error) {
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []string
|
||||
for rows.Next() {
|
||||
var schema string
|
||||
@@ -784,7 +1187,7 @@ ORDER BY CASE WHEN UPPER(s.SCHEMA_NAME) = UPPER(?) THEN 0 ELSE 1 END, s.SCHEMA_N
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
if rows.Next() {
|
||||
var schema string
|
||||
if err := rows.Scan(&schema); err != nil {
|
||||
@@ -816,7 +1219,7 @@ func (s *server) listTables(schema string, constraints metadataListConstraints)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []tableInfo
|
||||
for rows.Next() {
|
||||
var item tableInfo
|
||||
@@ -838,7 +1241,7 @@ func (s *server) listObjects(schema string, constraints metadataListConstraints)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []objectInfo
|
||||
for rows.Next() {
|
||||
var item objectInfo
|
||||
@@ -1021,7 +1424,7 @@ func (s *server) getColumns(schema, table string) ([]columnInfo, error) {
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []columnInfo
|
||||
for rows.Next() {
|
||||
var item columnInfo
|
||||
@@ -1056,7 +1459,7 @@ func (s *server) columnsFromSelect(schema, table string, primaryKeys map[string]
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
types, err := rows.ColumnTypes()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -1090,7 +1493,7 @@ func (s *server) primaryKeyColumns(schema, table string) (map[string]bool, error
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
result := map[string]bool{}
|
||||
for rows.Next() {
|
||||
var define string
|
||||
@@ -1117,7 +1520,7 @@ func (s *server) listIndexes(schema, table string) ([]indexInfo, error) {
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
|
||||
var result []indexInfo
|
||||
for rows.Next() {
|
||||
@@ -1160,7 +1563,7 @@ ORDER BY c.CONS_NAME`, []any{schema, table})
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []foreignKeyInfo
|
||||
for rows.Next() {
|
||||
var name, define, refTable string
|
||||
@@ -1199,7 +1602,7 @@ ORDER BY tr.TRIG_NAME`, []any{schema, table})
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var result []triggerInfo
|
||||
for rows.Next() {
|
||||
var item triggerInfo
|
||||
@@ -1228,7 +1631,7 @@ func (s *server) getObjectSource(schema, name, objectType string) (map[string]an
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var builder strings.Builder
|
||||
for rows.Next() {
|
||||
var line string
|
||||
@@ -1249,7 +1652,7 @@ func (s *server) getTableDDL(schema, table string) (string, error) {
|
||||
var ddl string
|
||||
rows, err := s.queryRows("SELECT TO_CHAR(DBMS_METADATA.GET_DDL('TABLE', ?, ?)) FROM DUAL", []any{strings.ToUpper(table), schema})
|
||||
if err == nil {
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
if rows.Next() {
|
||||
if scanErr := rows.Scan(&ddl); scanErr == nil && strings.TrimSpace(ddl) != "" {
|
||||
if err := rows.Err(); err != nil {
|
||||
@@ -1274,7 +1677,7 @@ func (s *server) getExplainInfo(sqlText string) (string, error) {
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
return "", err
|
||||
@@ -1303,12 +1706,14 @@ func (s *server) executeTransaction(params map[string]json.RawMessage) (queryRes
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
tx, err := db.Begin()
|
||||
ctx, cancel := s.beginActiveOperation()
|
||||
defer s.endActiveOperation(cancel)
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
if strings.TrimSpace(payload.Schema) != "" {
|
||||
if _, err := tx.Exec("SET SCHEMA " + quoteIdentifier(payload.Schema)); err != nil {
|
||||
if _, err := tx.ExecContext(ctx, "SET SCHEMA "+quoteIdentifier(payload.Schema)); err != nil {
|
||||
tx.Rollback()
|
||||
return queryResult{}, err
|
||||
}
|
||||
@@ -1320,7 +1725,7 @@ func (s *server) executeTransaction(params map[string]json.RawMessage) (queryRes
|
||||
if statement == "" {
|
||||
continue
|
||||
}
|
||||
result, err := tx.Exec(statement)
|
||||
result, err := tx.ExecContext(ctx, statement)
|
||||
if err != nil {
|
||||
tx.Rollback()
|
||||
return queryResult{}, err
|
||||
@@ -1370,7 +1775,7 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes
|
||||
}
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
columnTypes := columnTypeNames(rows)
|
||||
@@ -1382,14 +1787,14 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes
|
||||
result, err := readQuerySessionPage(session, pageSize)
|
||||
result.ExecutionTimeMS = time.Since(start).Milliseconds()
|
||||
if err != nil {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
return queryPageResult{}, err
|
||||
}
|
||||
if result.HasMore {
|
||||
sessionID := s.storeQuerySession(session)
|
||||
result.SessionID = &sessionID
|
||||
} else {
|
||||
rows.Close()
|
||||
s.closeRows(rows)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -1424,7 +1829,7 @@ func (s *server) closeQuerySession(sessionID string) bool {
|
||||
if session == nil {
|
||||
return false
|
||||
}
|
||||
session.rows.Close()
|
||||
s.closeRows(session.rows)
|
||||
delete(s.sessions, sessionID)
|
||||
return true
|
||||
}
|
||||
@@ -1497,7 +1902,9 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) {
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
execResult, err := db.Exec(sqlText)
|
||||
ctx, cancel := s.beginActiveOperation()
|
||||
defer s.endActiveOperation(cancel)
|
||||
execResult, err := db.ExecContext(ctx, sqlText)
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
@@ -1510,7 +1917,7 @@ func (s *server) executeSelect(sqlText string, maxRows int) (queryResult, error)
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
columns, err := rows.Columns()
|
||||
if err != nil {
|
||||
return queryResult{}, err
|
||||
@@ -1562,7 +1969,9 @@ func (s *server) setSchema(schema string) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = db.Exec("SET SCHEMA " + quoteIdentifier(schema))
|
||||
ctx, cancel := s.beginActiveOperation()
|
||||
defer s.endActiveOperation(cancel)
|
||||
_, err = db.ExecContext(ctx, "SET SCHEMA "+quoteIdentifier(schema))
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -1571,10 +1980,70 @@ func (s *server) queryRows(sqlText string, args []any) (*sql.Rows, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(args) == 0 {
|
||||
return db.Query(sqlText)
|
||||
ctx, cancel := s.beginActiveOperation()
|
||||
rows, queryErr := db.QueryContext(ctx, sqlText, args...)
|
||||
s.activeCancelMu.Lock()
|
||||
s.activeCancel = nil
|
||||
if queryErr != nil {
|
||||
cancel()
|
||||
} else {
|
||||
// Paged queries may continue fetching after QueryContext returns.
|
||||
s.activeRows[rows] = cancel
|
||||
}
|
||||
return db.Query(sqlText, args...)
|
||||
s.activeCancelMu.Unlock()
|
||||
return rows, queryErr
|
||||
}
|
||||
|
||||
func (s *server) beginActiveOperation() (context.Context, context.CancelFunc) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
s.activeCancelMu.Lock()
|
||||
s.activeCancel = cancel
|
||||
s.activeCancelMu.Unlock()
|
||||
return ctx, cancel
|
||||
}
|
||||
|
||||
func (s *server) endActiveOperation(cancel context.CancelFunc) {
|
||||
cancel()
|
||||
s.activeCancelMu.Lock()
|
||||
s.activeCancel = nil
|
||||
s.activeCancelMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *server) cancelActiveQuery() {
|
||||
s.activeCancelMu.Lock()
|
||||
cancels := make([]context.CancelFunc, 0, len(s.activeRows)+1)
|
||||
if s.activeCancel != nil {
|
||||
cancels = append(cancels, s.activeCancel)
|
||||
}
|
||||
for _, cancel := range s.activeRows {
|
||||
cancels = append(cancels, cancel)
|
||||
}
|
||||
s.activeCancelMu.Unlock()
|
||||
for _, cancel := range cancels {
|
||||
cancel()
|
||||
}
|
||||
if len(cancels) > 0 && s.cancelDB != nil && s.databaseSessionID > 0 {
|
||||
// go-xugu-driver does not implement QueryerContext/ExecerContext and
|
||||
// blocks in network reads, so context cancellation alone cannot interrupt
|
||||
// an in-flight statement. Xugu's control procedure stops the target
|
||||
// session's current transaction while preserving the connection. Runtime
|
||||
// sessions share one control connection per database endpoint.
|
||||
_, _ = s.cancelDB.Exec(fmt.Sprintf("CALL DBMS_DBA.KILL_SESSION_TRANS(%d, %d)", s.nodeID, s.databaseSessionID))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *server) closeRows(rows *sql.Rows) error {
|
||||
if rows == nil {
|
||||
return nil
|
||||
}
|
||||
s.activeCancelMu.Lock()
|
||||
cancel := s.activeRows[rows]
|
||||
delete(s.activeRows, rows)
|
||||
s.activeCancelMu.Unlock()
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
return rows.Close()
|
||||
}
|
||||
|
||||
func objectSourceQuery(schema, name, objectType string) (string, []any, error) {
|
||||
@@ -1718,7 +2187,7 @@ func (s *server) tableIndexDDL(schema, table string) (string, error) {
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var ddl string
|
||||
if rows.Next() {
|
||||
if err := rows.Scan(&ddl); err != nil {
|
||||
@@ -1741,7 +2210,7 @@ WHERE UPPER(s.SCHEMA_NAME) = UPPER(?)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer rows.Close()
|
||||
defer s.closeRows(rows)
|
||||
var comment *string
|
||||
if rows.Next() {
|
||||
if err := rows.Scan(&comment); err != nil {
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -46,6 +48,119 @@ func TestHandshakeResponse(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeHandshakeAdvertisesMultiSessionProtocol(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":7,"method":"handshake","params":{}}`)
|
||||
if shutdown || resp.Error != nil {
|
||||
t.Fatalf("unexpected handshake response: shutdown=%v error=%v", shutdown, resp.Error)
|
||||
}
|
||||
data, err := json.Marshal(resp.Result)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var result struct {
|
||||
ProtocolVersion int `json:"protocolVersion"`
|
||||
Capabilities []string `json:"capabilities"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &result); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result.ProtocolVersion != multiSessionProtocolVersion || !contains(result.Capabilities, "multi_session") {
|
||||
t.Fatalf("unexpected runtime handshake: %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeMissingAgentSessionDoesNotUseQueryCursorSessionID(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":8,"method":"fetch_query_page","params":{"sessionId":"cursor-1"}}`)
|
||||
if shutdown {
|
||||
t.Fatal("fetch_query_page should not shut down the runtime")
|
||||
}
|
||||
if resp.Error == nil || !strings.Contains(resp.Error.Message, legacyAgentSessionID) {
|
||||
t.Fatalf("expected missing legacy agent session error, got %#v", resp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeCloseOneSessionKeepsOtherSessionRegistered(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
runtime.sessions["a"] = &agentSession{server: newServer()}
|
||||
runtime.sessions["b"] = &agentSession{server: newServer()}
|
||||
|
||||
if err := runtime.closeSession("a"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := runtime.session("a"); err == nil {
|
||||
t.Fatal("closed session should be removed")
|
||||
}
|
||||
if _, err := runtime.session("b"); err != nil {
|
||||
t.Fatalf("other session should remain registered: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRuntimeCancelSessionOnlyCancelsTargetSession(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
serverA := newServer()
|
||||
serverB := newServer()
|
||||
ctxA, cancelA := context.WithCancel(context.Background())
|
||||
ctxB, cancelB := context.WithCancel(context.Background())
|
||||
serverA.activeCancel = cancelA
|
||||
serverB.activeCancel = cancelB
|
||||
runtime.sessions["a"] = &agentSession{server: serverA}
|
||||
runtime.sessions["b"] = &agentSession{server: serverB}
|
||||
|
||||
resp, shutdown := runtime.handleLine(`{"jsonrpc":"2.0","id":9,"method":"cancel_session","params":{"agentSessionId":"a"}}`)
|
||||
if shutdown || resp.Error != nil {
|
||||
t.Fatalf("unexpected cancel response: shutdown=%v error=%v", shutdown, resp.Error)
|
||||
}
|
||||
select {
|
||||
case <-ctxA.Done():
|
||||
default:
|
||||
t.Fatal("target session was not canceled")
|
||||
}
|
||||
select {
|
||||
case <-ctxB.Done():
|
||||
t.Fatal("canceling session a should not cancel session b")
|
||||
default:
|
||||
}
|
||||
cancelB()
|
||||
}
|
||||
|
||||
func TestRuntimeRejectsSessionsBeyondLimit(t *testing.T) {
|
||||
runtime := newRuntimeServer()
|
||||
for index := 0; index < maxAgentSessions; index++ {
|
||||
runtime.sessions[fmt.Sprintf("session-%d", index)] = &agentSession{server: newServer()}
|
||||
}
|
||||
err := runtime.openSession("overflow", connectParams{})
|
||||
if err == nil || !strings.Contains(err.Error(), "session limit") {
|
||||
t.Fatalf("expected session limit error, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewXuguDatabaseSessionFindsOnlyNewSession(t *testing.T) {
|
||||
existing := xuguDatabaseSession{nodeID: 1, sessionID: 10}
|
||||
created := xuguDatabaseSession{nodeID: 1, sessionID: 11}
|
||||
result, err := newXuguDatabaseSession(
|
||||
map[xuguDatabaseSession]struct{}{existing: {}},
|
||||
map[xuguDatabaseSession]struct{}{existing: {}, created: {}},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if result != created {
|
||||
t.Fatalf("unexpected session: %+v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestXuguSessionAppNameIsStableAndDoesNotExposeSessionID(t *testing.T) {
|
||||
name := xuguSessionAppName("tab-session-secret")
|
||||
if name != xuguSessionAppName("tab-session-secret") {
|
||||
t.Fatal("app name should be stable")
|
||||
}
|
||||
if strings.Contains(name, "tab-session-secret") || !strings.HasPrefix(name, "DBX_") {
|
||||
t.Fatalf("unexpected app name: %s", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCloseMissingQuerySessionReturnsFalse(t *testing.T) {
|
||||
s := newServer()
|
||||
resp, shutdown := s.handleLine(`{"jsonrpc":"2.0","id":8,"method":"close_query_session","params":{"sessionId":"missing"}}`)
|
||||
|
||||
@@ -2,7 +2,7 @@ package com.dbx.agent.yashandb;
|
||||
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
|
||||
public final class YashandbAgent extends ConfiguredJdbcAgent {
|
||||
public static final JdbcAgentProfile YASHANDB_PROFILE = new JdbcAgentProfile(
|
||||
@@ -17,6 +17,6 @@ public final class YashandbAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new YashandbAgent()).run();
|
||||
new MultiSessionJsonRpcServer(YashandbAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
dependencies {
|
||||
implementation fileTree(dir: 'libs', include: ['*.jar'])
|
||||
implementation 'org.opengauss:opengauss-jdbc:7.0.0-RC3-og'
|
||||
}
|
||||
|
||||
tasks.named('shadowJar') {
|
||||
manifest {
|
||||
attributes('Agent-Label': 'openGauss', 'Main-Class': 'com.dbx.agent.opengauss.OpenGaussAgent')
|
||||
}
|
||||
}
|
||||
@@ -1,20 +0,0 @@
|
||||
package com.dbx.agent.opengauss;
|
||||
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.PostgresLikeAgent;
|
||||
import com.dbx.agent.PostgresLikeAgentProfile;
|
||||
|
||||
public final class OpenGaussAgent extends PostgresLikeAgent {
|
||||
public static final PostgresLikeAgentProfile OPENGAUSS_PROFILE = new PostgresLikeAgentProfile(
|
||||
"org.opengauss.Driver",
|
||||
"jdbc:opengauss://{host}:{port}/{database}"
|
||||
);
|
||||
|
||||
public OpenGaussAgent() {
|
||||
super(OPENGAUSS_PROFILE);
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new JsonRpcServer(new OpenGaussAgent()).run();
|
||||
}
|
||||
}
|
||||
@@ -1,21 +0,0 @@
|
||||
package com.dbx.agent.opengauss;
|
||||
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
|
||||
class OpenGaussAgentTest {
|
||||
@Test
|
||||
void buildsOpenGaussUrl() {
|
||||
ConnectParams params = new ConnectParams();
|
||||
params.setHost("db.example.com");
|
||||
params.setPort(15432);
|
||||
params.setDatabase("postgres");
|
||||
|
||||
assertEquals(
|
||||
"jdbc:opengauss://db.example.com:15432/postgres",
|
||||
OpenGaussAgent.OPENGAUSS_PROFILE.buildUrl(params)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -44,6 +44,7 @@ const driverStoreTab = ref("agent");
|
||||
|
||||
import { useSettingsStore } from "@/stores/settingsStore";
|
||||
import type { DriverStorePathInfo } from "@/lib/backend/api";
|
||||
import { driverRuntimeProtocolLabel } from "./driverRuntimeDisplay";
|
||||
const settingsStore = useSettingsStore();
|
||||
|
||||
type DriverStoreDirKind = "plugin" | "agent";
|
||||
@@ -1479,6 +1480,10 @@ watch(driverStoreTab, (tab) => {
|
||||
<span>{{ runtimeKindLabel(runtime) }}</span>
|
||||
<span class="text-muted-foreground/50">/</span>
|
||||
<span>{{ runtimeSourceLabel(runtime) }}</span>
|
||||
<template v-if="runtime.protocol_mode">
|
||||
<span class="text-muted-foreground/50">/</span>
|
||||
<span>{{ driverRuntimeProtocolLabel(runtime) }}</span>
|
||||
</template>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { driverRuntimeProtocolLabel } from "./driverRuntimeDisplay";
|
||||
|
||||
describe("driverRuntimeProtocolLabel", () => {
|
||||
it("shows aggregated multi-session counts", () => {
|
||||
expect(driverRuntimeProtocolLabel({ protocol_mode: "multi_session", active_sessions: 4 })).toBe("Multi-session · 4 sessions");
|
||||
});
|
||||
|
||||
it("labels legacy runtimes and hides unknown modes", () => {
|
||||
expect(driverRuntimeProtocolLabel({ protocol_mode: "legacy", active_sessions: null })).toBe("Legacy");
|
||||
expect(driverRuntimeProtocolLabel({ protocol_mode: null, active_sessions: null })).toBeNull();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,11 @@
|
||||
import type { DriverRuntimeInfo } from "@/lib/backend/tauri";
|
||||
|
||||
export function driverRuntimeProtocolLabel(runtime: Pick<DriverRuntimeInfo, "protocol_mode" | "active_sessions">): string | null {
|
||||
if (runtime.protocol_mode === "multi_session") {
|
||||
return `Multi-session · ${runtime.active_sessions ?? 0} sessions`;
|
||||
}
|
||||
if (runtime.protocol_mode === "legacy") {
|
||||
return "Legacy";
|
||||
}
|
||||
return null;
|
||||
}
|
||||
@@ -134,6 +134,8 @@ export interface DriverRuntimeInfo {
|
||||
can_stop: boolean;
|
||||
can_restart: boolean;
|
||||
control_unavailable_reason: string | null;
|
||||
protocol_mode: "multi_session" | "legacy" | null;
|
||||
active_sessions: number | null;
|
||||
}
|
||||
|
||||
export interface DriverRuntimeSummary {
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
{
|
||||
"protocolVersion": 2,
|
||||
"handshakeMethod": "handshake",
|
||||
"handshakeResponseFields": [
|
||||
"protocolVersion",
|
||||
"agentProtocolVersion",
|
||||
"capabilities"
|
||||
],
|
||||
"allCapabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"kv",
|
||||
"multi_session"
|
||||
],
|
||||
"capabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"multi_session"
|
||||
],
|
||||
"defaultSqlCapabilities": [
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"multi_session"
|
||||
],
|
||||
"commonMethods": [
|
||||
"handshake",
|
||||
"connect",
|
||||
"open_session",
|
||||
"close_session",
|
||||
"validate_session",
|
||||
"cancel_session",
|
||||
"test_connection",
|
||||
"validate_connection",
|
||||
"list_databases",
|
||||
"list_schemas",
|
||||
"list_tables",
|
||||
"list_objects",
|
||||
"list_data_types",
|
||||
"completion_assistant_search_v1",
|
||||
"get_object_source",
|
||||
"get_table_ddl",
|
||||
"get_columns",
|
||||
"list_indexes",
|
||||
"list_foreign_keys",
|
||||
"list_triggers",
|
||||
"execute_query",
|
||||
"execute_query_page",
|
||||
"fetch_query_page",
|
||||
"close_query_session",
|
||||
"start_table_read",
|
||||
"fetch_table_read_page",
|
||||
"close_table_read_session",
|
||||
"get_explain_info",
|
||||
"execute_batch",
|
||||
"execute_transaction",
|
||||
"disconnect",
|
||||
"shutdown"
|
||||
],
|
||||
"mongoLegacyMethods": [
|
||||
"list_databases",
|
||||
"list_collections",
|
||||
"find_documents",
|
||||
"find_documents_extended_json",
|
||||
"count_documents",
|
||||
"server_version",
|
||||
"create_index",
|
||||
"drop_indexes",
|
||||
"drop_collection",
|
||||
"insert_document",
|
||||
"update_document",
|
||||
"update_documents",
|
||||
"delete_document",
|
||||
"delete_documents"
|
||||
],
|
||||
"kvMethods": [
|
||||
"kv_list_prefix",
|
||||
"kv_get",
|
||||
"kv_put",
|
||||
"kv_delete"
|
||||
],
|
||||
"sessionField": "agentSessionId",
|
||||
"cursorSessionField": "sessionId"
|
||||
}
|
||||
@@ -6,7 +6,7 @@ use std::path::{Path, PathBuf};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec, AgentMethod};
|
||||
use crate::db::agent_driver::{AgentDriverClient, AgentLaunchSpec, AgentMethod, AgentRuntimeClient};
|
||||
use crate::models::connection::DatabaseType;
|
||||
|
||||
pub const DEFAULT_JRE_KEY: &str = "21";
|
||||
@@ -394,6 +394,9 @@ pub struct AgentManager {
|
||||
base_dir: PathBuf,
|
||||
app_version: String,
|
||||
pub(crate) daemons: Mutex<std::collections::HashMap<String, AgentDriverClient>>,
|
||||
pub(crate) connection_runtimes: Mutex<
|
||||
std::collections::HashMap<String, std::sync::Arc<tokio::sync::OnceCell<std::sync::Arc<AgentRuntimeClient>>>>,
|
||||
>,
|
||||
}
|
||||
|
||||
impl Default for AgentManager {
|
||||
@@ -414,8 +417,12 @@ impl AgentManager {
|
||||
}
|
||||
|
||||
pub fn new_with_base_dir_and_app_version(base_dir: PathBuf, app_version: impl Into<String>) -> Self {
|
||||
let mgr =
|
||||
Self { base_dir, app_version: app_version.into(), daemons: Mutex::new(std::collections::HashMap::new()) };
|
||||
let mgr = Self {
|
||||
base_dir,
|
||||
app_version: app_version.into(),
|
||||
daemons: Mutex::new(std::collections::HashMap::new()),
|
||||
connection_runtimes: Mutex::new(std::collections::HashMap::new()),
|
||||
};
|
||||
mgr.migrate_legacy_jre();
|
||||
mgr.cleanup_pending_jre_dirs();
|
||||
mgr.cleanup_orphan_jre_dirs();
|
||||
@@ -760,7 +767,13 @@ impl AgentManager {
|
||||
}
|
||||
|
||||
pub async fn active_daemon_keys(&self) -> Vec<String> {
|
||||
self.daemons.lock().await.keys().cloned().collect()
|
||||
let mut keys = self.daemons.lock().await.keys().cloned().collect::<std::collections::HashSet<_>>();
|
||||
for runtime_key in self.connection_runtimes.lock().await.keys() {
|
||||
if let Some((agent_key, _)) = runtime_key.split_once('|') {
|
||||
keys.insert(agent_key.to_string());
|
||||
}
|
||||
}
|
||||
keys.into_iter().collect()
|
||||
}
|
||||
|
||||
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
|
||||
@@ -788,6 +801,27 @@ impl AgentManager {
|
||||
crate::agent_runtime::spawn_connection_client(self, db_type, driver_profile, extra_java_args).await
|
||||
}
|
||||
|
||||
pub async fn spawn_shared_connection_client(
|
||||
&self,
|
||||
db_type: &DatabaseType,
|
||||
driver_profile: Option<&str>,
|
||||
extra_java_args: &[String],
|
||||
agent_session_id: String,
|
||||
connect_params: serde_json::Value,
|
||||
connect_timeout: std::time::Duration,
|
||||
) -> Result<AgentDriverClient, String> {
|
||||
crate::agent_runtime::spawn_shared_connection_client(
|
||||
self,
|
||||
db_type,
|
||||
driver_profile,
|
||||
extra_java_args,
|
||||
agent_session_id,
|
||||
connect_params,
|
||||
connect_timeout,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn call_daemon<T: serde::de::DeserializeOwned + Send + 'static>(
|
||||
&self,
|
||||
db_type: &DatabaseType,
|
||||
|
||||
@@ -3,7 +3,7 @@ use std::time::Duration;
|
||||
|
||||
use crate::agent_manager::{AgentManager, DEFAULT_JRE_KEY};
|
||||
use crate::database_capabilities;
|
||||
use crate::db::agent_driver::{AgentDriverClient, AgentMethod};
|
||||
use crate::db::agent_driver::{AgentDriverClient, AgentMethod, AgentRuntimeClient};
|
||||
use crate::models::connection::DatabaseType;
|
||||
|
||||
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
|
||||
@@ -16,10 +16,21 @@ pub fn is_agent_type(db_type: &DatabaseType) -> bool {
|
||||
|
||||
pub async fn stop_daemons(manager: &AgentManager) {
|
||||
manager.daemons.lock().await.clear();
|
||||
let runtimes = std::mem::take(&mut *manager.connection_runtimes.lock().await);
|
||||
for runtime in runtimes.into_values().filter_map(|cell| cell.get().cloned()) {
|
||||
runtime.kill();
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn stop_daemon_by_key(manager: &AgentManager, agent_key: &str) {
|
||||
manager.daemons.lock().await.remove(agent_key);
|
||||
let mut runtimes = manager.connection_runtimes.lock().await;
|
||||
let matching = runtimes.keys().filter(|key| key.starts_with(&format!("{agent_key}|"))).cloned().collect::<Vec<_>>();
|
||||
for key in matching {
|
||||
if let Some(runtime) = runtimes.remove(&key).and_then(|cell| cell.get().cloned()) {
|
||||
runtime.kill();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn restart_daemon_by_key(manager: &AgentManager, agent_key: &str) -> Result<(), String> {
|
||||
@@ -40,6 +51,54 @@ pub async fn spawn_connection_client(
|
||||
spawn_first_available_client(manager, &keys, extra_java_args).await
|
||||
}
|
||||
|
||||
pub async fn spawn_shared_connection_client(
|
||||
manager: &AgentManager,
|
||||
db_type: &DatabaseType,
|
||||
driver_profile: Option<&str>,
|
||||
extra_java_args: &[String],
|
||||
agent_session_id: String,
|
||||
connect_params: serde_json::Value,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<AgentDriverClient, String> {
|
||||
let keys = runtime_agent_key_candidates(db_type, driver_profile)
|
||||
.ok_or_else(|| format!("{:?} is not an agent-driven database type", db_type))?;
|
||||
let key = first_installed_agent_key(manager, &keys).unwrap_or(keys[0]);
|
||||
let state = manager.load_state();
|
||||
let jre_key = state.installed_drivers.get(key).map(|driver| driver.jre.as_str()).unwrap_or(DEFAULT_JRE_KEY);
|
||||
let launch = manager.resolve_agent_launch_spec_with_extra_args(&state, key, jre_key, extra_java_args)?;
|
||||
let runtime_key = shared_runtime_key(key, &launch);
|
||||
|
||||
let runtime_cell = {
|
||||
let mut runtimes = manager.connection_runtimes.lock().await;
|
||||
if runtimes.get(&runtime_key).and_then(|cell| cell.get()).is_some_and(|runtime| runtime.is_failed()) {
|
||||
runtimes.remove(&runtime_key);
|
||||
}
|
||||
runtimes.entry(runtime_key).or_insert_with(|| std::sync::Arc::new(tokio::sync::OnceCell::new())).clone()
|
||||
};
|
||||
let runtime =
|
||||
runtime_cell.get_or_try_init(|| AgentRuntimeClient::spawn(launch, manager.agent_app_version())).await?.clone();
|
||||
let mut session_params = connect_params;
|
||||
session_params
|
||||
.as_object_mut()
|
||||
.ok_or_else(|| "Agent connect parameters must be an object".to_string())?
|
||||
.insert("agentSessionId".to_string(), serde_json::Value::String(agent_session_id.clone()));
|
||||
runtime
|
||||
.call::<serde_json::Value>(AgentMethod::OpenSession.as_str(), session_params, Some(connect_timeout), None)
|
||||
.await?;
|
||||
runtime.increment_session_count();
|
||||
Ok(AgentDriverClient::shared_session(runtime, agent_session_id))
|
||||
}
|
||||
|
||||
fn shared_runtime_key(agent_key: &str, launch: &crate::db::agent_driver::AgentLaunchSpec) -> String {
|
||||
format!(
|
||||
"{}|{}|{}|{}",
|
||||
agent_key,
|
||||
launch.program.display(),
|
||||
launch.args.join("\u{1f}"),
|
||||
launch.working_dir.as_ref().map(|path| path.display().to_string()).unwrap_or_default()
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn call_daemon<T: DeserializeOwned + Send + 'static>(
|
||||
manager: &AgentManager,
|
||||
db_type: &DatabaseType,
|
||||
@@ -167,6 +226,7 @@ async fn spawn_client_for_key(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[test]
|
||||
fn prestosql_does_not_use_agent_driver() {
|
||||
@@ -177,4 +237,17 @@ mod tests {
|
||||
fn trino_uses_only_trino_agent_driver() {
|
||||
assert_eq!(runtime_agent_key_candidates(&DatabaseType::Trino, None).unwrap(), vec!["trino"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shared_runtime_key_includes_launch_fingerprint() {
|
||||
let base = crate::db::agent_driver::AgentLaunchSpec::new(PathBuf::from("oracle-agent"))
|
||||
.with_args(["--mode", "stdio"])
|
||||
.with_working_dir(PathBuf::from("/tmp/oracle"));
|
||||
let different_args = crate::db::agent_driver::AgentLaunchSpec::new(PathBuf::from("oracle-agent"))
|
||||
.with_args(["--mode", "debug"])
|
||||
.with_working_dir(PathBuf::from("/tmp/oracle"));
|
||||
|
||||
assert_eq!(shared_runtime_key("oracle", &base), shared_runtime_key("oracle", &base));
|
||||
assert_ne!(shared_runtime_key("oracle", &base), shared_runtime_key("oracle", &different_args));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1264,71 +1264,161 @@ impl AppState {
|
||||
agent_connection_pool_database_type!() => {
|
||||
let connect_params =
|
||||
agent_connect_params(&db_config, &host, port, db_config.effective_database().unwrap_or(""));
|
||||
// Kerberos JVM properties are connection-scoped; shared agent daemons must not inherit them.
|
||||
let mut client = self
|
||||
.agent_manager
|
||||
.spawn_with_extra_java_args(
|
||||
&db_config.db_type,
|
||||
db_config.driver_profile.as_deref(),
|
||||
&db_config.agent_java_options,
|
||||
)
|
||||
.await?;
|
||||
let connect_result = client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::Connect,
|
||||
connect_params,
|
||||
Some(agent_connect_timeout(&db_config)),
|
||||
)
|
||||
.await;
|
||||
if let Err(err) = connect_result {
|
||||
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
|
||||
if !alternate_configs.is_empty() {
|
||||
log::warn!(
|
||||
if db_config.db_type != DatabaseType::Etcd && db_config.db_type != DatabaseType::ZooKeeper {
|
||||
let agent_session_id = uuid::Uuid::new_v4().simple().to_string();
|
||||
let initial_result = self
|
||||
.agent_manager
|
||||
.spawn_shared_connection_client(
|
||||
&db_config.db_type,
|
||||
db_config.driver_profile.as_deref(),
|
||||
&db_config.agent_java_options,
|
||||
agent_session_id.clone(),
|
||||
connect_params,
|
||||
agent_connect_timeout(&db_config),
|
||||
)
|
||||
.await;
|
||||
let client = match initial_result {
|
||||
Ok(client) => client,
|
||||
Err(err) => {
|
||||
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
|
||||
if alternate_configs.is_empty() {
|
||||
if err.contains("does not support multi_session protocol v2") {
|
||||
let mut client = self
|
||||
.agent_manager
|
||||
.spawn_with_extra_java_args(
|
||||
&db_config.db_type,
|
||||
db_config.driver_profile.as_deref(),
|
||||
&db_config.agent_java_options,
|
||||
)
|
||||
.await?;
|
||||
client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::Connect,
|
||||
agent_connect_params(
|
||||
&db_config,
|
||||
&host,
|
||||
port,
|
||||
db_config.effective_database().unwrap_or(""),
|
||||
),
|
||||
Some(agent_connect_timeout(&db_config)),
|
||||
)
|
||||
.await?;
|
||||
client
|
||||
} else {
|
||||
return Err(oracle_error_with_driver_hint(&db_config, &err));
|
||||
}
|
||||
} else {
|
||||
let mut fallback_errors = Vec::new();
|
||||
let mut connected = None;
|
||||
for alternate_config in alternate_configs {
|
||||
let label =
|
||||
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap_or_else(|| "alternate".to_string());
|
||||
let alternate_params = agent_connect_params(
|
||||
&alternate_config,
|
||||
&host,
|
||||
port,
|
||||
alternate_config.effective_database().unwrap_or(""),
|
||||
);
|
||||
match self
|
||||
.agent_manager
|
||||
.spawn_shared_connection_client(
|
||||
&alternate_config.db_type,
|
||||
alternate_config.driver_profile.as_deref(),
|
||||
&alternate_config.agent_java_options,
|
||||
agent_session_id.clone(),
|
||||
alternate_params,
|
||||
agent_connect_timeout(&alternate_config),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(client) => {
|
||||
connected = Some(client);
|
||||
break;
|
||||
}
|
||||
Err(alternate_err) => fallback_errors.push(format!("{label}: {alternate_err}")),
|
||||
}
|
||||
}
|
||||
connected.ok_or_else(|| {
|
||||
format!(
|
||||
"{err}\n\nFallback with alternate Oracle connection descriptors failed: {}",
|
||||
fallback_errors.join("\n")
|
||||
)
|
||||
})?
|
||||
}
|
||||
}
|
||||
};
|
||||
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
|
||||
} else {
|
||||
// Kerberos JVM properties are connection-scoped; shared agent daemons must not inherit them.
|
||||
let mut client = self
|
||||
.agent_manager
|
||||
.spawn_with_extra_java_args(
|
||||
&db_config.db_type,
|
||||
db_config.driver_profile.as_deref(),
|
||||
&db_config.agent_java_options,
|
||||
)
|
||||
.await?;
|
||||
let connect_result = client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::Connect,
|
||||
connect_params,
|
||||
Some(agent_connect_timeout(&db_config)),
|
||||
)
|
||||
.await;
|
||||
if let Err(err) = connect_result {
|
||||
let alternate_configs = oracle_alternate_connect_configs(&db_config, &err);
|
||||
if !alternate_configs.is_empty() {
|
||||
log::warn!(
|
||||
"Oracle connect failed with {:?} descriptor: {}. Retrying with Oracle JDBC URL variants: {:?}.",
|
||||
db_config.oracle_connection_type,
|
||||
err,
|
||||
oracle_alternate_connect_config_labels(&alternate_configs)
|
||||
);
|
||||
let mut fallback_errors = Vec::new();
|
||||
let mut connected = false;
|
||||
for alternate_config in alternate_configs {
|
||||
let label = oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap_or_else(|| "alternate".to_string());
|
||||
match client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::Connect,
|
||||
agent_connect_params(
|
||||
&alternate_config,
|
||||
&host,
|
||||
port,
|
||||
alternate_config.effective_database().unwrap_or(""),
|
||||
),
|
||||
Some(agent_connect_timeout(&alternate_config)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
connected = true;
|
||||
break;
|
||||
}
|
||||
Err(alternate_err) => {
|
||||
fallback_errors.push(format!("{label}: {alternate_err}"));
|
||||
let mut fallback_errors = Vec::new();
|
||||
let mut connected = false;
|
||||
for alternate_config in alternate_configs {
|
||||
let label =
|
||||
oracle_alternate_connect_config_labels(std::slice::from_ref(&alternate_config))
|
||||
.into_iter()
|
||||
.next()
|
||||
.unwrap_or_else(|| "alternate".to_string());
|
||||
match client
|
||||
.call_method_with_timeout::<serde_json::Value>(
|
||||
AgentMethod::Connect,
|
||||
agent_connect_params(
|
||||
&alternate_config,
|
||||
&host,
|
||||
port,
|
||||
alternate_config.effective_database().unwrap_or(""),
|
||||
),
|
||||
Some(agent_connect_timeout(&alternate_config)),
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(_) => {
|
||||
connected = true;
|
||||
break;
|
||||
}
|
||||
Err(alternate_err) => {
|
||||
fallback_errors.push(format!("{label}: {alternate_err}"));
|
||||
}
|
||||
}
|
||||
}
|
||||
if !connected {
|
||||
return Err(format!(
|
||||
"{err}\n\nFallback with alternate Oracle JDBC URLs failed: {}",
|
||||
fallback_errors.join("\n")
|
||||
));
|
||||
}
|
||||
} else {
|
||||
return Err(oracle_error_with_driver_hint(&db_config, &err));
|
||||
}
|
||||
if !connected {
|
||||
return Err(format!(
|
||||
"{err}\n\nFallback with alternate Oracle JDBC URLs failed: {}",
|
||||
fallback_errors.join("\n")
|
||||
));
|
||||
}
|
||||
} else {
|
||||
return Err(oracle_error_with_driver_hint(&db_config, &err));
|
||||
}
|
||||
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
|
||||
}
|
||||
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client)))
|
||||
}
|
||||
DatabaseType::PrestoSql => {
|
||||
let jdbc_config = prestosql_jdbc_config_for_endpoint(&db_config, &host, port);
|
||||
@@ -4003,6 +4093,17 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oracle_uses_isolated_pool_keys_for_tab_sessions() {
|
||||
let mut config = mysql_config(Some("ORCLPDB1"));
|
||||
config.db_type = DatabaseType::Oracle;
|
||||
|
||||
assert_eq!(
|
||||
super::session_scoped_pool_key_for(Some(&config), "oracle-conn".to_string(), Some("table-tab-1")),
|
||||
"oracle-conn:session:table-tab-1"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn other_agent_single_connection_types_keep_database_scoped_pool_keys() {
|
||||
assert_eq!(
|
||||
@@ -4338,6 +4439,25 @@ mod tests {
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closing_oracle_table_tab_keeps_connection_scoped_pool() {
|
||||
let (state, dir) = test_app_state().await;
|
||||
let mut config = mysql_config(Some("ORCLPDB1"));
|
||||
config.id = "oracle-conn".to_string();
|
||||
config.db_type = DatabaseType::Oracle;
|
||||
state.configs.write().await.insert(config.id.clone(), config);
|
||||
|
||||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||||
state.connections.write().await.insert("oracle-conn".to_string(), PoolKind::Sqlite(pool));
|
||||
state.pool_activity.write().await.insert("oracle-conn".to_string(), super::PoolActivity::now());
|
||||
|
||||
assert!(!state.close_client_session_pool("oracle-conn", Some("APP"), "table-tab-1").await.unwrap());
|
||||
assert!(state.connections.read().await.contains_key("oracle-conn"));
|
||||
assert!(state.pool_activity.read().await.contains_key("oracle-conn"));
|
||||
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_validate_connection_unknown_method_is_not_stale() {
|
||||
assert!(super::is_agent_validate_connection_unsupported(
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
use std::collections::VecDeque;
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::io::{BufRead, BufReader, BufWriter, Write};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Child, ChildStderr, ChildStdin, ChildStdout, Command, Stdio};
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
@@ -10,23 +11,261 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
pub const AGENT_PROTOCOL_VERSION: u32 = 1;
|
||||
type PendingAgentResponse = tokio::sync::oneshot::Sender<Result<Value, String>>;
|
||||
|
||||
pub struct AgentRuntimeClient {
|
||||
child: Arc<Mutex<Child>>,
|
||||
stdin: Arc<Mutex<BufWriter<ChildStdin>>>,
|
||||
pending: Arc<Mutex<HashMap<u64, PendingAgentResponse>>>,
|
||||
stderr_tail: Arc<Mutex<StderrTail>>,
|
||||
next_id: AtomicU64,
|
||||
active_sessions: AtomicU64,
|
||||
failed: Arc<AtomicBool>,
|
||||
handshake: AgentHandshake,
|
||||
}
|
||||
|
||||
impl AgentRuntimeClient {
|
||||
pub async fn spawn(launch: AgentLaunchSpec, app_version: &str) -> Result<Arc<Self>, String> {
|
||||
let mut command = crate::process::new_std_command(&launch.program);
|
||||
command.args(&launch.args).stdin(Stdio::piped()).stdout(Stdio::piped()).stderr(Stdio::piped());
|
||||
if let Some(working_dir) = &launch.working_dir {
|
||||
command.current_dir(working_dir);
|
||||
}
|
||||
remove_agent_proxy_env(&mut command);
|
||||
|
||||
let mut child =
|
||||
command.spawn().map_err(|e| format!("Failed to spawn agent process {}: {e}", launch_display(&launch)))?;
|
||||
let child_stdin = child.stdin.take().ok_or("Failed to capture agent stdin")?;
|
||||
let child_stdout = child.stdout.take().ok_or("Failed to capture agent stdout")?;
|
||||
let child_stderr = child.stderr.take().ok_or("Failed to capture agent stderr")?;
|
||||
let stderr_tail = Arc::new(Mutex::new(StderrTail::default()));
|
||||
start_stderr_collector(child_stderr, stderr_tail.clone());
|
||||
|
||||
let mut stdout = BufReader::new(child_stdout);
|
||||
let stdout = tokio::time::timeout(
|
||||
Duration::from_secs(STARTUP_TIMEOUT_SECS),
|
||||
tokio::task::spawn_blocking(move || loop {
|
||||
let line = read_agent_line(&mut stdout, "startup line")?;
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match serde_json::from_str::<Value>(trimmed) {
|
||||
Ok(value) if value.get("ready") == Some(&Value::Bool(true)) => return Ok(stdout),
|
||||
Ok(_) => return Err(format!("Agent did not send ready signal, got: {line}")),
|
||||
Err(_) => log::warn!("[agent:stdout] ignoring non-JSON line during startup: {trimmed}"),
|
||||
}
|
||||
}),
|
||||
)
|
||||
.await
|
||||
.map_err(|_| format!("Agent startup timed out ({STARTUP_TIMEOUT_SECS}s)"))?
|
||||
.map_err(|e| format!("Agent startup task failed: {e}"))??;
|
||||
|
||||
let runtime = Arc::new(Self {
|
||||
child: Arc::new(Mutex::new(child)),
|
||||
stdin: Arc::new(Mutex::new(BufWriter::new(child_stdin))),
|
||||
pending: Arc::new(Mutex::new(HashMap::new())),
|
||||
stderr_tail,
|
||||
next_id: AtomicU64::new(0),
|
||||
active_sessions: AtomicU64::new(0),
|
||||
failed: Arc::new(AtomicBool::new(false)),
|
||||
handshake: AgentHandshake { protocol_version: 0, agent_protocol_version: 0, capabilities: Vec::new() },
|
||||
});
|
||||
runtime.start_response_reader(stdout);
|
||||
let handshake = runtime
|
||||
.call::<AgentHandshake>(
|
||||
AgentMethod::Handshake.as_str(),
|
||||
agent_handshake_params(app_version),
|
||||
Some(Duration::from_secs(RPC_TIMEOUT_SECS)),
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
if handshake.protocol_version < 2 || !handshake.supports(AgentCapability::MultiSession) {
|
||||
runtime.kill();
|
||||
return Err("Agent runtime does not support multi_session protocol v2".to_string());
|
||||
}
|
||||
let runtime =
|
||||
Arc::try_unwrap(runtime).map_err(|_| "Agent runtime initialization is still referenced".to_string())?;
|
||||
Ok(Arc::new(Self { handshake, ..runtime }))
|
||||
}
|
||||
|
||||
fn start_response_reader(self: &Arc<Self>, mut stdout: BufReader<ChildStdout>) {
|
||||
let pending = self.pending.clone();
|
||||
let failed = self.failed.clone();
|
||||
std::thread::spawn(move || loop {
|
||||
let line = match read_agent_line(&mut stdout, "response") {
|
||||
Ok(line) => line,
|
||||
Err(err) => {
|
||||
failed.store(true, Ordering::Release);
|
||||
fail_pending_requests(&pending, err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let response: Value = match serde_json::from_str(line.trim()) {
|
||||
Ok(response) => response,
|
||||
Err(err) => {
|
||||
failed.store(true, Ordering::Release);
|
||||
fail_pending_requests(&pending, format!("Invalid JSON response from agent: {err}"));
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(id) = response.get("id").and_then(Value::as_u64) else {
|
||||
continue;
|
||||
};
|
||||
if let Some(sender) = pending.lock().expect("agent pending response lock poisoned").remove(&id) {
|
||||
let _ = sender.send(Ok(response));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub async fn call<T: DeserializeOwned>(
|
||||
&self,
|
||||
method: &str,
|
||||
params: Value,
|
||||
timeout_duration: Option<Duration>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
) -> Result<T, String> {
|
||||
if self.failed.load(Ordering::Acquire) {
|
||||
return Err("Agent runtime is unavailable".to_string());
|
||||
}
|
||||
let id = self.next_id.fetch_add(1, Ordering::Relaxed) + 1;
|
||||
let request = serde_json::json!({ "jsonrpc": "2.0", "id": id, "method": method, "params": params });
|
||||
let request_line =
|
||||
serde_json::to_string(&request).map_err(|e| format!("Failed to serialize JSON-RPC request: {e}"))?;
|
||||
let (sender, receiver) = tokio::sync::oneshot::channel();
|
||||
self.pending.lock().map_err(|_| "Agent pending response lock poisoned".to_string())?.insert(id, sender);
|
||||
let write_result =
|
||||
self.stdin.lock().map_err(|_| "Agent stdin lock poisoned".to_string()).and_then(|mut writer| {
|
||||
writer
|
||||
.write_all(request_line.as_bytes())
|
||||
.and_then(|_| writer.write_all(b"\n"))
|
||||
.and_then(|_| writer.flush())
|
||||
.map_err(|e| format!("Failed to write agent request: {e}"))
|
||||
});
|
||||
if let Err(err) = write_result {
|
||||
self.pending.lock().expect("agent pending response lock poisoned").remove(&id);
|
||||
return Err(err);
|
||||
}
|
||||
|
||||
let receive = async { receiver.await.map_err(|_| "Agent response channel closed".to_string())? };
|
||||
let response = match (timeout_duration, cancel_token) {
|
||||
(Some(duration), Some(token)) => tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
self.cancel_session_request(¶ms).await;
|
||||
Err("Query canceled".to_string())
|
||||
},
|
||||
result = tokio::time::timeout(duration, receive) => result.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
|
||||
},
|
||||
(Some(duration), None) => tokio::time::timeout(duration, receive)
|
||||
.await
|
||||
.map_err(|_| format!("Agent RPC call timed out ({}s)", duration.as_secs()))?,
|
||||
(None, Some(token)) => tokio::select! {
|
||||
_ = token.cancelled() => {
|
||||
self.cancel_session_request(¶ms).await;
|
||||
Err("Query canceled".to_string())
|
||||
},
|
||||
result = receive => result,
|
||||
},
|
||||
(None, None) => receive.await,
|
||||
};
|
||||
if response.is_err() {
|
||||
self.pending.lock().expect("agent pending response lock poisoned").remove(&id);
|
||||
}
|
||||
decode_agent_response(response?)
|
||||
}
|
||||
|
||||
async fn cancel_session_request(&self, params: &Value) {
|
||||
let Some(agent_session_id) = params.get("agentSessionId").and_then(Value::as_str) else {
|
||||
return;
|
||||
};
|
||||
let _ = Box::pin(self.call::<Value>(
|
||||
AgentMethod::CancelSession.as_str(),
|
||||
serde_json::json!({ "agentSessionId": agent_session_id }),
|
||||
Some(Duration::from_secs(5)),
|
||||
None,
|
||||
))
|
||||
.await;
|
||||
}
|
||||
|
||||
pub fn handshake(&self) -> &AgentHandshake {
|
||||
&self.handshake
|
||||
}
|
||||
|
||||
pub fn is_failed(&self) -> bool {
|
||||
self.failed.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub fn active_session_count(&self) -> u64 {
|
||||
self.active_sessions.load(Ordering::Acquire)
|
||||
}
|
||||
|
||||
pub fn increment_session_count(&self) {
|
||||
self.active_sessions.fetch_add(1, Ordering::AcqRel);
|
||||
}
|
||||
|
||||
pub fn decrement_session_count(runtime: &Arc<Self>) {
|
||||
let previous = runtime
|
||||
.active_sessions
|
||||
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |value| Some(value.saturating_sub(1)))
|
||||
.unwrap_or_default();
|
||||
if previous <= 1 {
|
||||
let runtime = runtime.clone();
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_secs(SHARED_RUNTIME_IDLE_GRACE_SECS)).await;
|
||||
if runtime.active_session_count() == 0 {
|
||||
runtime.kill();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
pub fn kill(&self) {
|
||||
self.failed.store(true, Ordering::Release);
|
||||
if let Ok(mut child) = self.child.lock() {
|
||||
let _ = child.kill();
|
||||
}
|
||||
fail_pending_requests(&self.pending, "Agent runtime terminated".to_string());
|
||||
}
|
||||
}
|
||||
|
||||
fn decode_agent_response<T: DeserializeOwned>(response: Value) -> Result<T, String> {
|
||||
if let Some(err) = response.get("error") {
|
||||
let message = err.get("message").and_then(Value::as_str).unwrap_or("Unknown agent error");
|
||||
let code = err.get("code").and_then(Value::as_i64).unwrap_or(-1);
|
||||
return Err(format!("Agent RPC error ({code}): {message}"));
|
||||
}
|
||||
let result =
|
||||
response.get("result").ok_or_else(|| "Agent response missing both 'result' and 'error'".to_string())?;
|
||||
serde_json::from_value(result.clone()).map_err(|e| format!("Failed to deserialize agent result: {e}"))
|
||||
}
|
||||
|
||||
fn fail_pending_requests(pending: &Arc<Mutex<HashMap<u64, PendingAgentResponse>>>, error: String) {
|
||||
let requests = std::mem::take(&mut *pending.lock().expect("agent pending response lock poisoned"));
|
||||
for (_, sender) in requests {
|
||||
let _ = sender.send(Err(error.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
pub const AGENT_PROTOCOL_VERSION: u32 = 2;
|
||||
const RPC_TIMEOUT_SECS: u64 = 30;
|
||||
const STARTUP_TIMEOUT_SECS: u64 = 15;
|
||||
const STDERR_TAIL_LINES: usize = 20;
|
||||
const AGENT_EXIT_DIAGNOSTIC_WAIT_MS: u64 = 1_000;
|
||||
const AGENT_EXIT_DIAGNOSTIC_POLL_MS: u64 = 10;
|
||||
const SHARED_RUNTIME_IDLE_GRACE_SECS: u64 = 30;
|
||||
const AGENT_JAVA_OPTS_ENV: &str = "DBX_AGENT_JAVA_OPTS";
|
||||
const AGENT_JAVA_TOO_OLD_MESSAGE: &str =
|
||||
"Agent requires Java 21, but DBX started it with an older Java runtime. Use DBX managed JRE 21 or select a Java 21 executable in Driver Manager.";
|
||||
|
||||
pub struct AgentDriverClient {
|
||||
child: Child,
|
||||
child: Option<Child>,
|
||||
stdin: Option<BufWriter<ChildStdin>>,
|
||||
stdout: Option<BufReader<ChildStdout>>,
|
||||
stderr_tail: Arc<Mutex<StderrTail>>,
|
||||
handshake: Option<AgentHandshake>,
|
||||
next_id: u64,
|
||||
shared_runtime: Option<Arc<AgentRuntimeClient>>,
|
||||
agent_session_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -93,10 +332,11 @@ pub enum AgentCapability {
|
||||
Transaction,
|
||||
Ddl,
|
||||
Kv,
|
||||
MultiSession,
|
||||
}
|
||||
|
||||
impl AgentCapability {
|
||||
pub const ALL: [Self; 8] = [
|
||||
pub const ALL: [Self; 9] = [
|
||||
Self::Connect,
|
||||
Self::TestConnection,
|
||||
Self::Metadata,
|
||||
@@ -105,6 +345,7 @@ impl AgentCapability {
|
||||
Self::Transaction,
|
||||
Self::Ddl,
|
||||
Self::Kv,
|
||||
Self::MultiSession,
|
||||
];
|
||||
|
||||
pub fn as_str(self) -> &'static str {
|
||||
@@ -117,6 +358,7 @@ impl AgentCapability {
|
||||
Self::Transaction => "transaction",
|
||||
Self::Ddl => "ddl",
|
||||
Self::Kv => "kv",
|
||||
Self::MultiSession => "multi_session",
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -125,6 +367,10 @@ impl AgentCapability {
|
||||
pub enum AgentMethod {
|
||||
Handshake,
|
||||
Connect,
|
||||
OpenSession,
|
||||
CloseSession,
|
||||
ValidateSession,
|
||||
CancelSession,
|
||||
TestConnection,
|
||||
ValidateConnection,
|
||||
ListDatabases,
|
||||
@@ -154,9 +400,13 @@ pub enum AgentMethod {
|
||||
}
|
||||
|
||||
impl AgentMethod {
|
||||
pub const ALL: [Self; 28] = [
|
||||
pub const ALL: [Self; 32] = [
|
||||
Self::Handshake,
|
||||
Self::Connect,
|
||||
Self::OpenSession,
|
||||
Self::CloseSession,
|
||||
Self::ValidateSession,
|
||||
Self::CancelSession,
|
||||
Self::TestConnection,
|
||||
Self::ValidateConnection,
|
||||
Self::ListDatabases,
|
||||
@@ -189,6 +439,10 @@ impl AgentMethod {
|
||||
match self {
|
||||
Self::Handshake => "handshake",
|
||||
Self::Connect => "connect",
|
||||
Self::OpenSession => "open_session",
|
||||
Self::CloseSession => "close_session",
|
||||
Self::ValidateSession => "validate_session",
|
||||
Self::CancelSession => "cancel_session",
|
||||
Self::TestConnection => "test_connection",
|
||||
Self::ValidateConnection => "validate_connection",
|
||||
Self::ListDatabases => "list_databases",
|
||||
@@ -425,7 +679,29 @@ impl AgentDriverClient {
|
||||
}
|
||||
};
|
||||
|
||||
Ok(Self { child, stdin: Some(stdin), stdout: Some(ready_stdout), stderr_tail, handshake: None, next_id: 0 })
|
||||
Ok(Self {
|
||||
child: Some(child),
|
||||
stdin: Some(stdin),
|
||||
stdout: Some(ready_stdout),
|
||||
stderr_tail,
|
||||
handshake: None,
|
||||
next_id: 0,
|
||||
shared_runtime: None,
|
||||
agent_session_id: None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn shared_session(runtime: Arc<AgentRuntimeClient>, agent_session_id: String) -> Self {
|
||||
Self {
|
||||
child: None,
|
||||
stdin: None,
|
||||
stdout: None,
|
||||
stderr_tail: runtime.stderr_tail.clone(),
|
||||
handshake: Some(runtime.handshake().clone()),
|
||||
next_id: 0,
|
||||
shared_runtime: Some(runtime),
|
||||
agent_session_id: Some(agent_session_id),
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a JSON-RPC 2.0 request and wait for the response.
|
||||
@@ -458,6 +734,20 @@ impl AgentDriverClient {
|
||||
timeout_duration: Option<Duration>,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
) -> Result<T, String> {
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
let mut params = params;
|
||||
if method != AgentMethod::Handshake.as_str()
|
||||
&& method != AgentMethod::TestConnection.as_str()
|
||||
&& method != AgentMethod::Shutdown.as_str()
|
||||
{
|
||||
let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?;
|
||||
params
|
||||
.as_object_mut()
|
||||
.ok_or_else(|| "Agent RPC parameters must be an object".to_string())?
|
||||
.insert("agentSessionId".to_string(), Value::String(session_id.clone()));
|
||||
}
|
||||
return runtime.call(method, params, timeout_duration, cancel_token).await;
|
||||
}
|
||||
self.next_id += 1;
|
||||
let id = self.next_id;
|
||||
|
||||
@@ -587,6 +877,31 @@ impl AgentDriverClient {
|
||||
self.call_method(AgentMethod::Connect, params).await
|
||||
}
|
||||
|
||||
pub async fn open_session(&mut self, agent_session_id: &str, mut params: Value) -> Result<Value, String> {
|
||||
params
|
||||
.as_object_mut()
|
||||
.ok_or_else(|| "Agent session parameters must be an object".to_string())?
|
||||
.insert("agentSessionId".to_string(), Value::String(agent_session_id.to_string()));
|
||||
self.call_method(AgentMethod::OpenSession, params).await
|
||||
}
|
||||
|
||||
pub async fn close_session(&mut self, agent_session_id: &str) -> Result<Value, String> {
|
||||
self.call_method(AgentMethod::CloseSession, serde_json::json!({ "agentSessionId": agent_session_id })).await
|
||||
}
|
||||
|
||||
pub async fn validate_session(
|
||||
&mut self,
|
||||
agent_session_id: &str,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<Value, String> {
|
||||
self.call_method_with_timeout(
|
||||
AgentMethod::ValidateSession,
|
||||
serde_json::json!({ "agentSessionId": agent_session_id }),
|
||||
timeout_duration,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn test_connection(&mut self, params: Value) -> Result<Value, String> {
|
||||
self.call_method(AgentMethod::TestConnection, params).await
|
||||
}
|
||||
@@ -596,6 +911,18 @@ impl AgentDriverClient {
|
||||
}
|
||||
|
||||
pub async fn disconnect(&mut self) -> Result<Value, String> {
|
||||
if self.shared_runtime.is_some() {
|
||||
let session_id = self.agent_session_id.as_ref().ok_or("Shared Agent session id is missing")?.clone();
|
||||
let result =
|
||||
self.call_method(AgentMethod::CloseSession, serde_json::json!({ "agentSessionId": session_id })).await;
|
||||
if result.is_ok() {
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
AgentRuntimeClient::decrement_session_count(runtime);
|
||||
}
|
||||
self.agent_session_id = None;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
self.call_method(AgentMethod::Disconnect, serde_json::json!({})).await
|
||||
}
|
||||
|
||||
@@ -1105,6 +1432,10 @@ impl AgentDriverClient {
|
||||
|
||||
/// Send a shutdown message to the agent and wait for the process to exit.
|
||||
pub async fn shutdown(&mut self) {
|
||||
if self.shared_runtime.is_some() {
|
||||
let _ = self.disconnect().await;
|
||||
return;
|
||||
}
|
||||
// Try to send a shutdown RPC; ignore errors if the agent is already gone
|
||||
let shutdown_result: Result<Value, String> = self.call_method(AgentMethod::Shutdown, Value::Null).await;
|
||||
if let Err(e) = &shutdown_result {
|
||||
@@ -1115,7 +1446,8 @@ impl AgentDriverClient {
|
||||
self.stdin.take();
|
||||
|
||||
// Wait for the child to exit
|
||||
match self.child.wait() {
|
||||
let Some(child) = self.child.as_mut() else { return };
|
||||
match child.wait() {
|
||||
Ok(status) => log::info!("Agent process exited with {status}"),
|
||||
Err(e) => log::warn!("Failed to wait for agent process: {e}"),
|
||||
}
|
||||
@@ -1123,9 +1455,14 @@ impl AgentDriverClient {
|
||||
|
||||
/// Forcefully kill the agent process.
|
||||
pub fn kill(&mut self) {
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
runtime.kill();
|
||||
return;
|
||||
}
|
||||
self.stdin.take();
|
||||
self.stdout.take();
|
||||
if let Err(e) = self.child.kill() {
|
||||
let Some(child) = self.child.as_mut() else { return };
|
||||
if let Err(e) = child.kill() {
|
||||
log::warn!("Failed to kill agent process: {e}");
|
||||
}
|
||||
// Reap the child to avoid zombie processes.
|
||||
@@ -1133,7 +1470,7 @@ impl AgentDriverClient {
|
||||
// hanging in Drop during async cleanup. Poll up to 100ms for the
|
||||
// process to exit after kill().
|
||||
for _ in 0..10 {
|
||||
match self.child.try_wait() {
|
||||
match child.try_wait() {
|
||||
Ok(Some(_status)) => return,
|
||||
Ok(None) => std::thread::sleep(Duration::from_millis(10)),
|
||||
Err(e) => {
|
||||
@@ -1143,13 +1480,28 @@ impl AgentDriverClient {
|
||||
}
|
||||
}
|
||||
// Final blocking wait as a last resort
|
||||
if let Err(e) = self.child.wait() {
|
||||
if let Err(e) = child.wait() {
|
||||
log::warn!("Final wait failed for agent process: {e}");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn pid(&self) -> u32 {
|
||||
self.child.id()
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
return runtime.child.lock().map(|child| child.id()).unwrap_or_default();
|
||||
}
|
||||
self.child.as_ref().map(Child::id).unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn protocol_mode(&self) -> &'static str {
|
||||
if self.shared_runtime.is_some() {
|
||||
"multi_session"
|
||||
} else {
|
||||
"legacy"
|
||||
}
|
||||
}
|
||||
|
||||
pub fn active_session_count(&self) -> u64 {
|
||||
self.shared_runtime.as_ref().map(|runtime| runtime.active_session_count()).unwrap_or(1)
|
||||
}
|
||||
|
||||
pub fn stderr_tail_snapshot(&self) -> String {
|
||||
@@ -1438,13 +1790,22 @@ impl AgentDriverClient {
|
||||
fn format_agent_process_error(&mut self, base: &str) -> String {
|
||||
// Runtime RPC errors are common SQL/driver paths. Do not wait for the
|
||||
// child to exit unless startup diagnostics already expect the process to die.
|
||||
format_agent_process_error(base, child_exit_status(&mut self.child), &stderr_tail_snapshot(&self.stderr_tail))
|
||||
let exit_status = self.child.as_mut().and_then(child_exit_status);
|
||||
format_agent_process_error(base, exit_status, &stderr_tail_snapshot(&self.stderr_tail))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for AgentDriverClient {
|
||||
fn drop(&mut self) {
|
||||
self.kill();
|
||||
// Shared-session clients do not own the runtime process. Session closure
|
||||
// is explicit in pool cleanup; dropping one tab must not kill other tabs.
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
if self.agent_session_id.take().is_some() {
|
||||
AgentRuntimeClient::decrement_session_count(runtime);
|
||||
}
|
||||
} else {
|
||||
self.kill();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1456,12 +1817,16 @@ mod tests {
|
||||
agent_schema_table_params, agent_supports_capability, agent_transaction_params, format_agent_process_error,
|
||||
format_agent_startup_error, is_unsupported_handshake_error, mongo_collection_params, mongo_database_params,
|
||||
mongo_document_id_params, parse_agent_java_opts, read_agent_line, start_stderr_collector, AgentCapability,
|
||||
AgentDriverClient, AgentHandshake, AgentKvMethod, AgentMethod, AgentTableReadCloseParams,
|
||||
AgentTableReadPageParams, AgentTableReadStartParams, MongoAgentMethod, StderrTail, AGENT_PROTOCOL_VERSION,
|
||||
AgentDriverClient, AgentHandshake, AgentKvMethod, AgentLaunchSpec, AgentMethod, AgentRuntimeClient,
|
||||
AgentTableReadCloseParams, AgentTableReadPageParams, AgentTableReadStartParams, MongoAgentMethod, StderrTail,
|
||||
AGENT_PROTOCOL_VERSION,
|
||||
};
|
||||
use std::io::Cursor;
|
||||
use std::io::Write;
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Duration;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[test]
|
||||
fn agent_java_args_include_oracle_network_compatibility_flags() {
|
||||
@@ -1638,12 +2003,14 @@ mod tests {
|
||||
.spawn()
|
||||
.expect("child should start");
|
||||
let mut client = AgentDriverClient {
|
||||
child,
|
||||
child: Some(child),
|
||||
stdin: None,
|
||||
stdout: None,
|
||||
stderr_tail: Arc::new(Mutex::new(StderrTail::default())),
|
||||
handshake: None,
|
||||
next_id: 0,
|
||||
shared_runtime: None,
|
||||
agent_session_id: None,
|
||||
};
|
||||
|
||||
let started_at = std::time::Instant::now();
|
||||
@@ -1654,6 +2021,129 @@ mod tests {
|
||||
assert!(!message.contains("agent process exited"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn multiplexed_runtime_correlates_out_of_order_responses() {
|
||||
let script_path = std::env::temp_dir().join(format!("dbx-agent-runtime-test-{}.py", uuid::Uuid::new_v4()));
|
||||
let mut script = std::fs::File::create(&script_path).unwrap();
|
||||
script
|
||||
.write_all(
|
||||
br#"import json, sys, threading, time
|
||||
print(json.dumps({'ready': True}), flush=True)
|
||||
def respond(req):
|
||||
if req['method'] == 'handshake':
|
||||
result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
|
||||
else:
|
||||
time.sleep(0.05 if req['params']['value'] == 1 else 0.0)
|
||||
result = req['params']['value']
|
||||
print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': result}), flush=True)
|
||||
for line in sys.stdin:
|
||||
threading.Thread(target=respond, args=(json.loads(line),), daemon=True).start()
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
drop(script);
|
||||
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let first = runtime.call::<u64>("echo", serde_json::json!({"value": 1}), Some(Duration::from_secs(2)), None);
|
||||
let second = runtime.call::<u64>("echo", serde_json::json!({"value": 2}), Some(Duration::from_secs(2)), None);
|
||||
let (first, second) = tokio::join!(first, second);
|
||||
|
||||
assert_eq!(first.unwrap(), 1);
|
||||
assert_eq!(second.unwrap(), 2);
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn canceling_one_runtime_request_keeps_other_requests_alive() {
|
||||
let script_path = std::env::temp_dir().join(format!("dbx-agent-cancel-test-{}.py", uuid::Uuid::new_v4()));
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
r#"import json, sys, threading, time
|
||||
print(json.dumps({'ready': True}), flush=True)
|
||||
def respond(req):
|
||||
if req['method'] == 'handshake':
|
||||
result = {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
|
||||
else:
|
||||
time.sleep(req['params'].get('delay', 0))
|
||||
result = req['params'].get('value')
|
||||
print(json.dumps({'jsonrpc': '2.0', 'id': req['id'], 'result': result}), flush=True)
|
||||
for line in sys.stdin:
|
||||
threading.Thread(target=respond, args=(json.loads(line),), daemon=True).start()
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let token = CancellationToken::new();
|
||||
let canceled = runtime.call::<u64>(
|
||||
"echo",
|
||||
serde_json::json!({"value": 1, "delay": 0.2}),
|
||||
Some(Duration::from_secs(2)),
|
||||
Some(token.clone()),
|
||||
);
|
||||
let healthy = runtime.call::<u64>(
|
||||
"echo",
|
||||
serde_json::json!({"value": 2, "delay": 0}),
|
||||
Some(Duration::from_secs(2)),
|
||||
None,
|
||||
);
|
||||
token.cancel();
|
||||
let (canceled, healthy) = tokio::join!(canceled, healthy);
|
||||
assert_eq!(canceled.unwrap_err(), "Query canceled");
|
||||
assert_eq!(healthy.unwrap(), 2);
|
||||
assert_eq!(
|
||||
runtime
|
||||
.call::<u64>("echo", serde_json::json!({"value": 3}), Some(Duration::from_secs(2)), None)
|
||||
.await
|
||||
.unwrap(),
|
||||
3
|
||||
);
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn runtime_eof_fails_pending_requests_and_marks_runtime_failed() {
|
||||
let script_path = std::env::temp_dir().join(format!("dbx-agent-eof-test-{}.py", uuid::Uuid::new_v4()));
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
r#"import json, sys
|
||||
print(json.dumps({'ready': True}), flush=True)
|
||||
for line in sys.stdin:
|
||||
req = json.loads(line)
|
||||
if req['method'] == 'handshake':
|
||||
print(json.dumps({'jsonrpc':'2.0','id':req['id'],'result':{'protocolVersion':2,'agentProtocolVersion':2,'capabilities':['multi_session']}}), flush=True)
|
||||
else:
|
||||
sys.exit(0)
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = runtime
|
||||
.call::<serde_json::Value>("crash", serde_json::json!({}), Some(Duration::from_secs(2)), None)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.contains("end of stream") || error.contains("response channel closed"));
|
||||
assert!(runtime.is_failed());
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stderr_tail_keeps_recent_lines_only() {
|
||||
let mut stderr_tail = StderrTail::with_capacity(3);
|
||||
@@ -1697,7 +2187,8 @@ mod tests {
|
||||
assert_eq!(AgentCapability::Transaction.as_str(), "transaction");
|
||||
assert_eq!(AgentCapability::Ddl.as_str(), "ddl");
|
||||
assert_eq!(AgentCapability::Kv.as_str(), "kv");
|
||||
assert_eq!(AgentCapability::ALL.len(), 8);
|
||||
assert_eq!(AgentCapability::MultiSession.as_str(), "multi_session");
|
||||
assert_eq!(AgentCapability::ALL.len(), 9);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -1915,7 +2406,7 @@ mod tests {
|
||||
#[test]
|
||||
fn agent_protocol_matches_contract_file() {
|
||||
let contract: serde_json::Value =
|
||||
serde_json::from_str(include_str!("../../assets/agent-protocol-v1.json")).unwrap();
|
||||
serde_json::from_str(include_str!("../../assets/agent-protocol-v2.json")).unwrap();
|
||||
|
||||
assert_eq!(contract["protocolVersion"], AGENT_PROTOCOL_VERSION);
|
||||
assert_eq!(contract["handshakeMethod"], AgentMethod::Handshake.as_str());
|
||||
@@ -1991,6 +2482,7 @@ mod tests {
|
||||
AgentCapability::PagedQuery,
|
||||
AgentCapability::Transaction,
|
||||
AgentCapability::Ddl,
|
||||
AgentCapability::MultiSession,
|
||||
]
|
||||
.iter()
|
||||
.map(|capability| capability.as_str())
|
||||
|
||||
@@ -23,6 +23,10 @@ pub struct DriverRuntimeInfo {
|
||||
pub can_stop: bool,
|
||||
pub can_restart: bool,
|
||||
pub control_unavailable_reason: Option<String>,
|
||||
#[serde(default)]
|
||||
pub protocol_mode: Option<String>,
|
||||
#[serde(default)]
|
||||
pub active_sessions: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
|
||||
@@ -48,6 +52,8 @@ struct RuntimeSeed {
|
||||
can_stop: bool,
|
||||
can_restart: bool,
|
||||
control_unavailable_reason: Option<String>,
|
||||
protocol_mode: Option<String>,
|
||||
active_sessions: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
@@ -65,16 +71,27 @@ pub async fn collect_driver_runtime_summary(state: &AppState) -> DriverRuntimeSu
|
||||
}
|
||||
|
||||
pub async fn stop_driver_runtime(state: &AppState, runtime_id: &str) -> Result<(), String> {
|
||||
let agent_key =
|
||||
runtime_id.strip_prefix("agent:").ok_or_else(|| "Only agent daemon runtimes can be stopped".to_string())?;
|
||||
let agent_key = runtime_agent_key(runtime_id).ok_or_else(|| "Only agent runtimes can be stopped".to_string())?;
|
||||
state.agent_manager.stop_daemon_by_key(agent_key).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn restart_driver_runtime(state: &AppState, runtime_id: &str) -> Result<(), String> {
|
||||
let agent_key =
|
||||
runtime_id.strip_prefix("agent:").ok_or_else(|| "Only agent daemon runtimes can be restarted".to_string())?;
|
||||
state.agent_manager.restart_daemon_by_key(agent_key).await
|
||||
let agent_key = runtime_agent_key(runtime_id).ok_or_else(|| "Only agent runtimes can be restarted".to_string())?;
|
||||
if runtime_id.starts_with("agent-runtime:") {
|
||||
// Shared connection runtimes are recreated lazily with their exact launch
|
||||
// fingerprint; eagerly spawning the generic daemon could use incompatible options.
|
||||
state.agent_manager.stop_daemon_by_key(agent_key).await;
|
||||
Ok(())
|
||||
} else {
|
||||
state.agent_manager.restart_daemon_by_key(agent_key).await
|
||||
}
|
||||
}
|
||||
|
||||
fn runtime_agent_key(runtime_id: &str) -> Option<&str> {
|
||||
runtime_id.strip_prefix("agent:").or_else(|| {
|
||||
runtime_id.strip_prefix("agent-runtime:").and_then(|value| value.rsplit_once(':').map(|(key, _)| key))
|
||||
})
|
||||
}
|
||||
|
||||
async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
|
||||
@@ -97,6 +114,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
|
||||
can_stop: false,
|
||||
can_restart: state.agent_manager.is_driver_installed(key),
|
||||
control_unavailable_reason: None,
|
||||
protocol_mode: None,
|
||||
active_sessions: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -120,6 +139,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
|
||||
can_stop: true,
|
||||
can_restart: true,
|
||||
control_unavailable_reason: None,
|
||||
protocol_mode: Some("legacy".to_string()),
|
||||
active_sessions: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -140,21 +161,31 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
|
||||
.unwrap_or("agent");
|
||||
let client = client.lock().await;
|
||||
let last_error = non_empty(client.stderr_tail_snapshot());
|
||||
let protocol_mode = client.protocol_mode();
|
||||
let runtime_id = if protocol_mode == "multi_session" {
|
||||
format!("agent-runtime:{key}:{}", client.pid())
|
||||
} else {
|
||||
format!("agent-connection:{pool_key}")
|
||||
};
|
||||
seeds.insert(
|
||||
format!("agent-connection:{pool_key}"),
|
||||
runtime_id.clone(),
|
||||
RuntimeSeed {
|
||||
id: format!("agent-connection:{pool_key}"),
|
||||
id: runtime_id,
|
||||
driver_key: key.to_string(),
|
||||
label: agent_catalog::label_for_key(key).unwrap_or(key).to_string(),
|
||||
kind: "agent".to_string(),
|
||||
source: "connection".to_string(),
|
||||
source: if protocol_mode == "multi_session" { "shared-runtime" } else { "connection" }
|
||||
.to_string(),
|
||||
status: "running".to_string(),
|
||||
pid: Some(client.pid()),
|
||||
version: agent_state.installed_drivers.get(key).map(|driver| driver.version.clone()),
|
||||
last_error,
|
||||
can_stop: false,
|
||||
can_restart: false,
|
||||
control_unavailable_reason: Some("connection-owned".to_string()),
|
||||
can_stop: protocol_mode == "multi_session",
|
||||
can_restart: protocol_mode == "multi_session",
|
||||
control_unavailable_reason: (protocol_mode != "multi_session")
|
||||
.then(|| "connection-owned".to_string()),
|
||||
protocol_mode: Some(protocol_mode.to_string()),
|
||||
active_sessions: Some(client.active_session_count()),
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -174,6 +205,8 @@ async fn collect_runtime_seeds(state: &AppState) -> Vec<RuntimeSeed> {
|
||||
can_stop: false,
|
||||
can_restart: false,
|
||||
control_unavailable_reason: Some("connection-owned".to_string()),
|
||||
protocol_mode: None,
|
||||
active_sessions: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
@@ -248,6 +281,8 @@ fn build_summary(seeds: Vec<RuntimeSeed>, stats: &HashMap<u32, ProcessStats>) ->
|
||||
can_stop: seed.can_stop,
|
||||
can_restart: seed.can_restart,
|
||||
control_unavailable_reason: seed.control_unavailable_reason,
|
||||
protocol_mode: seed.protocol_mode,
|
||||
active_sessions: seed.active_sessions,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
@@ -274,7 +309,7 @@ fn non_empty(value: String) -> Option<String> {
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{build_summary, ProcessStats, RuntimeSeed};
|
||||
use super::{build_summary, runtime_agent_key, ProcessStats, RuntimeSeed};
|
||||
use std::collections::HashMap;
|
||||
|
||||
fn seed(id: &str, status: &str, pid: Option<u32>) -> RuntimeSeed {
|
||||
@@ -291,6 +326,8 @@ mod tests {
|
||||
can_stop: true,
|
||||
can_restart: true,
|
||||
control_unavailable_reason: None,
|
||||
protocol_mode: None,
|
||||
active_sessions: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -306,4 +343,11 @@ mod tests {
|
||||
assert_eq!(summary.health, "healthy");
|
||||
assert_eq!(summary.runtimes[0].memory_bytes, Some(1024));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_legacy_and_shared_agent_runtime_ids() {
|
||||
assert_eq!(runtime_agent_key("agent:oracle"), Some("oracle"));
|
||||
assert_eq!(runtime_agent_key("agent-runtime:oracle:123"), Some("oracle"));
|
||||
assert_eq!(runtime_agent_key("plugin-connection:test"), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -31,7 +31,6 @@ const AGENT_RELEASE_PATHS = [
|
||||
"agents/common/build.gradle",
|
||||
"agents/common/src/main/",
|
||||
"agents/drivers/",
|
||||
"agents/opengauss/",
|
||||
];
|
||||
|
||||
const args = process.argv.slice(2);
|
||||
|
||||
Reference in New Issue
Block a user