fix(agent): recover JDBC connections after query timeouts
This commit is contained in:
miracle
GitHub
parent
0d484f5c74
commit
0e414e638e
|
|
@ -557,24 +557,34 @@ jobs:
|
|||
set -euo pipefail
|
||||
for version in 3.13 4.3; do
|
||||
name="dbx-rabbitmq-${version//./-}"
|
||||
cookie="dbx-ci-${GITHUB_RUN_ID:-local}-${version//./-}"
|
||||
# RabbitMQ is disposable in CI. Keep its data on a fresh tmpfs
|
||||
# owned by the image's rabbitmq user so .erlang.cookie is readable.
|
||||
docker rm -fv "$name" >/dev/null 2>&1 || true
|
||||
docker run -d --name "$name" \
|
||||
--user 999:999 \
|
||||
--tmpfs /var/lib/rabbitmq:rw,exec,uid=999,gid=999,mode=700 \
|
||||
-e RABBITMQ_DEFAULT_USER=dbx \
|
||||
-e RABBITMQ_DEFAULT_PASS=dbx-password \
|
||||
-e RABBITMQ_ERLANG_COOKIE="$cookie" \
|
||||
-p 5672:5672 -p 15672:15672 \
|
||||
"rabbitmq:${version}-management"
|
||||
cleanup() {
|
||||
docker rm -f "$name" >/dev/null 2>&1 || true
|
||||
docker rm -fv "$name" >/dev/null 2>&1 || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
ready=false
|
||||
for _ in $(seq 1 60); do
|
||||
if docker exec "$name" rabbitmq-diagnostics -q ping >/dev/null 2>&1; then
|
||||
if docker exec "$name" rabbitmq-diagnostics -q check_running >/dev/null 2>&1 \
|
||||
&& curl --fail --silent --noproxy '*' --user dbx:dbx-password \
|
||||
http://127.0.0.1:15672/api/overview >/dev/null; then
|
||||
ready=true
|
||||
break
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
if [ "$ready" != "true" ]; then
|
||||
docker inspect "$name" --format 'image={{.Config.Image}} user={{.Config.User}} status={{.State.Status}} exit={{.State.ExitCode}}' || true
|
||||
docker logs "$name"
|
||||
exit 1
|
||||
fi
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
private boolean requestActive;
|
||||
private boolean leasePinnedAtRequestStart;
|
||||
private boolean sessionAffinity;
|
||||
private boolean pooledConnectionPoisoned;
|
||||
|
||||
@Override
|
||||
public final Connection getConnection() {
|
||||
|
|
@ -47,6 +48,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
sessionAffinity = false;
|
||||
requestActive = false;
|
||||
leasePinnedAtRequestStart = false;
|
||||
pooledConnectionPoisoned = false;
|
||||
loadDriver(params);
|
||||
configuredDatabase = params.getDatabase();
|
||||
connectParams = params;
|
||||
|
|
@ -256,6 +258,7 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
sessionAffinity = false;
|
||||
requestActive = false;
|
||||
leasePinnedAtRequestStart = false;
|
||||
pooledConnectionPoisoned = false;
|
||||
identifierQuote = "";
|
||||
});
|
||||
}
|
||||
|
|
@ -271,24 +274,54 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
return poolRegistry != null;
|
||||
}
|
||||
|
||||
final synchronized void beginPooledRequest() throws Exception {
|
||||
if (poolRegistry == null) {
|
||||
return;
|
||||
}
|
||||
if (connectParams == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
requestActive = true;
|
||||
leasePinnedAtRequestStart = pooledLease != null;
|
||||
if (pooledLease == null) {
|
||||
try {
|
||||
pooledLease = borrowPooledConnection();
|
||||
} catch (Exception error) {
|
||||
requestActive = false;
|
||||
throw error;
|
||||
final synchronized boolean quarantinePooledConnection() {
|
||||
pooledConnectionPoisoned = true;
|
||||
return requestActive && pooledLease != null && pooledLease.quarantine();
|
||||
}
|
||||
|
||||
final void beginPooledRequest() throws Exception {
|
||||
ConnectParams params;
|
||||
String identity;
|
||||
JdbcConnectionPoolRegistry registry;
|
||||
synchronized (this) {
|
||||
if (poolRegistry == null) {
|
||||
return;
|
||||
}
|
||||
if (connectParams == null || poolIdentity == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
requestActive = true;
|
||||
leasePinnedAtRequestStart = pooledLease != null;
|
||||
if (pooledLease != null) {
|
||||
connection = pooledLease.connection();
|
||||
return;
|
||||
}
|
||||
params = connectParams;
|
||||
identity = poolIdentity;
|
||||
registry = poolRegistry;
|
||||
}
|
||||
|
||||
JdbcConnectionPoolRegistry.Lease borrowed;
|
||||
try {
|
||||
borrowed = borrowPooledConnection(registry, identity, params);
|
||||
} catch (Exception error) {
|
||||
synchronized (this) {
|
||||
requestActive = false;
|
||||
leasePinnedAtRequestStart = false;
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
|
||||
synchronized (this) {
|
||||
if (pooledConnectionPoisoned) {
|
||||
requestActive = false;
|
||||
leasePinnedAtRequestStart = false;
|
||||
borrowed.evict();
|
||||
throw new IllegalStateException("JDBC Session was quarantined while waiting for a connection");
|
||||
}
|
||||
pooledLease = borrowed;
|
||||
connection = borrowed.connection();
|
||||
}
|
||||
connection = pooledLease.connection();
|
||||
}
|
||||
|
||||
final synchronized void finishPooledRequest(
|
||||
|
|
@ -301,6 +334,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
return;
|
||||
}
|
||||
requestActive = false;
|
||||
if (pooledConnectionPoisoned) {
|
||||
releasePooledConnection(true);
|
||||
return;
|
||||
}
|
||||
if (succeeded && requiresSessionAffinity) {
|
||||
sessionAffinity = true;
|
||||
JdbcSchemaSwitcher.forget(connection);
|
||||
|
|
@ -322,7 +359,14 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
}
|
||||
|
||||
final synchronized void releaseIdlePooledConnection(JdbcExecutor executor) {
|
||||
if (poolRegistry == null || requestActive || sessionAffinity || pooledLease == null) {
|
||||
if (poolRegistry == null || requestActive || pooledLease == null) {
|
||||
return;
|
||||
}
|
||||
if (pooledConnectionPoisoned) {
|
||||
releasePooledConnection(true);
|
||||
return;
|
||||
}
|
||||
if (sessionAffinity) {
|
||||
return;
|
||||
}
|
||||
if (executor.hasOpenSessions() || executor.hasActiveStatements()) {
|
||||
|
|
@ -489,18 +533,20 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
|
||||
private Connection openInitializedConnection(ConnectParams params) throws Exception {
|
||||
Connection opened = openConnection(params);
|
||||
boolean initialized = false;
|
||||
try {
|
||||
afterPhysicalConnect(params, opened);
|
||||
initialized = true;
|
||||
return opened;
|
||||
} finally {
|
||||
if (!initialized) {
|
||||
try {
|
||||
opened.close();
|
||||
} catch (Exception ignored) {
|
||||
}
|
||||
} catch (Exception error) {
|
||||
try {
|
||||
opened.close();
|
||||
} catch (Exception closeError) {
|
||||
error.addSuppressed(closeError);
|
||||
throw AgentRpcError.resource(
|
||||
"close",
|
||||
new JdbcConnectionPoolRegistry.PhysicalConnectionStateUnknownException(error)
|
||||
);
|
||||
}
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -509,12 +555,24 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
|
|||
if (params == null || poolIdentity == null || poolRegistry == null) {
|
||||
throw new IllegalStateException("Not connected");
|
||||
}
|
||||
return poolRegistry.borrow(poolIdentity, () -> openInitializedConnection(params));
|
||||
return borrowPooledConnection(poolRegistry, poolIdentity, params);
|
||||
}
|
||||
|
||||
private JdbcConnectionPoolRegistry.Lease borrowPooledConnection(
|
||||
JdbcConnectionPoolRegistry registry,
|
||||
String identity,
|
||||
ConnectParams params
|
||||
) throws Exception {
|
||||
return registry.borrow(
|
||||
identity,
|
||||
JdbcSessionRole.from(params.getSessionRole()),
|
||||
() -> openInitializedConnection(params)
|
||||
);
|
||||
}
|
||||
|
||||
private void closeCurrentConnection() throws Exception {
|
||||
if (pooledLease != null) {
|
||||
boolean evict = sessionAffinity || !preparePooledConnectionForReturn();
|
||||
boolean evict = pooledConnectionPoisoned || sessionAffinity || !preparePooledConnectionForReturn();
|
||||
releasePooledConnection(evict);
|
||||
} else if (connection != null) {
|
||||
connection.close();
|
||||
|
|
|
|||
|
|
@ -0,0 +1,135 @@
|
|||
package com.dbx.agent;
|
||||
|
||||
import com.google.gson.JsonObject;
|
||||
|
||||
import java.sql.SQLException;
|
||||
import java.sql.SQLRecoverableException;
|
||||
import java.sql.SQLTransientConnectionException;
|
||||
import java.util.Locale;
|
||||
|
||||
final class AgentRpcError extends RuntimeException {
|
||||
private final String category;
|
||||
private final boolean retryable;
|
||||
private final String disposition;
|
||||
private final String stage;
|
||||
|
||||
private AgentRpcError(
|
||||
String message,
|
||||
String category,
|
||||
boolean retryable,
|
||||
String disposition,
|
||||
String stage,
|
||||
Throwable cause
|
||||
) {
|
||||
super(message, cause);
|
||||
this.category = category;
|
||||
this.retryable = retryable;
|
||||
this.disposition = disposition;
|
||||
this.stage = stage;
|
||||
}
|
||||
|
||||
static AgentRpcError resource(String stage, Throwable cause) {
|
||||
return new AgentRpcError(
|
||||
"Agent runtime resource limit reached",
|
||||
"resource",
|
||||
false,
|
||||
"replace_runtime",
|
||||
stage,
|
||||
cause
|
||||
);
|
||||
}
|
||||
|
||||
static AgentRpcError backpressure(String stage, Throwable cause) {
|
||||
return new AgentRpcError(
|
||||
"Agent request capacity is temporarily exhausted",
|
||||
"resource",
|
||||
true,
|
||||
"keep",
|
||||
stage,
|
||||
cause
|
||||
);
|
||||
}
|
||||
|
||||
static JsonObject toJson(Throwable error, String method, String agentSessionId) {
|
||||
AgentRpcError classified = classify(error, method);
|
||||
JsonObject rpcError = new JsonObject();
|
||||
rpcError.addProperty("code", -1);
|
||||
rpcError.addProperty("message", message(error));
|
||||
JsonObject data = new JsonObject();
|
||||
data.addProperty("category", classified.category);
|
||||
data.addProperty("retryable", classified.retryable);
|
||||
data.addProperty("sessionDisposition", classified.disposition);
|
||||
data.addProperty("stage", classified.stage);
|
||||
if (agentSessionId != null && !agentSessionId.trim().isEmpty()) {
|
||||
data.addProperty("agentSessionId", agentSessionId);
|
||||
}
|
||||
rpcError.add("data", data);
|
||||
return rpcError;
|
||||
}
|
||||
|
||||
private static AgentRpcError classify(Throwable error, String method) {
|
||||
AgentRpcError explicit = find(error, AgentRpcError.class);
|
||||
if (explicit != null) {
|
||||
return explicit;
|
||||
}
|
||||
SQLException sqlError = find(error, SQLException.class);
|
||||
if (sqlError != null) {
|
||||
String sqlState = sqlError.getSQLState();
|
||||
String stage = stage(method);
|
||||
boolean connectionError = "connect".equals(stage)
|
||||
|| "validate".equals(stage)
|
||||
|| sqlError instanceof SQLRecoverableException
|
||||
|| sqlError instanceof SQLTransientConnectionException
|
||||
|| (sqlState != null && sqlState.toUpperCase(Locale.ROOT).startsWith("08"));
|
||||
boolean operationRetryable = connectionError && ("connect".equals(stage) || "validate".equals(stage));
|
||||
String disposition = connectionError && !"connect".equals(stage) ? "quarantine" : "keep";
|
||||
return new AgentRpcError(
|
||||
message(error),
|
||||
connectionError ? "connection" : "sql",
|
||||
operationRetryable,
|
||||
disposition,
|
||||
stage,
|
||||
error
|
||||
);
|
||||
}
|
||||
return new AgentRpcError(message(error), "protocol", false, "keep", stage(method), error);
|
||||
}
|
||||
|
||||
private static String stage(String method) {
|
||||
if (method == null) {
|
||||
return "request";
|
||||
}
|
||||
if (AgentProtocol.METHOD_CONNECT.equals(method) || AgentProtocol.METHOD_OPEN_SESSION.equals(method)) {
|
||||
return "connect";
|
||||
}
|
||||
if (AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method) || AgentProtocol.METHOD_VALIDATE_SESSION.equals(method)) {
|
||||
return "validate";
|
||||
}
|
||||
if (AgentProtocol.METHOD_CANCEL_SESSION.equals(method)) {
|
||||
return "cancel";
|
||||
}
|
||||
if (AgentProtocol.METHOD_CLOSE_SESSION.equals(method) || AgentProtocol.METHOD_DISCONNECT.equals(method)) {
|
||||
return "close";
|
||||
}
|
||||
if (AgentProtocol.METHOD_FETCH_QUERY_PAGE.equals(method)
|
||||
|| AgentProtocol.METHOD_FETCH_TABLE_READ_PAGE.equals(method)) {
|
||||
return "fetch";
|
||||
}
|
||||
return "execute";
|
||||
}
|
||||
|
||||
private static String message(Throwable error) {
|
||||
return error.getMessage() == null ? error.toString() : error.getMessage();
|
||||
}
|
||||
|
||||
private static <T extends Throwable> T find(Throwable error, Class<T> type) {
|
||||
Throwable current = error;
|
||||
while (current != null) {
|
||||
if (type.isInstance(current)) {
|
||||
return type.cast(current);
|
||||
}
|
||||
current = current.getCause();
|
||||
}
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
|
@ -22,6 +22,7 @@ public final class ConnectParams {
|
|||
private String client_key_path;
|
||||
private String gbase_server;
|
||||
private String informix_server;
|
||||
private String sessionRole;
|
||||
|
||||
public ConnectParams() {
|
||||
this("", 0, "", "", "", "", "", false, "", Collections.emptyList());
|
||||
|
|
@ -200,6 +201,14 @@ public final class ConnectParams {
|
|||
this.informix_server = informix_server;
|
||||
}
|
||||
|
||||
public String getSessionRole() {
|
||||
return sessionRole;
|
||||
}
|
||||
|
||||
public void setSessionRole(String sessionRole) {
|
||||
this.sessionRole = sessionRole;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean equals(Object other) {
|
||||
if (this == other) return true;
|
||||
|
|
@ -221,13 +230,14 @@ public final class ConnectParams {
|
|||
&& Objects.equals(client_cert_path, that.client_cert_path)
|
||||
&& Objects.equals(client_key_path, that.client_key_path)
|
||||
&& Objects.equals(gbase_server, that.gbase_server)
|
||||
&& Objects.equals(informix_server, that.informix_server);
|
||||
&& Objects.equals(informix_server, that.informix_server)
|
||||
&& Objects.equals(sessionRole, that.sessionRole);
|
||||
}
|
||||
|
||||
@Override
|
||||
public int hashCode() {
|
||||
return Objects.hash(host, port, database, username, password, url_params, connection_string,
|
||||
port_explicit, mysql_compat_mode, jdbc_driver_class, jdbc_driver_paths, ssl, ca_cert_path, client_cert_path, client_key_path, gbase_server, informix_server);
|
||||
port_explicit, mysql_compat_mode, jdbc_driver_class, jdbc_driver_paths, ssl, ca_cert_path, client_cert_path, client_key_path, gbase_server, informix_server, sessionRole);
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -249,6 +259,7 @@ public final class ConnectParams {
|
|||
+ ", client_key_path=" + client_key_path
|
||||
+ ", gbase_server=" + gbase_server
|
||||
+ ", informix_server=" + informix_server
|
||||
+ ", sessionRole=" + sessionRole
|
||||
+ ")";
|
||||
}
|
||||
}
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,10 @@
|
|||
package com.dbx.agent;
|
||||
|
||||
enum JdbcSessionRole {
|
||||
WORKLOAD,
|
||||
METADATA;
|
||||
|
||||
static JdbcSessionRole from(String value) {
|
||||
return "metadata".equalsIgnoreCase(value == null ? "" : value.trim()) ? METADATA : WORKLOAD;
|
||||
}
|
||||
}
|
||||
|
|
@ -121,6 +121,11 @@ public final class JsonRpcServer {
|
|||
}
|
||||
}
|
||||
|
||||
boolean quarantine() {
|
||||
AbstractJdbcAgent jdbcAgent = pooledJdbcAgent();
|
||||
return jdbcAgent != null && jdbcAgent.quarantinePooledConnection();
|
||||
}
|
||||
|
||||
void expireIdleResources() {
|
||||
jdbcExecutor.expireIdleResources();
|
||||
releaseIdlePooledConnection();
|
||||
|
|
|
|||
|
|
@ -13,22 +13,30 @@ import java.util.Map;
|
|||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.RejectedExecutionException;
|
||||
import java.util.concurrent.ScheduledExecutorService;
|
||||
import java.util.concurrent.SynchronousQueue;
|
||||
import java.util.concurrent.ThreadPoolExecutor;
|
||||
import java.util.concurrent.ThreadFactory;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.concurrent.locks.ReentrantLock;
|
||||
import java.util.function.Consumer;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
||||
private static final String LEGACY_SESSION_ID = "__legacy__";
|
||||
private static final int MAX_SESSIONS = 256;
|
||||
private static final int MAX_REQUEST_THREADS = 64;
|
||||
private static final int MAX_CLEANUP_THREADS = 16;
|
||||
private static final long MAINTENANCE_INTERVAL_MILLIS = 60_000L;
|
||||
|
||||
private final Supplier<? extends DatabaseAgent> agentFactory;
|
||||
private final Supplier<? extends SessionRpcHandler> sessionHandlerFactory;
|
||||
private final Map<String, Session> sessions = new ConcurrentHashMap<>();
|
||||
private final ExecutorService requests = Executors.newCachedThreadPool();
|
||||
private final ExecutorService requests;
|
||||
private final ExecutorService cleanup;
|
||||
private final JdbcConnectionPoolRegistry poolRegistry;
|
||||
private final Gson gson = new Gson();
|
||||
private final PrintStream protocolOutput = System.out;
|
||||
|
|
@ -38,29 +46,43 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
private ScheduledExecutorService maintenance;
|
||||
|
||||
public MultiSessionJsonRpcServer(Supplier<? extends DatabaseAgent> agentFactory) {
|
||||
this(agentFactory, new JdbcConnectionPoolRegistry());
|
||||
this(agentFactory, new JdbcConnectionPoolRegistry(), RuntimeLimits.defaults());
|
||||
}
|
||||
|
||||
MultiSessionJsonRpcServer(
|
||||
Supplier<? extends DatabaseAgent> agentFactory,
|
||||
JdbcConnectionPoolRegistry.PoolSettings poolSettings
|
||||
) {
|
||||
this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings));
|
||||
this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings), RuntimeLimits.defaults());
|
||||
}
|
||||
|
||||
MultiSessionJsonRpcServer(
|
||||
Supplier<? extends DatabaseAgent> agentFactory,
|
||||
JdbcConnectionPoolRegistry.PoolSettings poolSettings,
|
||||
RuntimeLimits runtimeLimits
|
||||
) {
|
||||
this(agentFactory, new JdbcConnectionPoolRegistry(poolSettings), runtimeLimits);
|
||||
}
|
||||
|
||||
private MultiSessionJsonRpcServer(
|
||||
Supplier<? extends DatabaseAgent> agentFactory,
|
||||
JdbcConnectionPoolRegistry poolRegistry
|
||||
JdbcConnectionPoolRegistry poolRegistry,
|
||||
RuntimeLimits runtimeLimits
|
||||
) {
|
||||
this.agentFactory = agentFactory;
|
||||
this.sessionHandlerFactory = null;
|
||||
this.poolRegistry = poolRegistry;
|
||||
this.requests = boundedExecutor(runtimeLimits.maximumRequestThreads, "dbx-agent-request");
|
||||
this.cleanup = boundedExecutor(runtimeLimits.maximumCleanupThreads, "dbx-agent-cleanup");
|
||||
}
|
||||
|
||||
private MultiSessionJsonRpcServer(Supplier<? extends SessionRpcHandler> sessionHandlerFactory, boolean customHandler) {
|
||||
this.agentFactory = null;
|
||||
this.sessionHandlerFactory = sessionHandlerFactory;
|
||||
this.poolRegistry = new JdbcConnectionPoolRegistry();
|
||||
RuntimeLimits runtimeLimits = RuntimeLimits.defaults();
|
||||
this.requests = boundedExecutor(runtimeLimits.maximumRequestThreads, "dbx-agent-request");
|
||||
this.cleanup = boundedExecutor(runtimeLimits.maximumCleanupThreads, "dbx-agent-cleanup");
|
||||
}
|
||||
|
||||
/** Creates a protocol v2 server for a non-JDBC, session-scoped agent. */
|
||||
|
|
@ -82,7 +104,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
writeResponse(handleRequest(request));
|
||||
return;
|
||||
}
|
||||
requests.submit(() -> writeResponse(handleRequest(request)));
|
||||
executeRequest(request, this::writeResponse);
|
||||
}
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
|
|
@ -95,6 +117,14 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
return gson.toJson(handleRequest(JsonParser.parseString(line).getAsJsonObject()));
|
||||
}
|
||||
|
||||
void executeRequest(JsonObject request, Consumer<JsonObject> responseConsumer) {
|
||||
try {
|
||||
requests.execute(() -> responseConsumer.accept(handleRequest(request)));
|
||||
} catch (RejectedExecutionException error) {
|
||||
responseConsumer.accept(errorResponse(request.get("id"), AgentRpcError.backpressure("request", error)));
|
||||
}
|
||||
}
|
||||
|
||||
private JsonObject handleRequest(JsonObject request) {
|
||||
JsonElement id = request.get("id");
|
||||
String method = request.get("method").getAsString();
|
||||
|
|
@ -134,10 +164,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
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);
|
||||
response.add("error", AgentRpcError.toJson(error, method, stringOrNull(params, "agentSessionId")));
|
||||
}
|
||||
return response;
|
||||
}
|
||||
|
|
@ -165,7 +192,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
return session.connect(params);
|
||||
} catch (Exception error) {
|
||||
sessions.remove(sessionId, session);
|
||||
session.close();
|
||||
session.quarantineAndClose(cleanup);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
|
@ -173,7 +200,13 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
private Object closeSession(String sessionId) {
|
||||
Session session = sessions.remove(sessionId);
|
||||
if (session != null) {
|
||||
session.close();
|
||||
boolean replaceRuntime = session.quarantineAndClose(cleanup);
|
||||
if (replaceRuntime) {
|
||||
throw AgentRpcError.resource(
|
||||
"close",
|
||||
new IllegalStateException("JDBC quarantine operation limit reached")
|
||||
);
|
||||
}
|
||||
}
|
||||
return Collections.singletonMap("ok", true);
|
||||
}
|
||||
|
|
@ -206,7 +239,11 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
|
||||
private void closeAllSessions() {
|
||||
for (String sessionId : sessions.keySet()) {
|
||||
closeSession(sessionId);
|
||||
try {
|
||||
closeSession(sessionId);
|
||||
} catch (AgentRpcError ignored) {
|
||||
// Sessions are already detached; process shutdown remains the final cleanup boundary.
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -263,7 +300,16 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
}
|
||||
closeAllSessions();
|
||||
requests.shutdown();
|
||||
requests.shutdownNow();
|
||||
cleanup.shutdown();
|
||||
try {
|
||||
if (!cleanup.awaitTermination(2, TimeUnit.SECONDS)) {
|
||||
cleanup.shutdownNow();
|
||||
}
|
||||
} catch (InterruptedException error) {
|
||||
Thread.currentThread().interrupt();
|
||||
cleanup.shutdownNow();
|
||||
}
|
||||
poolRegistry.close();
|
||||
}
|
||||
|
||||
|
|
@ -275,6 +321,47 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
};
|
||||
}
|
||||
|
||||
private static ExecutorService boundedExecutor(int maximumThreads, String threadName) {
|
||||
return new ThreadPoolExecutor(
|
||||
0,
|
||||
maximumThreads,
|
||||
60L,
|
||||
TimeUnit.SECONDS,
|
||||
new SynchronousQueue<>(),
|
||||
daemonThreadFactory(threadName),
|
||||
new ThreadPoolExecutor.AbortPolicy()
|
||||
);
|
||||
}
|
||||
|
||||
private static JsonObject errorResponse(JsonElement id, Throwable error) {
|
||||
JsonObject response = new JsonObject();
|
||||
response.addProperty("jsonrpc", "2.0");
|
||||
response.add("id", id);
|
||||
response.add("error", AgentRpcError.toJson(error, "request", null));
|
||||
return response;
|
||||
}
|
||||
|
||||
private static String stringOrNull(JsonObject params, String key) {
|
||||
return params.has(key) && !params.get(key).isJsonNull() ? params.get(key).getAsString() : null;
|
||||
}
|
||||
|
||||
static final class RuntimeLimits {
|
||||
private final int maximumRequestThreads;
|
||||
private final int maximumCleanupThreads;
|
||||
|
||||
RuntimeLimits(int maximumRequestThreads, int maximumCleanupThreads) {
|
||||
if (maximumRequestThreads <= 0 || maximumCleanupThreads <= 0) {
|
||||
throw new IllegalArgumentException("Agent runtime thread limits must be positive");
|
||||
}
|
||||
this.maximumRequestThreads = maximumRequestThreads;
|
||||
this.maximumCleanupThreads = maximumCleanupThreads;
|
||||
}
|
||||
|
||||
private static RuntimeLimits defaults() {
|
||||
return new RuntimeLimits(MAX_REQUEST_THREADS, MAX_CLEANUP_THREADS);
|
||||
}
|
||||
}
|
||||
|
||||
private static String requiredSessionId(JsonObject params) {
|
||||
if (!params.has("agentSessionId") || params.get("agentSessionId").getAsString().trim().isEmpty()) {
|
||||
throw new IllegalArgumentException("agentSessionId is required");
|
||||
|
|
@ -293,6 +380,8 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
private final JsonRpcServer server;
|
||||
private final SessionRpcHandler handler;
|
||||
private final ReentrantLock lock = new ReentrantLock();
|
||||
private final AtomicReference<State> state = new AtomicReference<>(State.ACTIVE);
|
||||
private final AtomicBoolean cleanupScheduled = new AtomicBoolean();
|
||||
|
||||
private Session(JsonRpcServer server) {
|
||||
this.server = server;
|
||||
|
|
@ -305,8 +394,10 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
|
||||
private Object handle(String method, JsonObject params) throws Exception {
|
||||
requireActive();
|
||||
lock.lock();
|
||||
try {
|
||||
requireActive();
|
||||
return handler == null ? server.dispatchForRuntime(method, params) : handler.handle(method, params);
|
||||
} finally {
|
||||
lock.unlock();
|
||||
|
|
@ -314,8 +405,10 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
|
||||
private Object connect(JsonObject params) throws Exception {
|
||||
requireActive();
|
||||
lock.lock();
|
||||
try {
|
||||
requireActive();
|
||||
return handler == null
|
||||
? server.dispatchForRuntime(AgentProtocol.METHOD_CONNECT, params)
|
||||
: handler.connect(params);
|
||||
|
|
@ -324,9 +417,31 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
}
|
||||
|
||||
private boolean quarantineAndClose(ExecutorService cleanup) {
|
||||
state.compareAndSet(State.ACTIVE, State.QUARANTINED);
|
||||
boolean replaceRuntime = server != null && server.quarantine();
|
||||
if (!cleanupScheduled.compareAndSet(false, true)) {
|
||||
return replaceRuntime;
|
||||
}
|
||||
try {
|
||||
cleanup.execute(this::closeWhenIdle);
|
||||
} catch (RejectedExecutionException error) {
|
||||
throw AgentRpcError.resource("close", error);
|
||||
}
|
||||
return replaceRuntime;
|
||||
}
|
||||
|
||||
private void closeWhenIdle() {
|
||||
close();
|
||||
}
|
||||
|
||||
private void close() {
|
||||
state.compareAndSet(State.ACTIVE, State.QUARANTINED);
|
||||
lock.lock();
|
||||
try {
|
||||
if (state.get() == State.CLOSED) {
|
||||
return;
|
||||
}
|
||||
if (handler != null) {
|
||||
handler.close();
|
||||
} else {
|
||||
|
|
@ -334,6 +449,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
} catch (Exception ignored) {
|
||||
} finally {
|
||||
state.set(State.CLOSED);
|
||||
lock.unlock();
|
||||
}
|
||||
}
|
||||
|
|
@ -347,7 +463,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
|
||||
private void expireIdleResources() {
|
||||
if (handler != null) {
|
||||
if (state.get() != State.ACTIVE || handler != null) {
|
||||
return;
|
||||
}
|
||||
if (!lock.tryLock()) {
|
||||
|
|
@ -361,7 +477,7 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
}
|
||||
|
||||
private void expireIdleResources(long nowMillis, long idleTimeoutMillis) {
|
||||
if (handler != null) {
|
||||
if (state.get() != State.ACTIVE || handler != null) {
|
||||
return;
|
||||
}
|
||||
if (!lock.tryLock()) {
|
||||
|
|
@ -373,5 +489,17 @@ public final class MultiSessionJsonRpcServer implements AutoCloseable {
|
|||
lock.unlock();
|
||||
}
|
||||
}
|
||||
|
||||
private void requireActive() {
|
||||
if (state.get() != State.ACTIVE) {
|
||||
throw new IllegalStateException("Agent session is quarantined");
|
||||
}
|
||||
}
|
||||
|
||||
private enum State {
|
||||
ACTIVE,
|
||||
QUARANTINED,
|
||||
CLOSED
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -151,7 +151,7 @@ class CommonJavaCompatibilityTest {
|
|||
}
|
||||
|
||||
@Test
|
||||
void multiSessionServerCreatesAndClosesIndependentAgents() {
|
||||
void multiSessionServerCreatesAndClosesIndependentAgents() throws Exception {
|
||||
java.util.List<TrackingAgent> created = new java.util.ArrayList<>();
|
||||
MultiSessionJsonRpcServer server = new MultiSessionJsonRpcServer(() -> {
|
||||
TrackingAgent agent = new TrackingAgent();
|
||||
|
|
@ -172,6 +172,7 @@ class CommonJavaCompatibilityTest {
|
|||
assertEquals(1, created.get(1).connectCount);
|
||||
|
||||
server.handleRequest("{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"close_session\",\"params\":{\"agentSessionId\":\"a\"}}");
|
||||
awaitCondition(() -> created.get(0).disconnectCount == 1);
|
||||
assertEquals(1, created.get(0).disconnectCount);
|
||||
assertEquals(0, created.get(1).disconnectCount);
|
||||
}
|
||||
|
|
@ -691,7 +692,7 @@ class CommonJavaCompatibilityTest {
|
|||
|
||||
private static final class TrackingAgent extends MinimalAgent {
|
||||
private int connectCount;
|
||||
private int disconnectCount;
|
||||
private volatile int disconnectCount;
|
||||
|
||||
@Override
|
||||
public void connect(ConnectParams params) {
|
||||
|
|
@ -1078,6 +1079,14 @@ class CommonJavaCompatibilityTest {
|
|||
return false;
|
||||
}
|
||||
|
||||
private static void awaitCondition(java.util.function.BooleanSupplier condition) throws InterruptedException {
|
||||
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
|
||||
while (!condition.getAsBoolean() && System.nanoTime() < deadline) {
|
||||
Thread.sleep(10L);
|
||||
}
|
||||
assertTrue(condition.getAsBoolean());
|
||||
}
|
||||
|
||||
private static JsonObject protocolContract(String resourcePath) {
|
||||
InputStream stream = CommonJavaCompatibilityTest.class.getResourceAsStream(resourcePath);
|
||||
if (stream == null) {
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -4,7 +4,7 @@ Protocol v2 allows one Agent process to serve multiple isolated database session
|
|||
|
||||
## Session lifecycle
|
||||
|
||||
- `open_session` creates one logical database session. Parameters contain the normal connection fields plus `agentSessionId`.
|
||||
- `open_session` creates one logical database session. Parameters contain the normal connection fields plus `agentSessionId` and an optional `sessionRole`.
|
||||
- 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.
|
||||
|
|
@ -13,6 +13,8 @@ Protocol v2 allows one Agent process to serve multiple isolated database session
|
|||
|
||||
`agentSessionId` identifies a logical database connection. Existing `sessionId` fields remain pagination cursor identifiers and must not be used as logical connection identifiers.
|
||||
|
||||
`sessionRole` is `workload` by default. DBX sends `metadata` for object-tree, completion, and other read-only metadata sessions. New runtimes use this role to preserve metadata checkout capacity; older runtimes may ignore the field.
|
||||
|
||||
## 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`.
|
||||
|
|
@ -31,6 +33,22 @@ Etcd and ZooKeeper retain the legacy path because they use the key-value Agent p
|
|||
|
||||
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.
|
||||
|
||||
JSON-RPC failures may include structured recovery data:
|
||||
|
||||
```json
|
||||
{
|
||||
"category": "timeout|canceled|connection|protocol|resource|sql",
|
||||
"retryable": false,
|
||||
"sessionDisposition": "keep|quarantine|replace_runtime",
|
||||
"agentSessionId": "optional-session-id",
|
||||
"stage": "checkout|connect|validate|execute|fetch|cancel|close"
|
||||
}
|
||||
```
|
||||
|
||||
`keep` preserves the logical session, `quarantine` removes only that session from routing, and `replace_runtime` requires DBX to atomically remove every pool sharing the runtime before terminating it. Agent code reports the disposition but must not independently terminate a shared runtime because it does not own DBX routing state. Temporary workload checkout backpressure uses `category=resource`, `retryable=true`, and `sessionDisposition=keep`; only unrecoverable runtime or cleanup saturation requests `replace_runtime`.
|
||||
|
||||
The complete JDBC pool checkout runs under a bounded runtime executor, including HikariCP idle-connection validation, physical connection creation, and driver setup. Workload admission, the runtime-wide physical connection budget, physical creation, and checkout consume one absolute deadline rather than restarting the timeout at each stage. Connection return, eviction, and physical close use separate bounded executors so they cannot deadlock checkout or creation. If a driver call outlives its boundary, or cleanup cannot confirm the physical connection state, the connection identity is poisoned and returns `category=resource` with `sessionDisposition=replace_runtime` on the current or next checkout. A late connection must be evicted and closed instead of published, and DBX must not replay the timed-out user operation automatically.
|
||||
|
||||
## Driver author guidance
|
||||
|
||||
Use `MultiSessionJsonRpcServer(YourAgent::new)` for Java SQL Agents so each logical session receives a new `DatabaseAgent` with isolated connection state. The shared runtime owns the physical JDBC pools. 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.
|
||||
|
|
|
|||
|
|
@ -103,12 +103,20 @@ func TestRabbitMQIntegration(t *testing.T) {
|
|||
t.Fatalf("unexpected messages %#v", messages)
|
||||
}
|
||||
|
||||
stats, err := service.getTopicStats(jsonObject{"topic": queue, "virtual_host": vhost})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stats.(jsonObject)["totalMessages"] != int64(1) {
|
||||
t.Fatalf("unexpected stats %#v", stats)
|
||||
var stats any
|
||||
statsDeadline := time.Now().Add(10 * time.Second)
|
||||
for {
|
||||
stats, err = service.getTopicStats(jsonObject{"topic": queue, "virtual_host": vhost})
|
||||
if err == nil && stats.(jsonObject)["totalMessages"] == int64(1) {
|
||||
break
|
||||
}
|
||||
if time.Now().After(statsDeadline) {
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Fatalf("unexpected stats %#v", stats)
|
||||
}
|
||||
time.Sleep(250 * time.Millisecond)
|
||||
}
|
||||
config, err := service.getTopicConfig(jsonObject{"topic": queue, "virtual_host": vhost})
|
||||
if err != nil {
|
||||
|
|
|
|||
|
|
@ -721,6 +721,261 @@ describe("connectionStore metadata loading", () => {
|
|||
expect(store.treeNodes[0]?.children?.[0]?.children?.map((node) => node.label)).toEqual(["public", "tree.extensions"]);
|
||||
});
|
||||
|
||||
it("preserves the last successful tree snapshot when a forced metadata refresh fails", async () => {
|
||||
const listSchemaInfos = vi.fn().mockRejectedValue(new Error("Agent RPC call timed out (5s)"));
|
||||
const deleteSchemaCachePrefix = vi.fn().mockResolvedValue(undefined);
|
||||
|
||||
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
|
||||
vi.doMock("@/lib/backend/api", () => ({
|
||||
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
|
||||
deleteSchemaCachePrefix,
|
||||
listInstalledAgents: vi.fn().mockResolvedValue([]),
|
||||
listSchemaInfos,
|
||||
loadSchemaCache: vi.fn().mockResolvedValue(null),
|
||||
saveConnections: vi.fn().mockResolvedValue(undefined),
|
||||
saveSchemaCache: vi.fn().mockResolvedValue(undefined),
|
||||
saveSidebarLayout: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const { useConnectionStore } = await import("@/stores/connectionStore");
|
||||
const store = useConnectionStore();
|
||||
const connection = postgresConnection();
|
||||
const previousSchema: TreeNode = {
|
||||
id: `${connection.id}:app:public`,
|
||||
label: "public",
|
||||
type: "schema",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
schema: "public",
|
||||
isExpanded: false,
|
||||
children: [],
|
||||
};
|
||||
const databaseNode: TreeNode = {
|
||||
id: `${connection.id}:app`,
|
||||
label: "app",
|
||||
type: "database",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
isExpanded: true,
|
||||
children: [previousSchema],
|
||||
};
|
||||
store.connections = [connection];
|
||||
store.connectedIds.add(connection.id);
|
||||
store.treeNodes = [
|
||||
{
|
||||
id: connection.id,
|
||||
label: connection.name,
|
||||
type: "connection",
|
||||
connectionId: connection.id,
|
||||
isExpanded: true,
|
||||
children: [databaseNode],
|
||||
},
|
||||
];
|
||||
|
||||
await expect(store.refreshTreeNode(databaseNode)).rejects.toThrow("Agent RPC call timed out (5s)");
|
||||
|
||||
expect(databaseNode.children).toEqual([previousSchema]);
|
||||
expect(databaseNode.isExpanded).toBe(true);
|
||||
expect(store.connectionErrors[connection.id]).toBe("Agent RPC call timed out (5s)");
|
||||
expect(deleteSchemaCachePrefix).toHaveBeenCalledWith("pg-1:app:");
|
||||
});
|
||||
|
||||
it("does not let an older refresh resume after a newer refresh succeeds", async () => {
|
||||
let resolveOlderMetadata!: (value: Array<{ name: string; comment: null }>) => void;
|
||||
const olderMetadata = new Promise<Array<{ name: string; comment: null }>>((resolve) => {
|
||||
resolveOlderMetadata = resolve;
|
||||
});
|
||||
const deleteSchemaCachePrefix = vi.fn().mockResolvedValue(undefined);
|
||||
const listSchemaInfos = vi
|
||||
.fn()
|
||||
.mockImplementationOnce(() => olderMetadata)
|
||||
.mockResolvedValue([{ name: "latest", comment: null }]);
|
||||
|
||||
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
|
||||
vi.doMock("@/lib/backend/api", () => ({
|
||||
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
|
||||
deleteSchemaCachePrefix,
|
||||
listInstalledAgents: vi.fn().mockResolvedValue([]),
|
||||
listSchemaInfos,
|
||||
loadSchemaCache: vi.fn().mockResolvedValue(null),
|
||||
saveConnections: vi.fn().mockResolvedValue(undefined),
|
||||
saveSchemaCache: vi.fn().mockResolvedValue(undefined),
|
||||
saveSidebarLayout: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const { useConnectionStore } = await import("@/stores/connectionStore");
|
||||
const store = useConnectionStore();
|
||||
const connection = postgresConnection();
|
||||
const databaseNode: TreeNode = {
|
||||
id: `${connection.id}:app`,
|
||||
label: "app",
|
||||
type: "database",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
isExpanded: true,
|
||||
children: [],
|
||||
};
|
||||
store.connections = [connection];
|
||||
store.connectedIds.add(connection.id);
|
||||
store.treeNodes = [
|
||||
{
|
||||
id: connection.id,
|
||||
label: connection.name,
|
||||
type: "connection",
|
||||
connectionId: connection.id,
|
||||
isExpanded: true,
|
||||
children: [databaseNode],
|
||||
},
|
||||
];
|
||||
|
||||
const olderRefresh = store.refreshTreeNode(databaseNode);
|
||||
await vi.waitFor(() => expect(listSchemaInfos).toHaveBeenCalledTimes(1));
|
||||
await store.refreshTreeNode(databaseNode);
|
||||
resolveOlderMetadata([{ name: "stale", comment: null }]);
|
||||
await olderRefresh;
|
||||
|
||||
expect(listSchemaInfos).toHaveBeenCalledTimes(2);
|
||||
expect(databaseNode.children?.map((node) => node.label)).toEqual(["latest", "tree.extensions"]);
|
||||
});
|
||||
|
||||
it("does not let an older refresh failure overwrite a newer successful refresh", async () => {
|
||||
let rejectOlderMetadata!: (reason: Error) => void;
|
||||
const olderMetadata = new Promise<Array<{ name: string; comment: null }>>((_, reject) => {
|
||||
rejectOlderMetadata = reject;
|
||||
});
|
||||
const listSchemaInfos = vi
|
||||
.fn()
|
||||
.mockImplementationOnce(() => olderMetadata)
|
||||
.mockResolvedValue([{ name: "latest", comment: null }]);
|
||||
|
||||
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
|
||||
vi.doMock("@/lib/backend/api", () => ({
|
||||
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
|
||||
deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined),
|
||||
listInstalledAgents: vi.fn().mockResolvedValue([]),
|
||||
listSchemaInfos,
|
||||
loadSchemaCache: vi.fn().mockResolvedValue(null),
|
||||
saveConnections: vi.fn().mockResolvedValue(undefined),
|
||||
saveSchemaCache: vi.fn().mockResolvedValue(undefined),
|
||||
saveSidebarLayout: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const { useConnectionStore } = await import("@/stores/connectionStore");
|
||||
const store = useConnectionStore();
|
||||
const connection = postgresConnection();
|
||||
const databaseNode: TreeNode = {
|
||||
id: `${connection.id}:app`,
|
||||
label: "app",
|
||||
type: "database",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
isExpanded: true,
|
||||
children: [],
|
||||
};
|
||||
store.connections = [connection];
|
||||
store.connectedIds.add(connection.id);
|
||||
store.treeNodes = [
|
||||
{
|
||||
id: connection.id,
|
||||
label: connection.name,
|
||||
type: "connection",
|
||||
connectionId: connection.id,
|
||||
isExpanded: true,
|
||||
children: [databaseNode],
|
||||
},
|
||||
];
|
||||
|
||||
const olderRefresh = store.refreshTreeNode(databaseNode);
|
||||
await vi.waitFor(() => expect(listSchemaInfos).toHaveBeenCalledTimes(1));
|
||||
await store.refreshTreeNode(databaseNode);
|
||||
rejectOlderMetadata(new Error("connection closed"));
|
||||
await expect(olderRefresh).rejects.toThrow("connection closed");
|
||||
|
||||
expect(databaseNode.children?.map((node) => node.label)).toEqual(["latest", "tree.extensions"]);
|
||||
expect(store.connectionErrors[connection.id]).toBeUndefined();
|
||||
expect(store.connectedIds.has(connection.id)).toBe(true);
|
||||
});
|
||||
|
||||
it("does not restore a pre-disconnect snapshot into a same-id reconnected node", async () => {
|
||||
let rejectMetadata!: (reason: Error) => void;
|
||||
const pendingMetadata = new Promise<Array<{ name: string; comment: null }>>((_, reject) => {
|
||||
rejectMetadata = reject;
|
||||
});
|
||||
const listSchemaInfos = vi.fn().mockReturnValue(pendingMetadata);
|
||||
|
||||
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
|
||||
vi.doMock("@/lib/backend/api", () => ({
|
||||
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
|
||||
deleteSchemaCachePrefix: vi.fn().mockResolvedValue(undefined),
|
||||
disconnectDb: vi.fn().mockResolvedValue(undefined),
|
||||
listInstalledAgents: vi.fn().mockResolvedValue([]),
|
||||
listSchemaInfos,
|
||||
loadSchemaCache: vi.fn().mockResolvedValue(null),
|
||||
saveConnections: vi.fn().mockResolvedValue(undefined),
|
||||
saveSchemaCache: vi.fn().mockResolvedValue(undefined),
|
||||
saveSidebarLayout: vi.fn().mockResolvedValue(undefined),
|
||||
}));
|
||||
|
||||
const { useConnectionStore } = await import("@/stores/connectionStore");
|
||||
const store = useConnectionStore();
|
||||
const connection = postgresConnection();
|
||||
const databaseId = `${connection.id}:app`;
|
||||
const staleSchema: TreeNode = {
|
||||
id: `${databaseId}:stale`,
|
||||
label: "stale",
|
||||
type: "schema",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
schema: "stale",
|
||||
children: [],
|
||||
};
|
||||
const databaseNode: TreeNode = {
|
||||
id: databaseId,
|
||||
label: "app",
|
||||
type: "database",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
isExpanded: true,
|
||||
children: [staleSchema],
|
||||
};
|
||||
const connectionNode: TreeNode = {
|
||||
id: connection.id,
|
||||
label: connection.name,
|
||||
type: "connection",
|
||||
connectionId: connection.id,
|
||||
isExpanded: true,
|
||||
children: [databaseNode],
|
||||
};
|
||||
store.connections = [connection];
|
||||
store.connectedIds.add(connection.id);
|
||||
store.treeNodes = [connectionNode];
|
||||
|
||||
const staleRefresh = store.refreshTreeNode(databaseNode);
|
||||
await vi.waitFor(() => expect(listSchemaInfos).toHaveBeenCalledTimes(1));
|
||||
await store.disconnect(connection.id);
|
||||
|
||||
const freshSchema: TreeNode = {
|
||||
id: `${databaseId}:fresh`,
|
||||
label: "fresh",
|
||||
type: "schema",
|
||||
connectionId: connection.id,
|
||||
database: "app",
|
||||
schema: "fresh",
|
||||
children: [],
|
||||
};
|
||||
const reconnectedDatabaseNode: TreeNode = {
|
||||
...databaseNode,
|
||||
children: [freshSchema],
|
||||
};
|
||||
connectionNode.children = [reconnectedDatabaseNode];
|
||||
store.connectedIds.add(connection.id);
|
||||
|
||||
rejectMetadata(new Error("disconnected refresh"));
|
||||
await expect(staleRefresh).rejects.toThrow("disconnected refresh");
|
||||
|
||||
expect(reconnectedDatabaseNode.children?.map((child) => child.label)).toEqual(["fresh"]);
|
||||
});
|
||||
|
||||
it.each(["opengauss", "kingbase"] as const)("reloads %s sidebar schemas when system visibility changes", async (dbType) => {
|
||||
const listSchemaInfos = vi.fn().mockResolvedValue([
|
||||
{ name: "information_schema", comment: null },
|
||||
|
|
|
|||
|
|
@ -419,6 +419,8 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
const connectionGroupPaths = computed(() => buildConnectionGroupPathMap(sidebarLayout.value));
|
||||
let layoutPersistTimer: ReturnType<typeof setTimeout> | null = null;
|
||||
const staleTreeRefreshIds = new Set<string>();
|
||||
const activeTreeRefreshGenerations = new Map<string, number>();
|
||||
let nextTreeRefreshGeneration = 0;
|
||||
const metadataLoadCoordinator = new MetadataLoadCoordinator((event) => {
|
||||
console.debug("[DBX][metadata-load:coordinator]", event);
|
||||
});
|
||||
|
|
@ -913,7 +915,8 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
}
|
||||
|
||||
// Metadata loaders keep this internal: match connection-loss errors before recording generic errors.
|
||||
function recordMetadataLoadError(connectionId: string, error: unknown) {
|
||||
function recordMetadataLoadError(connectionId: string, error: unknown, load?: TreeNodeLoadHandle) {
|
||||
if (load && !load.isCurrent()) return;
|
||||
if (recordConnectionLostError(connectionId, error)) return;
|
||||
recordConnectionError(connectionId, error);
|
||||
}
|
||||
|
|
@ -2976,7 +2979,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
if (cacheHit) {
|
||||
// Render the last known metadata immediately; network validation and refresh
|
||||
// continue in the background so opening a connection never waits on them.
|
||||
void loadDatabases(connectionId, { ...options, force: true }).catch((error) => recordMetadataLoadError(connectionId, error));
|
||||
void loadDatabases(connectionId, { ...options, force: true }).catch(() => undefined);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
|
@ -3053,7 +3056,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
let dorisCatalogs: CatalogInfo[] | null = null;
|
||||
if (connectionIsDorisFamilyCatalogCapable(config)) {
|
||||
dorisCatalogs = await withMetadataLoadTimeout(connectionId, api.listDorisCatalogs(connectionId), "catalogs").catch((error: unknown) => {
|
||||
recordMetadataLoadError(connectionId, error);
|
||||
recordMetadataLoadError(connectionId, error, load);
|
||||
return null;
|
||||
});
|
||||
}
|
||||
|
|
@ -3134,7 +3137,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
if (liveNode) liveNode.isExpanded = true;
|
||||
if (options?.force) void loadSidebarDatabaseStorage(connectionId, { force: true });
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3203,7 +3206,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3258,7 +3261,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3295,7 +3298,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3349,7 +3352,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
const liveNode = treeNodeLoadTarget(load);
|
||||
if (liveNode) liveNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3393,7 +3396,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3460,7 +3463,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3514,7 +3517,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3550,7 +3553,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3586,7 +3589,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
setChildren(targetNode, isMilvus && database ? collectionChildren : withSavedSqlRoot(connectionId, collectionChildren, targetNode));
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3631,7 +3634,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
setChildren(targetNode, children);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3708,7 +3711,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
await savePersistedTreeChildren(cacheKey, children);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3747,7 +3750,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
await savePersistedTreeChildren(cacheKey, children);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3783,7 +3786,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3821,7 +3824,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3866,7 +3869,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(node.connectionId, e);
|
||||
recordMetadataLoadError(node.connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -3922,7 +3925,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
setChildren(targetNode, databaseNodes);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4004,7 +4007,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
}
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4031,7 +4034,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
const nodeId = schema ? `${connectionId}:${database}:${schema}` : `${connectionId}:${database}`;
|
||||
const cacheKey = schemaCacheKey(connectionId, database, schema || "", "objects-simple-v6");
|
||||
if (await hydrateTreeNodeFromCache(findNode(treeNodes.value, nodeId), cacheKey)) {
|
||||
void loadTables(connectionId, database, schema, { ...options, force: true }).catch((error) => recordMetadataLoadError(connectionId, error));
|
||||
void loadTables(connectionId, database, schema, { ...options, force: true }).catch(() => undefined);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
|
@ -4136,7 +4139,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
});
|
||||
}
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4162,9 +4165,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
});
|
||||
if (!options?.force && !searchFilter && !options?.sidebarTableSearchParentId && !tableNameFilterForScope) {
|
||||
if (await hydrateTreeNodeFromCache(node, objectGroupCacheKey(node))) {
|
||||
void loadObjectGroupChildren(node, { ...options, force: true }).catch((error) => {
|
||||
if (node.connectionId) recordMetadataLoadError(node.connectionId, error);
|
||||
});
|
||||
void loadObjectGroupChildren(node, { ...options, force: true }).catch(() => undefined);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
|
@ -4259,7 +4260,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
}
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(node.connectionId, e);
|
||||
recordMetadataLoadError(node.connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4383,7 +4384,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
await savePersistedTreeChildren(objectGroupCacheKey(targetParent), nextChildren);
|
||||
targetParent.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(parentConnectionId, e);
|
||||
recordMetadataLoadError(parentConnectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4418,7 +4419,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
targetNode.objectCount = children.length;
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4553,7 +4554,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
const finishedParent = treeNodeLoadTarget(load);
|
||||
if (finishedParent) finishedParent.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(parent.connectionId, e);
|
||||
recordMetadataLoadError(parent.connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4795,7 +4796,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4836,7 +4837,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4881,7 +4882,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4923,7 +4924,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4954,7 +4955,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -4985,7 +4986,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -5016,7 +5017,7 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
);
|
||||
targetNode.isExpanded = true;
|
||||
} catch (e) {
|
||||
recordMetadataLoadError(connectionId, e);
|
||||
recordMetadataLoadError(connectionId, e, load);
|
||||
throw e;
|
||||
} finally {
|
||||
finishTreeNodeLoad(load);
|
||||
|
|
@ -5114,19 +5115,22 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
}
|
||||
}
|
||||
|
||||
async function restoreExpandedChildren(node: TreeNode, expandedIds: Set<string>, options?: LoadTreeOptions) {
|
||||
async function restoreExpandedChildren(node: TreeNode, expandedIds: Set<string>, options?: LoadTreeOptions, isCurrent: () => boolean = () => true) {
|
||||
if (!isCurrent()) return;
|
||||
if (!node.children) return;
|
||||
for (const child of node.children) {
|
||||
if (!isCurrent()) return;
|
||||
if (!expandedIds.has(child.id)) continue;
|
||||
await loadTreeNodeChildren(child, options);
|
||||
await restoreExpandedChildren(child, expandedIds, options);
|
||||
if (!isCurrent()) return;
|
||||
await restoreExpandedChildren(child, expandedIds, options, isCurrent);
|
||||
}
|
||||
}
|
||||
|
||||
async function refreshTreeNode(node: TreeNode) {
|
||||
invalidateMetadataCachesForNode(node);
|
||||
if (objectTypesForGroupNode(node.type)) {
|
||||
clearLoadedChildrenCache(node.id);
|
||||
clearLoadedChildrenCache(node.id, { deletePersisted: false });
|
||||
await loadObjectGroupChildren(node, { force: true });
|
||||
return;
|
||||
}
|
||||
|
|
@ -5144,24 +5148,44 @@ export const useConnectionStore = defineStore("connection", () => {
|
|||
const previousChildren = node.children;
|
||||
const previousHiddenChildren = node.hiddenChildren;
|
||||
const previousObjectCount = node.objectCount;
|
||||
const previousExpanded = node.isExpanded;
|
||||
const previousLoadedIds = [...loadedTreeNodeChildrenIds.value].filter((id) => id === node.id || id.startsWith(`${node.id}:`));
|
||||
const previousConfirmedEmptyIds = [...confirmedEmptyTreeNodeIds.value].filter((id) => id === node.id || id.startsWith(`${node.id}:`));
|
||||
await clearPersistedTreeCacheForNode(node);
|
||||
clearLoadedChildrenCache(node.id);
|
||||
if (node.type !== "connection-group") {
|
||||
node.children = [];
|
||||
}
|
||||
const connectionRevision = node.connectionId ? connectionStateRevision(node.connectionId) : undefined;
|
||||
const refreshGeneration = ++nextTreeRefreshGeneration;
|
||||
activeTreeRefreshGenerations.set(node.id, refreshGeneration);
|
||||
const ownsRefreshGeneration = () => activeTreeRefreshGenerations.get(node.id) === refreshGeneration;
|
||||
const isCurrentRefresh = () => ownsRefreshGeneration() && (!node.connectionId || connectionStateRevision(node.connectionId) === connectionRevision);
|
||||
try {
|
||||
await clearPersistedTreeCacheForNode(node);
|
||||
if (!isCurrentRefresh()) return;
|
||||
clearLoadedChildrenCache(node.id);
|
||||
if (node.type !== "connection-group") {
|
||||
node.children = [];
|
||||
}
|
||||
await loadTreeNodeChildren(node, { force: true });
|
||||
await restoreExpandedChildren(node, expandedIds, { force: true });
|
||||
if (isCurrentRefresh()) {
|
||||
await restoreExpandedChildren(node, expandedIds, { force: true }, isCurrentRefresh);
|
||||
}
|
||||
} catch (error) {
|
||||
node.children = previousChildren;
|
||||
node.hiddenChildren = previousHiddenChildren;
|
||||
node.objectCount = previousObjectCount;
|
||||
clearLoadedChildrenCache(node.id, { deletePersisted: false });
|
||||
for (const id of previousLoadedIds) loadedTreeNodeChildrenIds.value.add(id);
|
||||
for (const id of previousConfirmedEmptyIds) confirmedEmptyTreeNodeIds.value.add(id);
|
||||
// A stale failure must never overwrite a newer successful (including empty) result.
|
||||
if (isCurrentRefresh()) {
|
||||
const target = treeNodeInSidebarTree(node);
|
||||
if (target) {
|
||||
target.children = previousChildren;
|
||||
target.hiddenChildren = previousHiddenChildren;
|
||||
target.objectCount = previousObjectCount;
|
||||
target.isExpanded = previousExpanded;
|
||||
clearLoadedChildrenCache(target.id, { deletePersisted: false });
|
||||
for (const id of previousLoadedIds) loadedTreeNodeChildrenIds.value.add(id);
|
||||
for (const id of previousConfirmedEmptyIds) confirmedEmptyTreeNodeIds.value.add(id);
|
||||
}
|
||||
}
|
||||
throw error;
|
||||
} finally {
|
||||
if (ownsRefreshGeneration()) {
|
||||
activeTreeRefreshGenerations.delete(node.id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,22 @@ const OCEANBASE_ORACLE_COMPATIBLE_OJDBC_VERSION_KEY: &str = "compatibleOjdbcVers
|
|||
const OCEANBASE_ORACLE_COMPATIBLE_OJDBC_VERSION_PARAM: &str = "compatibleOjdbcVersion=8";
|
||||
const ZOOKEEPER_MIN_CONNECTION_TIMEOUT_MS: u64 = 15_000;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum AgentSessionRole {
|
||||
#[default]
|
||||
Workload,
|
||||
Metadata,
|
||||
}
|
||||
|
||||
impl AgentSessionRole {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Workload => "workload",
|
||||
Self::Metadata => "metadata",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn agent_jdbc_driver_class(config: &ConnectionConfig) -> &str {
|
||||
let driver_class = config.jdbc_driver_class.as_deref().unwrap_or("");
|
||||
if config.db_type == DatabaseType::H2
|
||||
|
|
@ -19,6 +35,16 @@ fn agent_jdbc_driver_class(config: &ConnectionConfig) -> &str {
|
|||
}
|
||||
|
||||
pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, database: &str) -> serde_json::Value {
|
||||
agent_connect_params_with_role(config, host, port, database, AgentSessionRole::Workload)
|
||||
}
|
||||
|
||||
pub fn agent_connect_params_with_role(
|
||||
config: &ConnectionConfig,
|
||||
host: &str,
|
||||
port: u16,
|
||||
database: &str,
|
||||
session_role: AgentSessionRole,
|
||||
) -> serde_json::Value {
|
||||
let agent_database = if config.db_type == DatabaseType::MongoDb {
|
||||
mongo_agent_database(config, database)
|
||||
} else if matches!(config.db_type, DatabaseType::Oracle | DatabaseType::OceanbaseOracle) {
|
||||
|
|
@ -83,6 +109,7 @@ pub fn agent_connect_params(config: &ConnectionConfig, host: &str, port: u16, da
|
|||
"informix_server": config.informix_server,
|
||||
"jdbc_driver_class": agent_jdbc_driver_class(config),
|
||||
"jdbc_driver_paths": &config.jdbc_driver_paths,
|
||||
"sessionRole": session_role.as_str(),
|
||||
});
|
||||
if config.db_type == DatabaseType::ZooKeeper {
|
||||
params["connection_timeout_ms"] = serde_json::json!(
|
||||
|
|
@ -662,6 +689,24 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_connect_params_default_to_workload_session_role() {
|
||||
let params = agent_connect_params(&config(DatabaseType::H2, Some("test")), "127.0.0.1", 9092, "test");
|
||||
assert_eq!(params["sessionRole"], "workload");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metadata_agent_connect_params_include_metadata_session_role() {
|
||||
let params = agent_connect_params_with_role(
|
||||
&config(DatabaseType::H2, Some("test")),
|
||||
"127.0.0.1",
|
||||
9092,
|
||||
"test",
|
||||
AgentSessionRole::Metadata,
|
||||
);
|
||||
assert_eq!(params["sessionRole"], "metadata");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mongodb_database_falls_back_to_uri_database() {
|
||||
let mut cfg = config(DatabaseType::MongoDb, None);
|
||||
|
|
|
|||
|
|
@ -894,7 +894,7 @@ impl AgentManager {
|
|||
agent_session_id: String,
|
||||
connect_params: serde_json::Value,
|
||||
connect_timeout: std::time::Duration,
|
||||
) -> Result<AgentDriverClient, String> {
|
||||
) -> Result<AgentDriverClient, crate::agent_runtime::SharedConnectionOpenError> {
|
||||
crate::agent_runtime::spawn_shared_connection_client(
|
||||
self,
|
||||
db_type,
|
||||
|
|
|
|||
|
|
@ -3,9 +3,22 @@ use std::time::Duration;
|
|||
|
||||
use crate::agent_manager::{AgentManager, DEFAULT_JRE_KEY};
|
||||
use crate::database_capabilities;
|
||||
use crate::db::agent_driver::{AgentDriverClient, AgentMethod, AgentRuntimeClient};
|
||||
use crate::db::agent_driver::{
|
||||
agent_session_disposition, AgentDriverClient, AgentMethod, AgentRuntimeClient, AgentSessionDisposition,
|
||||
};
|
||||
use crate::models::connection::DatabaseType;
|
||||
|
||||
pub struct SharedConnectionOpenError {
|
||||
pub(crate) message: String,
|
||||
pub(crate) runtime: Option<std::sync::Arc<AgentRuntimeClient>>,
|
||||
}
|
||||
|
||||
impl From<String> for SharedConnectionOpenError {
|
||||
fn from(message: String) -> Self {
|
||||
Self { message, runtime: None }
|
||||
}
|
||||
}
|
||||
|
||||
pub fn db_type_to_agent_key(db_type: &DatabaseType, driver_profile: Option<&str>) -> Option<&'static str> {
|
||||
database_capabilities::agent_key(db_type, driver_profile)
|
||||
}
|
||||
|
|
@ -64,7 +77,7 @@ pub async fn spawn_shared_connection_client(
|
|||
agent_session_id: String,
|
||||
connect_params: serde_json::Value,
|
||||
connect_timeout: Duration,
|
||||
) -> Result<AgentDriverClient, String> {
|
||||
) -> Result<AgentDriverClient, SharedConnectionOpenError> {
|
||||
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]);
|
||||
|
|
@ -103,12 +116,22 @@ pub async fn spawn_shared_connection_client(
|
|||
.call::<serde_json::Value>(AgentMethod::OpenSession.as_str(), session_params, Some(connect_timeout), None)
|
||||
.await
|
||||
{
|
||||
let open_error = shared_connection_open_error(err, runtime.clone());
|
||||
forget_unused_runtime_after_failed_open(manager, &runtime_key, &runtime_cell, &runtime).await;
|
||||
return Err(err);
|
||||
return Err(open_error);
|
||||
}
|
||||
Ok(AgentDriverClient::shared_session(runtime, agent_session_id))
|
||||
}
|
||||
|
||||
fn shared_connection_open_error(
|
||||
message: String,
|
||||
runtime: std::sync::Arc<AgentRuntimeClient>,
|
||||
) -> SharedConnectionOpenError {
|
||||
let runtime =
|
||||
(agent_session_disposition(&message) == Some(AgentSessionDisposition::ReplaceRuntime)).then_some(runtime);
|
||||
SharedConnectionOpenError { message, runtime }
|
||||
}
|
||||
|
||||
async fn forget_unused_runtime_after_failed_open(
|
||||
manager: &AgentManager,
|
||||
runtime_key: &str,
|
||||
|
|
@ -316,8 +339,9 @@ for line in sys.stdin:
|
|||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
crate::db::agent_driver::AgentLaunchSpec::new("python3")
|
||||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||||
.with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
|
|
@ -385,6 +409,21 @@ for line in sys.stdin:
|
|||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replace_runtime_open_error_reports_runtime_without_killing_it_directly() {
|
||||
let (manager, _cell, runtime, script_path) = test_shared_runtime("replace-runtime-open-error").await;
|
||||
let error = "Agent RPC error (-1): capacity exhausted\nDBX_AGENT_ERROR_DATA:{\"category\":\"resource\",\"sessionDisposition\":\"replace_runtime\"}";
|
||||
|
||||
let open_error = shared_connection_open_error(error.to_string(), runtime.clone());
|
||||
|
||||
assert!(open_error.runtime.as_ref().is_some_and(|failed| std::sync::Arc::ptr_eq(failed, &runtime)));
|
||||
assert!(!runtime.is_failed());
|
||||
|
||||
runtime.kill();
|
||||
drop(manager);
|
||||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn failed_open_cleanup_cannot_remove_runtime_after_reservation() {
|
||||
let (manager, cell, runtime, script_path) = test_shared_runtime("failed-open-race").await;
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -2371,7 +2371,6 @@ mod tests {
|
|||
#[test]
|
||||
fn concurrent_prefetch_only_allowed_for_multi_connection_pools() {
|
||||
use crate::connection::PoolKind;
|
||||
use std::sync::Arc;
|
||||
|
||||
// ChClient::new 只构造 HTTP 客户端，不发起连接
|
||||
let clickhouse = PoolKind::ClickHouse(crate::db::clickhouse_driver::ChClient::new(
|
||||
|
|
@ -2383,8 +2382,7 @@ mod tests {
|
|||
assert!(concurrent_metadata_prefetch_allowed(Some(&clickhouse)));
|
||||
|
||||
// Agent（JDBC sidecar）请求超时覆盖排队时间，必须回退串行
|
||||
let agent =
|
||||
PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(crate::db::agent_driver::AgentDriverClient::test_stub())));
|
||||
let agent = PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub());
|
||||
assert!(!concurrent_metadata_prefetch_allowed(Some(&agent)));
|
||||
|
||||
assert!(!concurrent_metadata_prefetch_allowed(None));
|
||||
|
|
|
|||
|
|
@ -285,15 +285,78 @@ impl AgentRuntimeClient {
|
|||
|
||||
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}"));
|
||||
return Err(format_agent_rpc_error(err));
|
||||
}
|
||||
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}"))
|
||||
}
|
||||
|
||||
const AGENT_RPC_ERROR_DATA_MARKER: &str = "\nDBX_AGENT_ERROR_DATA:";
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum AgentSessionDisposition {
|
||||
Keep,
|
||||
Quarantine,
|
||||
ReplaceRuntime,
|
||||
}
|
||||
|
||||
pub fn agent_rpc_error_category(error: &str) -> Option<String> {
|
||||
agent_rpc_error_data(error)?.get("category")?.as_str().map(str::to_string)
|
||||
}
|
||||
|
||||
pub fn agent_rpc_error_session_id(error: &str) -> Option<String> {
|
||||
agent_rpc_error_data(error)?.get("agentSessionId")?.as_str().map(str::to_string)
|
||||
}
|
||||
|
||||
pub fn agent_session_disposition(error: &str) -> Option<AgentSessionDisposition> {
|
||||
match agent_rpc_error_data(error)?.get("sessionDisposition")?.as_str()? {
|
||||
"keep" => Some(AgentSessionDisposition::Keep),
|
||||
"quarantine" => Some(AgentSessionDisposition::Quarantine),
|
||||
"replace_runtime" => Some(AgentSessionDisposition::ReplaceRuntime),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn format_agent_rpc_error(error: &Value) -> String {
|
||||
let message = error.get("message").and_then(Value::as_str).unwrap_or("Unknown agent error");
|
||||
let code = error.get("code").and_then(Value::as_i64).unwrap_or(-1);
|
||||
let mut formatted = format!("Agent RPC error ({code}): {message}");
|
||||
if let Some(data) = error.get("data").filter(|data| data.is_object()) {
|
||||
formatted.push_str(AGENT_RPC_ERROR_DATA_MARKER);
|
||||
formatted.push_str(&data.to_string());
|
||||
}
|
||||
formatted
|
||||
}
|
||||
|
||||
fn agent_rpc_error_data(error: &str) -> Option<Value> {
|
||||
let (_, data) = error.rsplit_once(AGENT_RPC_ERROR_DATA_MARKER)?;
|
||||
serde_json::from_str(data).ok()
|
||||
}
|
||||
|
||||
fn with_agent_rpc_error_session_id(error: String, agent_session_id: Option<&str>) -> String {
|
||||
let Some(agent_session_id) = agent_session_id else {
|
||||
return error;
|
||||
};
|
||||
let Some((message, data)) = error.rsplit_once(AGENT_RPC_ERROR_DATA_MARKER) else {
|
||||
return format!(
|
||||
"{error}{AGENT_RPC_ERROR_DATA_MARKER}{}",
|
||||
serde_json::json!({ "agentSessionId": agent_session_id })
|
||||
);
|
||||
};
|
||||
let Ok(mut data) = serde_json::from_str::<Value>(data) else {
|
||||
return format!(
|
||||
"{error}{AGENT_RPC_ERROR_DATA_MARKER}{}",
|
||||
serde_json::json!({ "agentSessionId": agent_session_id })
|
||||
);
|
||||
};
|
||||
let Some(data) = data.as_object_mut() else {
|
||||
return error;
|
||||
};
|
||||
data.insert("agentSessionId".to_string(), Value::String(agent_session_id.to_string()));
|
||||
format!("{message}{AGENT_RPC_ERROR_DATA_MARKER}{}", Value::Object(data.clone()))
|
||||
}
|
||||
|
||||
fn deserialize_cached_agent_result<T: DeserializeOwned>(result: Result<Value, String>) -> Result<T, String> {
|
||||
result
|
||||
.and_then(|value| serde_json::from_value(value).map_err(|e| format!("Failed to deserialize agent result: {e}")))
|
||||
|
|
@ -329,6 +392,66 @@ pub struct AgentDriverClient {
|
|||
cached_query: Option<CachedAgentQuery>,
|
||||
}
|
||||
|
||||
/// Keeps serialized session RPC access separate from the process-level fail-stop handle.
|
||||
/// A stuck RPC may hold `client` indefinitely, but it must never prevent terminating the
|
||||
/// shared Agent runtime after the pool has been removed from routing.
|
||||
pub struct PooledAgentClient {
|
||||
client: tokio::sync::Mutex<AgentDriverClient>,
|
||||
shared_runtime: Option<Arc<AgentRuntimeClient>>,
|
||||
agent_session_id: Option<String>,
|
||||
}
|
||||
|
||||
impl PooledAgentClient {
|
||||
pub fn new(client: AgentDriverClient) -> Self {
|
||||
let shared_runtime = client.shared_runtime.clone();
|
||||
let agent_session_id = client.agent_session_id.clone();
|
||||
Self { client: tokio::sync::Mutex::new(client), shared_runtime, agent_session_id }
|
||||
}
|
||||
|
||||
pub async fn lock(&self) -> tokio::sync::MutexGuard<'_, AgentDriverClient> {
|
||||
self.client.lock().await
|
||||
}
|
||||
|
||||
pub fn try_lock(&self) -> Result<tokio::sync::MutexGuard<'_, AgentDriverClient>, tokio::sync::TryLockError> {
|
||||
self.client.try_lock()
|
||||
}
|
||||
|
||||
pub fn shares_runtime_with(&self, other: &Self) -> bool {
|
||||
match (&self.shared_runtime, &other.shared_runtime) {
|
||||
(Some(runtime), Some(other_runtime)) => Arc::ptr_eq(runtime, other_runtime),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn uses_runtime(&self, runtime: &Arc<AgentRuntimeClient>) -> bool {
|
||||
self.shared_runtime.as_ref().is_some_and(|current| Arc::ptr_eq(current, runtime))
|
||||
}
|
||||
|
||||
pub fn matches_session_id(&self, session_id: &str) -> bool {
|
||||
self.agent_session_id.as_deref() == Some(session_id)
|
||||
}
|
||||
|
||||
pub fn is_runtime_available(&self) -> bool {
|
||||
self.shared_runtime.as_ref().is_none_or(|runtime| !runtime.is_failed())
|
||||
}
|
||||
|
||||
/// Terminates a protocol-v2 shared runtime without waiting for the logical session lock.
|
||||
/// Legacy single-session clients can only be killed immediately when they are not busy.
|
||||
pub fn fail_stop(&self) -> bool {
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
runtime.kill();
|
||||
return true;
|
||||
}
|
||||
match self.client.try_lock() {
|
||||
Ok(mut client) => {
|
||||
client.kill();
|
||||
true
|
||||
}
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct AgentLaunchSpec {
|
||||
pub program: PathBuf,
|
||||
|
|
@ -936,18 +1059,22 @@ impl AgentDriverClient {
|
|||
cancel_token: Option<CancellationToken>,
|
||||
) -> Result<T, String> {
|
||||
if let Some(runtime) = &self.shared_runtime {
|
||||
let agent_session_id = self.agent_session_id.clone();
|
||||
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")?;
|
||||
let session_id = 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;
|
||||
return runtime
|
||||
.call(method, params, timeout_duration, cancel_token)
|
||||
.await
|
||||
.map_err(|error| with_agent_rpc_error_session_id(error, agent_session_id.as_deref()));
|
||||
}
|
||||
self.next_id += 1;
|
||||
let id = self.next_id;
|
||||
|
|
@ -993,9 +1120,7 @@ impl AgentDriverClient {
|
|||
};
|
||||
|
||||
let result = if let Some(err) = resp.get("error") {
|
||||
let msg = err.get("message").and_then(|m| m.as_str()).unwrap_or("Unknown agent error");
|
||||
let code = err.get("code").and_then(|c| c.as_i64()).unwrap_or(-1);
|
||||
Err(format!("Agent RPC error ({code}): {msg}"))
|
||||
Err(format_agent_rpc_error(err))
|
||||
} else if let Some(result_val) = resp.get("result") {
|
||||
serde_json::from_value::<T>(result_val.clone())
|
||||
.map_err(|e| format!("Failed to deserialize agent result: {e}"))
|
||||
|
|
@ -2178,14 +2303,15 @@ impl Drop for AgentDriverClient {
|
|||
mod tests {
|
||||
use super::{
|
||||
agent_close_query_session_params, agent_handshake_params, agent_java_args, agent_java_args_with_extra,
|
||||
agent_java_args_with_extra_opts, agent_object_source_params, agent_proxy_env_vars, agent_schema_params,
|
||||
agent_schema_table_params, agent_supports_capability, agent_transaction_params, format_agent_process_error,
|
||||
agent_java_args_with_extra_opts, agent_object_source_params, agent_proxy_env_vars, agent_rpc_error_category,
|
||||
agent_rpc_error_session_id, agent_schema_params, agent_schema_table_params, agent_session_disposition,
|
||||
agent_supports_capability, agent_transaction_params, decode_agent_response, format_agent_process_error,
|
||||
format_agent_startup_error, is_agent_rpc_response_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, validate_dameng_java_system_properties, AgentCapability,
|
||||
AgentDriverClient, AgentHandshake, AgentKvMethod, AgentLaunchSpec, AgentMethod, AgentRuntimeClient,
|
||||
AgentTableReadCloseParams, AgentTableReadPageParams, AgentTableReadStartParams, MongoAgentMethod, StderrTail,
|
||||
AGENT_PROTOCOL_VERSION,
|
||||
AgentSessionDisposition, AgentTableReadCloseParams, AgentTableReadPageParams, AgentTableReadStartParams,
|
||||
MongoAgentMethod, StderrTail, AGENT_PROTOCOL_VERSION,
|
||||
};
|
||||
use std::io::Cursor;
|
||||
use std::io::Write;
|
||||
|
|
@ -2194,6 +2320,30 @@ mod tests {
|
|||
use std::time::{Duration, Instant};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
#[test]
|
||||
fn structured_agent_error_data_survives_legacy_string_boundary() {
|
||||
let response = serde_json::json!({
|
||||
"error": {
|
||||
"code": -1,
|
||||
"message": "connection lost",
|
||||
"data": {
|
||||
"category": "connection",
|
||||
"retryable": true,
|
||||
"sessionDisposition": "quarantine",
|
||||
"agentSessionId": "session-generation-1",
|
||||
"stage": "execute"
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let error = decode_agent_response::<serde_json::Value>(response).unwrap_err();
|
||||
|
||||
assert!(error.starts_with("Agent RPC error (-1): connection lost"));
|
||||
assert_eq!(agent_rpc_error_category(&error).as_deref(), Some("connection"));
|
||||
assert_eq!(agent_rpc_error_session_id(&error).as_deref(), Some("session-generation-1"));
|
||||
assert_eq!(agent_session_disposition(&error), Some(AgentSessionDisposition::Quarantine));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_java_args_include_oracle_network_compatibility_flags() {
|
||||
let args = agent_java_args("/tmp/dbx-agent-oracle.jar");
|
||||
|
|
@ -2494,8 +2644,9 @@ for line in sys.stdin:
|
|||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new("python3").with_args([script_path.to_string_lossy().to_string()]),
|
||||
AgentLaunchSpec::new(python).with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
|
|
@ -2520,6 +2671,7 @@ for line in sys.stdin:
|
|||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.contains("Agent RPC call timed out"));
|
||||
assert_eq!(agent_rpc_error_session_id(&error).as_deref(), Some("timeout-session"));
|
||||
let started = Instant::now();
|
||||
client
|
||||
.call_with_timeout::<serde_json::Value>("probe", serde_json::json!({}), Some(Duration::from_millis(500)))
|
||||
|
|
@ -2536,6 +2688,7 @@ for line in sys.stdin:
|
|||
.await
|
||||
.unwrap_err();
|
||||
assert!(error.contains("Agent RPC call timed out"));
|
||||
assert_eq!(agent_rpc_error_session_id(&error).as_deref(), Some("timeout-session"));
|
||||
let started = Instant::now();
|
||||
client
|
||||
.call_with_timeout::<serde_json::Value>("probe", serde_json::json!({}), Some(Duration::from_millis(500)))
|
||||
|
|
@ -2776,6 +2929,63 @@ for line in sys.stdin:
|
|||
let _ = std::fs::remove_file(script_path);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disconnect_reports_runtime_replacement_without_owning_runtime_routing() {
|
||||
let script_path =
|
||||
std::env::temp_dir().join(format!("dbx-agent-runtime-cleanup-saturation-{}.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':
|
||||
response = {
|
||||
'jsonrpc': '2.0',
|
||||
'id': req['id'],
|
||||
'result': {'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}
|
||||
}
|
||||
elif req['method'] == 'close_session':
|
||||
response = {
|
||||
'jsonrpc': '2.0',
|
||||
'id': req['id'],
|
||||
'error': {
|
||||
'code': -1,
|
||||
'message': 'Agent runtime resource limit reached',
|
||||
'data': {
|
||||
'category': 'resource',
|
||||
'retryable': False,
|
||||
'sessionDisposition': 'replace_runtime',
|
||||
'stage': 'close'
|
||||
}
|
||||
}
|
||||
}
|
||||
else:
|
||||
response = {'jsonrpc': '2.0', 'id': req['id'], 'result': {}}
|
||||
print(json.dumps(response), flush=True)
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||||
let runtime = AgentRuntimeClient::spawn(
|
||||
AgentLaunchSpec::new(python).with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
runtime.increment_session_count();
|
||||
let mut client = AgentDriverClient::shared_session(runtime.clone(), "session-1".to_string());
|
||||
|
||||
let error = client.disconnect().await.unwrap_err();
|
||||
|
||||
assert_eq!(agent_session_disposition(&error), Some(AgentSessionDisposition::ReplaceRuntime));
|
||||
assert!(!runtime.is_failed());
|
||||
runtime.kill();
|
||||
drop(client);
|
||||
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()));
|
||||
|
|
|
|||
|
|
@ -545,6 +545,9 @@ pub fn agent_close_query_session_params(session_id: &str) -> serde_json::Value {
|
|||
}
|
||||
|
||||
pub fn is_connection_error(err: &str) -> bool {
|
||||
if crate::db::agent_driver::agent_rpc_error_category(err).as_deref() == Some("connection") {
|
||||
return true;
|
||||
}
|
||||
let lower = err.to_lowercase();
|
||||
if is_dbx_query_timeout_error(&lower) || is_agent_rpc_timeout_error(&lower) {
|
||||
return false;
|
||||
|
|
@ -571,6 +574,9 @@ pub fn is_connection_error(err: &str) -> bool {
|
|||
|| lower.contains("idle")
|
||||
|| lower.contains("agent stdin not available")
|
||||
|| lower.contains("agent stdout not available")
|
||||
|| lower.contains("agent runtime terminated")
|
||||
|| lower.contains("agent runtime is unavailable")
|
||||
|| lower.contains("agent runtime unavailable")
|
||||
|| lower.contains("failed to write to agent stdin")
|
||||
|| lower.contains("failed to flush agent stdin")
|
||||
|| lower.contains("communicating with the server")
|
||||
|
|
@ -590,6 +596,15 @@ fn is_schema_reset_cleanup_error(lower: &str) -> bool {
|
|||
}
|
||||
|
||||
fn should_discard_agent_pool_after_error(err: &str) -> bool {
|
||||
if matches!(
|
||||
crate::db::agent_driver::agent_session_disposition(err),
|
||||
Some(
|
||||
crate::db::agent_driver::AgentSessionDisposition::Quarantine
|
||||
| crate::db::agent_driver::AgentSessionDisposition::ReplaceRuntime
|
||||
)
|
||||
) {
|
||||
return true;
|
||||
}
|
||||
let lower = err.to_lowercase();
|
||||
is_dbx_query_timeout_error(&lower)
|
||||
|| is_agent_rpc_timeout_error(&lower)
|
||||
|
|
@ -601,6 +616,19 @@ fn should_discard_agent_pool_after_error(err: &str) -> bool {
|
|||
}
|
||||
|
||||
pub fn pool_error_action(db_type: Option<DatabaseType>, err: &str) -> PoolErrorAction {
|
||||
if db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type))
|
||||
&& matches!(
|
||||
crate::db::agent_driver::agent_session_disposition(err),
|
||||
Some(
|
||||
crate::db::agent_driver::AgentSessionDisposition::Quarantine
|
||||
| crate::db::agent_driver::AgentSessionDisposition::ReplaceRuntime
|
||||
)
|
||||
)
|
||||
{
|
||||
// The connection may be replaced, but the result of the user operation is unknown.
|
||||
// Discard the session without replaying SQL, DDL, writes, or transactions.
|
||||
return PoolErrorAction::Discard;
|
||||
}
|
||||
let lower = err.to_lowercase();
|
||||
if db::sqlserver::is_driver_panic_error(err)
|
||||
|| (is_dbx_query_timeout_error(&lower) && should_discard_pool_after_query_timeout(db_type))
|
||||
|
|
@ -680,6 +708,41 @@ pub fn should_discard_pool_after_error(db_type: Option<DatabaseType>, err: &str)
|
|||
matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)
|
||||
}
|
||||
|
||||
async fn discard_pool_after_error(state: &AppState, pool_key: &str, db_type: Option<DatabaseType>, error: &str) {
|
||||
let action = pool_error_action(db_type, error);
|
||||
if !matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) {
|
||||
return;
|
||||
}
|
||||
|
||||
let replace_agent_runtime = matches!(
|
||||
crate::db::agent_driver::agent_session_disposition(error),
|
||||
Some(crate::db::agent_driver::AgentSessionDisposition::ReplaceRuntime)
|
||||
);
|
||||
if replace_agent_runtime {
|
||||
state.detach_pool_by_key(pool_key, true).await;
|
||||
} else {
|
||||
state.remove_pool_by_key(pool_key).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn discard_agent_pool_after_error(
|
||||
state: &AppState,
|
||||
pool_key: &str,
|
||||
client: &Arc<crate::db::agent_driver::PooledAgentClient>,
|
||||
db_type: Option<DatabaseType>,
|
||||
error: &str,
|
||||
) {
|
||||
let action = pool_error_action(db_type, error);
|
||||
if !matches!(action, PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) {
|
||||
return;
|
||||
}
|
||||
let replace_agent_runtime = matches!(
|
||||
crate::db::agent_driver::agent_session_disposition(error),
|
||||
Some(crate::db::agent_driver::AgentSessionDisposition::ReplaceRuntime)
|
||||
);
|
||||
state.detach_agent_pool_if_current(pool_key, client, replace_agent_runtime).await;
|
||||
}
|
||||
|
||||
fn query_pool_error_action(db_type: Option<DatabaseType>, sql: &str, err: &str) -> PoolErrorAction {
|
||||
match pool_error_action(db_type, err) {
|
||||
// A connection error does not prove that the database did not receive
|
||||
|
|
@ -1264,6 +1327,7 @@ pub async fn do_execute(
|
|||
}
|
||||
PoolKind::Agent(client) => {
|
||||
let client = client.clone();
|
||||
let source_client = client.clone();
|
||||
let sql = sql_for_execution_context(pool_db_type, sql, schema);
|
||||
let database = database.map(|s| s.to_string());
|
||||
let schema = schema_for_execution_context(pool_db_type, schema).map(|s| s.to_string());
|
||||
|
|
@ -1301,10 +1365,12 @@ pub async fn do_execute(
|
|||
.await
|
||||
.map(|result| truncate_result_with_max_rows(result, max_rows));
|
||||
if matches!(result.as_ref(), Err(err) if err == QUERY_CANCELED) {
|
||||
state.remove_pool_by_key(pool_key).await;
|
||||
state.detach_agent_pool_if_current(pool_key, &source_client, false).await;
|
||||
}
|
||||
if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(pool_db_type, err)) {
|
||||
state.remove_pool_by_key(pool_key).await;
|
||||
if let Err(err) = result.as_ref() {
|
||||
if err != QUERY_CANCELED {
|
||||
discard_agent_pool_after_error(state, pool_key, &source_client, pool_db_type, err).await;
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
|
|
@ -1493,7 +1559,11 @@ pub async fn execute_sql_statement_with_options(
|
|||
)
|
||||
}
|
||||
Some(PoolErrorAction::Discard) => {
|
||||
state.remove_pool_by_key(&pool_key).await;
|
||||
// Agent execution owns structured quarantine/runtime replacement before
|
||||
// returning. Native drivers retain the existing caller-side cleanup.
|
||||
if !db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) {
|
||||
state.remove_pool_by_key(&pool_key).await;
|
||||
}
|
||||
with_sql_context(result)
|
||||
}
|
||||
_ => with_sql_context(result),
|
||||
|
|
@ -2169,6 +2239,7 @@ pub async fn execute_statements(
|
|||
}
|
||||
};
|
||||
if let Some(client) = agent_client {
|
||||
let source_client = client.clone();
|
||||
check_read_only_for_connection_multi(state, &pool_key, statements).await?;
|
||||
let db_type = connection_database_type_for_pool_key(state, &pool_key).await;
|
||||
let execution_schema = schema_for_execution_context(db_type, schema);
|
||||
|
|
@ -2183,6 +2254,7 @@ pub async fn execute_statements(
|
|||
let mut client = client.lock().await;
|
||||
let database = if database.trim().is_empty() { None } else { Some(database) };
|
||||
let result = execute_multi_agent(&mut client, database, statements, execution_schema, timeout_secs).await;
|
||||
drop(client);
|
||||
match result {
|
||||
Ok(result) => return Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result }),
|
||||
Err(err) => {
|
||||
|
|
@ -2191,12 +2263,8 @@ pub async fn execute_statements(
|
|||
"Agent does not support execute_batch; falling back to statement-by-statement execution"
|
||||
);
|
||||
} else {
|
||||
match pool_error_action(connection_database_type(state, connection_id).await, &err) {
|
||||
PoolErrorAction::ReconnectAndRetry | PoolErrorAction::Discard => {
|
||||
let _ = state.remove_pool_by_key(&pool_key).await;
|
||||
}
|
||||
PoolErrorAction::Keep => {}
|
||||
}
|
||||
let db_type = connection_database_type(state, connection_id).await;
|
||||
discard_agent_pool_after_error(state, &pool_key, &source_client, db_type, &err).await;
|
||||
return Err(query_error_with_omitted_sql_context(&err, sql_ctx));
|
||||
}
|
||||
}
|
||||
|
|
@ -2220,15 +2288,18 @@ pub async fn execute_statements(
|
|||
total_affected += result.affected_rows;
|
||||
}
|
||||
Err(e) => {
|
||||
match pool_error_action(connection_database_type(state, connection_id).await, &e) {
|
||||
let db_type = connection_database_type(state, connection_id).await;
|
||||
match pool_error_action(db_type, &e) {
|
||||
PoolErrorAction::ReconnectAndRetry => {
|
||||
let db_opt = if database.is_empty() { None } else { Some(database) };
|
||||
let _ = state.reconnect_pool(connection_id, db_opt).await;
|
||||
}
|
||||
PoolErrorAction::Discard => {
|
||||
PoolErrorAction::Discard
|
||||
if !db_type.is_some_and(|db_type| database_capabilities::is_agent_type(&db_type)) =>
|
||||
{
|
||||
let _ = state.remove_pool_by_key(&pool_key).await;
|
||||
}
|
||||
PoolErrorAction::Keep => {}
|
||||
PoolErrorAction::Discard | PoolErrorAction::Keep => {}
|
||||
}
|
||||
return Err(query_error_with_omitted_sql_context(
|
||||
&format!("Statement {} failed: {}. Previous {} statement(s) may have been committed.", i + 1, e, i),
|
||||
|
|
@ -2560,11 +2631,10 @@ pub async fn execute_statements_in_transaction_on_pool(
|
|||
PoolKind::Mysql(mp, _mode) => TxPath::Mysql(mp.clone(), false),
|
||||
PoolKind::Sqlite(sq) => TxPath::Sqlite(sq.clone()),
|
||||
PoolKind::CloudflareD1(client) => TxPath::CloudflareD1(client.clone()),
|
||||
PoolKind::ClickHouse(_)
|
||||
| PoolKind::Rqlite(_)
|
||||
| PoolKind::Turso(_)
|
||||
| PoolKind::SqlServer(_)
|
||||
| PoolKind::Agent(_) => TxPath::Explicit,
|
||||
PoolKind::ClickHouse(_) | PoolKind::Rqlite(_) | PoolKind::Turso(_) | PoolKind::SqlServer(_) => {
|
||||
TxPath::Explicit
|
||||
}
|
||||
PoolKind::Agent(client) => TxPath::Agent(client.clone()),
|
||||
PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::HBase(_) => TxPath::None,
|
||||
PoolKind::DuckDbWorker(_)
|
||||
| PoolKind::Redis(_)
|
||||
|
|
@ -2606,6 +2676,13 @@ pub async fn execute_statements_in_transaction_on_pool(
|
|||
)
|
||||
.await
|
||||
}
|
||||
Some(TxPath::Agent(client)) => {
|
||||
let result = exec_tx_agent_inner(client.clone(), db_type, Some(database), statements, schema, start).await;
|
||||
if let Err(error) = result.as_ref() {
|
||||
discard_agent_pool_after_error(state, pool_key, &client, db_type, error).await;
|
||||
}
|
||||
return result;
|
||||
}
|
||||
Some(TxPath::Explicit) => {
|
||||
let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await;
|
||||
exec_tx_explicit_inner(state, pool_key, mysql_dialect, Some(database), statements, schema, start).await
|
||||
|
|
@ -2618,9 +2695,7 @@ pub async fn execute_statements_in_transaction_on_pool(
|
|||
};
|
||||
|
||||
if let Err(err) = result.as_ref() {
|
||||
if matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry) {
|
||||
state.remove_pool_by_key(pool_key).await;
|
||||
}
|
||||
discard_pool_after_error(state, pool_key, db_type, err).await;
|
||||
}
|
||||
|
||||
result
|
||||
|
|
@ -2632,6 +2707,7 @@ enum TxPath {
|
|||
Mysql(mysql_async::Pool, bool),
|
||||
Sqlite(db::sqlite::SqliteHandle),
|
||||
CloudflareD1(db::cloudflare_d1_driver::CloudflareD1Client),
|
||||
Agent(Arc<crate::db::agent_driver::PooledAgentClient>),
|
||||
Explicit,
|
||||
None,
|
||||
}
|
||||
|
|
@ -2859,24 +2935,6 @@ async fn exec_tx_explicit_inner(
|
|||
schema: Option<&str>,
|
||||
start: std::time::Instant,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
let conns = state.connections.read().await;
|
||||
if let Some(crate::connection::PoolKind::Agent(client)) = conns.get(pool_key) {
|
||||
let db_type = connection_database_type_for_pool_key(state, pool_key).await;
|
||||
let execution_schema = schema_for_execution_context(db_type, schema);
|
||||
let rewritten_statements;
|
||||
let statements = if qualifies_unqualified_agent_relations(db_type) {
|
||||
rewritten_statements =
|
||||
statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::<Vec<_>>();
|
||||
rewritten_statements.as_slice()
|
||||
} else {
|
||||
statements
|
||||
};
|
||||
let mut client = client.lock().await;
|
||||
let result: db::QueryResult = client.execute_transaction(database, statements, execution_schema).await?;
|
||||
return Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result });
|
||||
}
|
||||
drop(conns);
|
||||
|
||||
do_execute(
|
||||
state,
|
||||
pool_key,
|
||||
|
|
@ -2938,6 +2996,28 @@ async fn exec_tx_explicit_inner(
|
|||
})
|
||||
}
|
||||
|
||||
async fn exec_tx_agent_inner(
|
||||
client: Arc<crate::db::agent_driver::PooledAgentClient>,
|
||||
db_type: Option<DatabaseType>,
|
||||
database: Option<&str>,
|
||||
statements: &[String],
|
||||
schema: Option<&str>,
|
||||
start: std::time::Instant,
|
||||
) -> Result<db::QueryResult, String> {
|
||||
let execution_schema = schema_for_execution_context(db_type, schema);
|
||||
let rewritten_statements;
|
||||
let statements = if qualifies_unqualified_agent_relations(db_type) {
|
||||
rewritten_statements =
|
||||
statements.iter().map(|sql| sql_for_execution_context(db_type, sql, schema)).collect::<Vec<_>>();
|
||||
rewritten_statements.as_slice()
|
||||
} else {
|
||||
statements
|
||||
};
|
||||
let mut client = client.lock().await;
|
||||
let result: db::QueryResult = client.execute_transaction(database, statements, execution_schema).await?;
|
||||
Ok(db::QueryResult { execution_time_ms: start.elapsed().as_millis(), ..result })
|
||||
}
|
||||
|
||||
async fn exec_tx_none_inner(
|
||||
state: &AppState,
|
||||
pool_key: &str,
|
||||
|
|
@ -3884,6 +3964,141 @@ for line in sys.stdin:
|
|||
}
|
||||
}
|
||||
|
||||
async fn agent_error_state(
|
||||
disposition: &str,
|
||||
) -> (AppState, std::path::PathBuf, std::sync::Arc<crate::db::agent_driver::AgentRuntimeClient>) {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-query-agent-error-{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let script_path = dir.join("agent.py");
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
format!(
|
||||
r#"import json, sys
|
||||
print(json.dumps({{'ready': True}}), flush=True)
|
||||
for line in sys.stdin:
|
||||
req = json.loads(line)
|
||||
if req['method'] == 'handshake':
|
||||
response = {{
|
||||
'jsonrpc': '2.0',
|
||||
'id': req['id'],
|
||||
'result': {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}}
|
||||
}}
|
||||
elif req['method'] in ('execute_query', 'execute_batch', 'execute_transaction'):
|
||||
response = {{
|
||||
'jsonrpc': '2.0',
|
||||
'id': req['id'],
|
||||
'error': {{
|
||||
'code': -1,
|
||||
'message': 'injected Agent failure',
|
||||
'data': {{
|
||||
'category': 'resource',
|
||||
'retryable': False,
|
||||
'sessionDisposition': '{disposition}',
|
||||
'stage': 'execute'
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
else:
|
||||
response = {{'jsonrpc': '2.0', 'id': req['id'], 'result': {{}}}}
|
||||
print(json.dumps(response), flush=True)
|
||||
"#
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||||
.with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
runtime.increment_session_count();
|
||||
|
||||
let storage = Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = AppState::new(storage);
|
||||
state.configs.write().await.insert("conn-1".to_string(), test_connection_config(DatabaseType::Dameng));
|
||||
state.connections.write().await.insert(
|
||||
"conn-1".to_string(),
|
||||
PoolKind::agent(crate::db::agent_driver::AgentDriverClient::shared_session(
|
||||
runtime.clone(),
|
||||
"session-1".to_string(),
|
||||
)),
|
||||
);
|
||||
|
||||
(state, dir, runtime)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_query_replace_runtime_error_detaches_pool_and_stops_runtime() {
|
||||
let (state, dir, runtime) = agent_error_state("replace_runtime").await;
|
||||
|
||||
let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err();
|
||||
|
||||
assert!(error.contains("injected Agent failure"));
|
||||
assert!(!state.connections.read().await.contains_key("conn-1"));
|
||||
assert!(runtime.is_failed());
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_transaction_replace_runtime_error_detaches_pool_and_stops_runtime() {
|
||||
let (state, dir, runtime) = agent_error_state("replace_runtime").await;
|
||||
|
||||
let error = execute_statements_in_transaction_on_pool(
|
||||
&state,
|
||||
"conn-1",
|
||||
"conn-1",
|
||||
"",
|
||||
&["UPDATE test_table SET value = 1".to_string()],
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.contains("injected Agent failure"));
|
||||
assert!(!state.connections.read().await.contains_key("conn-1"));
|
||||
assert!(runtime.is_failed());
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_batch_replace_runtime_error_detaches_pool_and_stops_runtime() {
|
||||
let (state, dir, runtime) = agent_error_state("replace_runtime").await;
|
||||
|
||||
let error =
|
||||
execute_statements(&state, "conn-1", "", &["UPDATE test_table SET value = 1".to_string()], None, None)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(error.contains("injected Agent failure"));
|
||||
assert!(!state.connections.read().await.contains_key("conn-1"));
|
||||
assert!(runtime.is_failed());
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn agent_quarantine_error_removes_only_target_pool() {
|
||||
let (state, dir, runtime) = agent_error_state("quarantine").await;
|
||||
|
||||
let error = execute_sql_statement(&state, "conn-1", "", "SELECT 1", None, None).await.unwrap_err();
|
||||
|
||||
assert!(error.contains("injected Agent failure"));
|
||||
assert!(!state.connections.read().await.contains_key("conn-1"));
|
||||
assert!(!runtime.is_failed());
|
||||
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
struct FakeMysqlBatchExecutor {
|
||||
outcomes: std::collections::VecDeque<Result<db::QueryResult, String>>,
|
||||
executed: Vec<String>,
|
||||
|
|
@ -5010,16 +5225,38 @@ for line in sys.stdin:
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn structured_agent_disposition_controls_pool_recovery() {
|
||||
let quarantined = "Agent RPC error (-1): lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}";
|
||||
let replace_runtime = "Agent RPC error (-1): saturated\nDBX_AGENT_ERROR_DATA:{\"category\":\"resource\",\"sessionDisposition\":\"replace_runtime\"}";
|
||||
|
||||
assert!(should_discard_agent_pool_after_error(quarantined));
|
||||
assert!(should_discard_agent_pool_after_error(replace_runtime));
|
||||
assert!(is_connection_error(quarantined));
|
||||
assert_eq!(pool_error_action(Some(DatabaseType::Oracle), quarantined), PoolErrorAction::Discard);
|
||||
assert_eq!(pool_error_action(Some(DatabaseType::Oracle), replace_runtime), PoolErrorAction::Discard);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unavailable_agent_pipes_are_reconnectable_errors() {
|
||||
assert!(should_discard_agent_pool_after_error("Agent stdin not available"));
|
||||
assert!(should_discard_agent_pool_after_error("Agent stdout not available"));
|
||||
assert!(is_connection_error("Agent stdin not available"));
|
||||
assert!(is_connection_error("Agent stdout not available"));
|
||||
assert!(is_connection_error("Agent runtime terminated"));
|
||||
assert!(is_connection_error("Agent runtime is unavailable"));
|
||||
assert_eq!(
|
||||
pool_error_action(Some(DatabaseType::Oracle), "Agent stdin not available"),
|
||||
PoolErrorAction::ReconnectAndRetry
|
||||
);
|
||||
assert_eq!(
|
||||
pool_error_action(Some(DatabaseType::Oracle), "Agent runtime terminated"),
|
||||
PoolErrorAction::ReconnectAndRetry
|
||||
);
|
||||
assert_eq!(
|
||||
pool_error_action(Some(DatabaseType::Oracle), "Agent runtime is unavailable"),
|
||||
PoolErrorAction::ReconnectAndRetry
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ impl EphemeralAgentMetadataSession {
|
|||
let client_session_id = ephemeral_agent_metadata_session_id(db_config.as_ref(), task_kind);
|
||||
let cleanup_guard = match client_session_id.as_deref() {
|
||||
Some(client_session_id) => {
|
||||
state.client_session_pool_cleanup_guard(connection_id, database, client_session_id).await
|
||||
state.metadata_session_pool_cleanup_guard(connection_id, database, client_session_id).await
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
|
@ -350,7 +350,7 @@ pub async fn list_sqlserver_linked_server_tables_core(
|
|||
/// use the MySQL protocol, so this is a defensive no-op); the caller's
|
||||
/// flat-sidebar fallback then renders the standard database list.
|
||||
pub async fn list_doris_catalogs_core(state: &AppState, connection_id: &str) -> Result<Vec<db::CatalogInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
if let Some(PoolKind::Mysql(p, _)) = connections.get(&pool_key) {
|
||||
|
|
@ -373,7 +373,7 @@ pub async fn list_doris_catalog_databases_core(
|
|||
connection_id: &str,
|
||||
catalog: &str,
|
||||
) -> Result<Vec<db::DatabaseInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -417,7 +417,7 @@ pub async fn list_doris_catalog_tables_core(
|
|||
object_types: Option<&[String]>,
|
||||
table_name_filter: Option<&TableNameFilter>,
|
||||
) -> Result<Vec<db::TableInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -440,7 +440,7 @@ pub async fn get_doris_catalog_columns_core(
|
|||
database: &str,
|
||||
table: &str,
|
||||
) -> Result<Vec<db::ColumnInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -463,7 +463,7 @@ pub async fn get_doris_catalog_table_ddl_core(
|
|||
database: &str,
|
||||
table: &str,
|
||||
) -> Result<String, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -485,7 +485,7 @@ pub async fn list_doris_catalog_indexes_core(
|
|||
database: &str,
|
||||
table: &str,
|
||||
) -> Result<Vec<db::IndexInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, None).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, None, None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -654,7 +654,7 @@ async fn list_schema_infos_once(
|
|||
connection_id: &str,
|
||||
database: &str,
|
||||
) -> Result<Vec<db::SchemaInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let show_system_schemas = db_config.as_ref().is_some_and(|config| config.show_system_schemas);
|
||||
{
|
||||
|
|
@ -674,7 +674,7 @@ pub async fn list_data_types_core(
|
|||
database: &str,
|
||||
) -> Result<Vec<String>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) {
|
||||
|
|
@ -725,7 +725,7 @@ async fn list_schemas_once(
|
|||
database: &str,
|
||||
apply_visible_filter: bool,
|
||||
) -> Result<Vec<String>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let show_system_schemas = db_config.as_ref().is_some_and(|config| config.show_system_schemas);
|
||||
let visible_schema_filter = visible_schema_filter(db_config.as_ref(), database, apply_visible_filter);
|
||||
|
|
@ -892,8 +892,13 @@ pub async fn list_vector_collections_core(
|
|||
connection_id: &str,
|
||||
database: &str,
|
||||
) -> Result<Vec<db::vector_driver::CollectionInfo>, String> {
|
||||
let pool_key =
|
||||
state.get_or_create_pool(connection_id, if database.is_empty() { None } else { Some(database) }).await?;
|
||||
let pool_key = state
|
||||
.get_or_create_metadata_pool_for_session(
|
||||
connection_id,
|
||||
if database.is_empty() { None } else { Some(database) },
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let client = {
|
||||
let connections = state.connections.read().await;
|
||||
match connections.get(&pool_key) {
|
||||
|
|
@ -911,8 +916,13 @@ pub async fn get_vector_collection_detail_core(
|
|||
database: &str,
|
||||
collection: &str,
|
||||
) -> Result<db::vector_driver::CollectionInfo, String> {
|
||||
let pool_key =
|
||||
state.get_or_create_pool(connection_id, if database.is_empty() { None } else { Some(database) }).await?;
|
||||
let pool_key = state
|
||||
.get_or_create_metadata_pool_for_session(
|
||||
connection_id,
|
||||
if database.is_empty() { None } else { Some(database) },
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
let client = {
|
||||
let connections = state.connections.read().await;
|
||||
match connections.get(&pool_key) {
|
||||
|
|
@ -958,7 +968,8 @@ async fn get_table_comment_core_for_session(
|
|||
client_session_id: Option<&str>,
|
||||
) -> Result<Option<String>, String> {
|
||||
retry_metadata_connection_for_session(state, connection_id, Some(database), client_session_id, || async {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -1656,7 +1667,7 @@ async fn load_oracle_table_comments_for_objects(
|
|||
}
|
||||
|
||||
async fn oracle_agent_list_object_statistics(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
timeout_duration: Option<Duration>,
|
||||
|
|
@ -1696,7 +1707,7 @@ async fn oracle_agent_list_object_statistics(
|
|||
}
|
||||
|
||||
async fn dameng_agent_list_object_statistics(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
timeout_duration: Option<Duration>,
|
||||
|
|
@ -1735,7 +1746,7 @@ async fn dameng_agent_list_object_statistics(
|
|||
}
|
||||
|
||||
async fn agent_list_object_statistics(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
sql: String,
|
||||
|
|
@ -1778,7 +1789,8 @@ async fn list_tables_once(
|
|||
table_name_filter: Option<&TableNameFilter>,
|
||||
client_session_id: Option<&str>,
|
||||
) -> Result<Vec<db::TableInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -2518,18 +2530,19 @@ mod tests {
|
|||
ephemeral_agent_metadata_session_id, external_driver_uses_mysql_ddl, filter_mongodb_agent_collections,
|
||||
filter_mysql_system_databases_for_config, filter_object_infos, filter_table_infos, filter_visible_schema_names,
|
||||
gbase8a_object_statistics_sql, is_agent_postgres_metadata_fallback_config, is_mysql_external_driver_config,
|
||||
is_retryable_metadata_error, metadata_name_or_comment_matches, mysql_external_driver_ddl_from_query_result,
|
||||
mysql_external_driver_ddl_sql, mysql_object_source_ddl_column_index, mysql_object_source_sql,
|
||||
mysql_table_metadata_catalog, normalize_information_schema_table_type, oracle_columns_from_query_result,
|
||||
oracle_columns_sql, oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result,
|
||||
is_retryable_metadata_error, metadata_error_action, metadata_name_or_comment_matches,
|
||||
mysql_external_driver_ddl_from_query_result, mysql_external_driver_ddl_sql,
|
||||
mysql_object_source_ddl_column_index, mysql_object_source_sql, mysql_table_metadata_catalog,
|
||||
normalize_information_schema_table_type, oracle_columns_from_query_result, oracle_columns_sql,
|
||||
oracle_object_statistics_dba_segments_sql, oracle_object_statistics_from_query_result,
|
||||
oracle_object_statistics_rows_only_sql, oracle_object_statistics_sql,
|
||||
oracle_object_statistics_user_segments_sql, oracle_table_comment_from_query_result, oracle_table_comment_sql,
|
||||
oracle_table_comments_sql, presto_like_columns_from_query_result, presto_like_information_schema_columns_sql,
|
||||
presto_like_information_schema_tables_sql, presto_like_tables_from_query_result,
|
||||
presto_like_information_schema_tables_sql, presto_like_tables_from_query_result, replace_metadata_runtime,
|
||||
should_query_oracle_columns_via_sql_first, table_comments_from_query_result, table_name_filter_matches,
|
||||
tdengine_table_comment_like_pattern, tdengine_table_comment_sql, tdengine_table_comments_sql,
|
||||
uses_mongodb_agent_collection_listing, visible_schema_filter, TableNameFilter, TDENGINE_COMMENT_SEARCH_TIMEOUT,
|
||||
TDENGINE_LIKE_PATTERN_MAX_BYTES,
|
||||
uses_mongodb_agent_collection_listing, visible_schema_filter, MetadataErrorAction, TableNameFilter,
|
||||
TDENGINE_COMMENT_SEARCH_TIMEOUT, TDENGINE_LIKE_PATTERN_MAX_BYTES,
|
||||
};
|
||||
use super::{list_databases_core, list_tables_core};
|
||||
use crate::connection::{AppState, PoolKind};
|
||||
|
|
@ -2868,10 +2881,184 @@ mod tests {
|
|||
assert!(is_retryable_metadata_error("Pool not found"));
|
||||
assert!(is_retryable_metadata_error("connection reset by peer"));
|
||||
assert!(is_retryable_metadata_error("Agent RPC error (-1): dm.jdbc.driver.DMException: 网络通信异常"));
|
||||
assert!(is_retryable_metadata_error(
|
||||
"Agent RPC error (-1): connection lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}"
|
||||
));
|
||||
assert!(!is_retryable_metadata_error(
|
||||
"Agent RPC error (-1): connection text in SQL error\nDBX_AGENT_ERROR_DATA:{\"category\":\"sql\",\"sessionDisposition\":\"keep\"}"
|
||||
));
|
||||
assert!(!is_retryable_metadata_error(
|
||||
"Agent RPC error (-1): connection kept\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"keep\"}"
|
||||
));
|
||||
assert!(!is_retryable_metadata_error(
|
||||
"Agent RPC error (-1): runtime saturated\nDBX_AGENT_ERROR_DATA:{\"category\":\"resource\",\"sessionDisposition\":\"replace_runtime\"}"
|
||||
));
|
||||
assert!(!is_retryable_metadata_error("Unknown column 'email' in 'field list'"));
|
||||
assert!(!is_retryable_metadata_error("Access denied for user"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metadata_error_action_applies_fail_stop_to_every_attempt() {
|
||||
let quarantine = "Agent RPC error (-1): connection lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}";
|
||||
let replace_runtime = "Agent RPC error (-1): runtime saturated\nDBX_AGENT_ERROR_DATA:{\"category\":\"resource\",\"sessionDisposition\":\"replace_runtime\"}";
|
||||
let sql = "Agent RPC error (-1): syntax error\nDBX_AGENT_ERROR_DATA:{\"category\":\"sql\",\"sessionDisposition\":\"keep\"}";
|
||||
let db_type = Some(DatabaseType::Dameng);
|
||||
|
||||
assert_eq!(metadata_error_action(db_type, quarantine, false), MetadataErrorAction::Retry);
|
||||
assert_eq!(metadata_error_action(db_type, quarantine, true), MetadataErrorAction::Discard);
|
||||
assert_eq!(
|
||||
metadata_error_action(db_type, "Agent RPC call timed out (30s)", false),
|
||||
MetadataErrorAction::Discard
|
||||
);
|
||||
assert_eq!(metadata_error_action(db_type, replace_runtime, false), MetadataErrorAction::ReplaceRuntime);
|
||||
assert_eq!(metadata_error_action(db_type, replace_runtime, true), MetadataErrorAction::ReplaceRuntime);
|
||||
assert_eq!(metadata_error_action(db_type, sql, false), MetadataErrorAction::Return);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn metadata_fail_stop_detaches_base_pool_without_client_session() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-schema-metadata-fail-stop-{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = crate::connection::AppState::new(storage);
|
||||
let mut config = test_connection_config(DatabaseType::Dameng);
|
||||
config.id = "conn".to_string();
|
||||
state.configs.write().await.insert(config.id.clone(), config);
|
||||
state.connections.write().await.insert(
|
||||
"conn:analytics:role:metadata".to_string(),
|
||||
super::PoolKind::agent(crate::db::agent_driver::AgentDriverClient::test_stub()),
|
||||
);
|
||||
|
||||
replace_metadata_runtime(&state, "conn", Some("analytics"), None).await;
|
||||
|
||||
assert!(!state.connections.read().await.contains_key("conn:analytics:role:metadata"));
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn metadata_timeout_detaches_pool_without_replaying_operation() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-schema-metadata-timeout-{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = crate::connection::AppState::new(storage);
|
||||
let mut config = test_connection_config(DatabaseType::Dameng);
|
||||
config.id = "conn".to_string();
|
||||
state.configs.write().await.insert(config.id.clone(), config);
|
||||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||||
state.connections.write().await.insert("conn:role:metadata".to_string(), super::PoolKind::Sqlite(pool));
|
||||
let mut attempts = 0;
|
||||
|
||||
let result = super::retry_metadata_connection_for_session(&state, "conn", None, None, || {
|
||||
attempts += 1;
|
||||
async { Err::<(), _>("Agent RPC call timed out (30s)".to_string()) }
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(result.unwrap_err(), "Agent RPC call timed out (30s)");
|
||||
assert_eq!(attempts, 1);
|
||||
assert!(!state.connections.read().await.contains_key("conn:role:metadata"));
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn metadata_second_quarantine_detaches_replacement_pool() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-schema-metadata-quarantine-{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = crate::connection::AppState::new(storage);
|
||||
let mut config = test_connection_config(DatabaseType::Sqlite);
|
||||
config.id = "conn".to_string();
|
||||
config.host = ":memory:".to_string();
|
||||
config.password.clear();
|
||||
config.database = None;
|
||||
state.configs.write().await.insert(config.id.clone(), config);
|
||||
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
|
||||
state.connections.write().await.insert("conn".to_string(), super::PoolKind::Sqlite(pool));
|
||||
let mut attempts = 0;
|
||||
let quarantine = "Agent RPC error (-1): connection lost\nDBX_AGENT_ERROR_DATA:{\"category\":\"connection\",\"sessionDisposition\":\"quarantine\"}";
|
||||
|
||||
let result = super::retry_metadata_connection_for_session(&state, "conn", None, None, || {
|
||||
attempts += 1;
|
||||
async { Err::<(), _>(quarantine.to_string()) }
|
||||
})
|
||||
.await;
|
||||
|
||||
assert_eq!(result.unwrap_err(), quarantine);
|
||||
assert_eq!(attempts, 2);
|
||||
assert!(!state.connections.read().await.contains_key("conn"));
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn table_ddl_timeout_detaches_metadata_pool_without_replay() {
|
||||
let dir = std::env::temp_dir().join(format!("dbx-schema-table-ddl-timeout-{}", uuid::Uuid::new_v4()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let script_path = dir.join("table-ddl-timeout-agent.py");
|
||||
let call_count_path = dir.join("table-ddl-call-count");
|
||||
let call_count = serde_json::to_string(&call_count_path.to_string_lossy()).unwrap();
|
||||
std::fs::write(
|
||||
&script_path,
|
||||
format!(
|
||||
r#"import json, pathlib, sys
|
||||
call_count = pathlib.Path({call_count})
|
||||
print(json.dumps({{'ready': True}}), flush=True)
|
||||
for line in sys.stdin:
|
||||
req = json.loads(line)
|
||||
if req['method'] == 'handshake':
|
||||
result = {{'protocolVersion': 2, 'agentProtocolVersion': 2, 'capabilities': ['multi_session']}}
|
||||
response = {{'jsonrpc': '2.0', 'id': req['id'], 'result': result}}
|
||||
elif req['method'] in ('validate_session', 'validate_connection'):
|
||||
response = {{'jsonrpc': '2.0', 'id': req['id'], 'result': {{}}}}
|
||||
else:
|
||||
previous = int(call_count.read_text()) if call_count.exists() else 0
|
||||
call_count.write_text(str(previous + 1))
|
||||
response = {{
|
||||
'jsonrpc': '2.0',
|
||||
'id': req['id'],
|
||||
'error': {{
|
||||
'code': -1,
|
||||
'message': 'metadata timed out',
|
||||
'data': {{
|
||||
'category': 'timeout',
|
||||
'retryable': False,
|
||||
'sessionDisposition': 'quarantine',
|
||||
'stage': 'execute'
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
print(json.dumps(response), flush=True)
|
||||
"#
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
let python = if cfg!(windows) { "python" } else { "python3" };
|
||||
let runtime = crate::db::agent_driver::AgentRuntimeClient::spawn(
|
||||
crate::db::agent_driver::AgentLaunchSpec::new(python)
|
||||
.with_args([script_path.to_string_lossy().to_string()]),
|
||||
"test",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
runtime.increment_session_count();
|
||||
let client =
|
||||
crate::db::agent_driver::AgentDriverClient::shared_session(runtime.clone(), "metadata-session".to_string());
|
||||
let storage = crate::storage::Storage::open(&dir.join("storage.db")).await.unwrap();
|
||||
let state = crate::connection::AppState::new(storage);
|
||||
let mut config = test_connection_config(DatabaseType::Dameng);
|
||||
config.id = "conn".to_string();
|
||||
state.configs.write().await.insert(config.id.clone(), config);
|
||||
let pool_key = "conn:analytics:role:metadata";
|
||||
state.connections.write().await.insert(pool_key.to_string(), super::PoolKind::agent(client));
|
||||
|
||||
let error = super::get_table_ddl_core(&state, "conn", "analytics", "APP", "EVENTS", None).await.unwrap_err();
|
||||
|
||||
assert_eq!(crate::db::agent_driver::agent_rpc_error_category(&error).as_deref(), Some("timeout"));
|
||||
assert_eq!(std::fs::read_to_string(call_count_path).unwrap(), "1");
|
||||
assert!(!state.connections.read().await.contains_key(pool_key));
|
||||
runtime.kill();
|
||||
let _ = std::fs::remove_dir_all(dir);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn visible_schema_filter_only_applies_when_requested() {
|
||||
let mut config = test_connection_config(DatabaseType::Oracle);
|
||||
|
|
@ -4021,7 +4208,7 @@ async fn close_ephemeral_agent_metadata_session(
|
|||
let Some(client_session_id) = client_session_id else {
|
||||
return true;
|
||||
};
|
||||
match state.close_client_session_pool(connection_id, database, client_session_id).await {
|
||||
match state.close_metadata_session_pool(connection_id, database, client_session_id).await {
|
||||
Ok(_) => true,
|
||||
Err(error) => {
|
||||
log::warn!(
|
||||
|
|
@ -4047,7 +4234,9 @@ pub async fn completion_assistant_search_core(
|
|||
request.max_results
|
||||
);
|
||||
retry_metadata_connection(state, &request.connection_id, Some(&request.database), || async {
|
||||
let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?;
|
||||
let pool_key = state
|
||||
.get_or_create_metadata_pool_for_session(&request.connection_id, Some(&request.database), None)
|
||||
.await?;
|
||||
log::debug!("[schema][completion_assistant:start] {request_summary}");
|
||||
{
|
||||
let connections = state.connections.read().await;
|
||||
|
|
@ -4280,7 +4469,7 @@ async fn list_object_statistics_once(
|
|||
database: &str,
|
||||
schema: &str,
|
||||
) -> Result<Vec<db::ObjectStatistics>, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
try_sqlserver!(connections, &pool_key, list_object_statistics, schema);
|
||||
|
|
@ -4370,7 +4559,8 @@ async fn list_objects_once(
|
|||
object_types: Option<&[String]>,
|
||||
client_session_id: Option<&str>,
|
||||
) -> Result<ObjectListOutcome, String> {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let (mysql_limit, mysql_offset) =
|
||||
if filter.is_none_or(|value| value.trim().is_empty()) { (limit, offset) } else { (None, None) };
|
||||
|
|
@ -4567,7 +4757,8 @@ async fn list_completion_objects_once(
|
|||
schema: &str,
|
||||
client_session_id: Option<&str>,
|
||||
) -> Result<Vec<db::ObjectInfo>, String> {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
let connections = state.connections.read().await;
|
||||
|
|
@ -4726,17 +4917,120 @@ where
|
|||
F: FnMut() -> Fut,
|
||||
Fut: Future<Output = Result<T, String>>,
|
||||
{
|
||||
let result = operation().await;
|
||||
match result {
|
||||
Err(error) if is_retryable_metadata_error(&error) => {
|
||||
state.reconnect_pool_for_session(connection_id, database, client_session_id).await?;
|
||||
operation().await
|
||||
let db_type = {
|
||||
let configs = state.configs.read().await;
|
||||
configs.get(connection_id).map(|config| config.db_type)
|
||||
};
|
||||
let mut retried = false;
|
||||
loop {
|
||||
let result = operation().await;
|
||||
let action = result
|
||||
.as_ref()
|
||||
.err()
|
||||
.map(|error| metadata_error_action(db_type, error, retried))
|
||||
.unwrap_or(MetadataErrorAction::Return);
|
||||
match action {
|
||||
MetadataErrorAction::ReplaceRuntime => {
|
||||
state
|
||||
.detach_metadata_pool_after_error(
|
||||
connection_id,
|
||||
database,
|
||||
client_session_id,
|
||||
result.as_ref().err().expect("replace-runtime action requires an error"),
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
return result;
|
||||
}
|
||||
MetadataErrorAction::Discard => {
|
||||
state
|
||||
.detach_metadata_pool_after_error(
|
||||
connection_id,
|
||||
database,
|
||||
client_session_id,
|
||||
result.as_ref().err().expect("discard action requires an error"),
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
return result;
|
||||
}
|
||||
MetadataErrorAction::Retry => {
|
||||
retried = true;
|
||||
if let Err(error) =
|
||||
state.reconnect_metadata_pool_for_session(connection_id, database, client_session_id).await
|
||||
{
|
||||
match metadata_error_action(db_type, &error, true) {
|
||||
MetadataErrorAction::ReplaceRuntime => {
|
||||
state
|
||||
.detach_metadata_pool_after_error(
|
||||
connection_id,
|
||||
database,
|
||||
client_session_id,
|
||||
&error,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
MetadataErrorAction::Retry | MetadataErrorAction::Discard => {
|
||||
state
|
||||
.detach_metadata_pool_after_error(
|
||||
connection_id,
|
||||
database,
|
||||
client_session_id,
|
||||
&error,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
MetadataErrorAction::Return => {}
|
||||
}
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
MetadataErrorAction::Return => return result,
|
||||
}
|
||||
_ => result,
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
enum MetadataErrorAction {
|
||||
Retry,
|
||||
Discard,
|
||||
ReplaceRuntime,
|
||||
Return,
|
||||
}
|
||||
|
||||
fn metadata_error_action(db_type: Option<DatabaseType>, error: &str, retried: bool) -> MetadataErrorAction {
|
||||
if crate::db::agent_driver::agent_session_disposition(error)
|
||||
== Some(crate::db::agent_driver::AgentSessionDisposition::ReplaceRuntime)
|
||||
{
|
||||
MetadataErrorAction::ReplaceRuntime
|
||||
} else if !retried && is_retryable_metadata_error(error) {
|
||||
MetadataErrorAction::Retry
|
||||
} else if should_discard_pool_after_error(db_type, error) {
|
||||
MetadataErrorAction::Discard
|
||||
} else {
|
||||
MetadataErrorAction::Return
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
async fn replace_metadata_runtime(
|
||||
state: &AppState,
|
||||
connection_id: &str,
|
||||
database: Option<&str>,
|
||||
client_session_id: Option<&str>,
|
||||
) {
|
||||
state.replace_runtime_for_metadata_pool(connection_id, database, client_session_id).await;
|
||||
}
|
||||
|
||||
fn is_retryable_metadata_error(error: &str) -> bool {
|
||||
let category = crate::db::agent_driver::agent_rpc_error_category(error);
|
||||
if let Some(category) = category {
|
||||
return category == "connection"
|
||||
&& crate::db::agent_driver::agent_session_disposition(error)
|
||||
== Some(crate::db::agent_driver::AgentSessionDisposition::Quarantine);
|
||||
}
|
||||
error == "Pool not found" || crate::query::is_connection_error(error)
|
||||
}
|
||||
|
||||
|
|
@ -4791,7 +5085,7 @@ async fn get_columns_core_for_session_inner(
|
|||
let context_session_id = if use_client_session_context { client_session_id } else { None };
|
||||
retry_metadata_connection_for_session(state, connection_id, Some(database), client_session_id, || async {
|
||||
let pool_key = state
|
||||
.get_or_create_pool_for_session(connection_id, Some(database), client_session_id)
|
||||
.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id)
|
||||
.await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
|
|
@ -5074,7 +5368,7 @@ pub async fn get_sqlserver_column_metadata_core(
|
|||
table: &str,
|
||||
) -> Result<Vec<db::sqlserver::SqlServerColumnMetadata>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let connections = state.connections.read().await;
|
||||
try_sqlserver!(connections, &pool_key, get_column_metadata, schema, table);
|
||||
Err("SQL Server column metadata requires a native SQL Server connection".to_string())
|
||||
|
|
@ -5158,7 +5452,8 @@ async fn list_indexes_core_for_session(
|
|||
client_session_id: Option<&str>,
|
||||
) -> Result<Vec<db::IndexInfo>, String> {
|
||||
retry_metadata_connection_for_session(state, connection_id, Some(database), client_session_id, || async {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -5238,7 +5533,8 @@ async fn list_foreign_keys_core_for_session(
|
|||
client_session_id: Option<&str>,
|
||||
) -> Result<Vec<db::ForeignKeyInfo>, String> {
|
||||
retry_metadata_connection_for_session(state, connection_id, Some(database), client_session_id, || async {
|
||||
let pool_key = state.get_or_create_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let pool_key =
|
||||
state.get_or_create_metadata_pool_for_session(connection_id, Some(database), client_session_id).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -5286,7 +5582,7 @@ pub async fn list_triggers_core(
|
|||
return Ok(vec![]);
|
||||
}
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -5332,7 +5628,7 @@ pub async fn list_constraints_core(
|
|||
table: &str,
|
||||
) -> Result<Vec<db::ConstraintInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
|
||||
|
|
@ -5353,7 +5649,7 @@ pub async fn list_partitions_core(
|
|||
table: &str,
|
||||
) -> Result<Vec<db::PartitionInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
|
||||
|
|
@ -5374,7 +5670,7 @@ pub async fn list_subpartitions_core(
|
|||
table: &str,
|
||||
) -> Result<Vec<db::SubpartitionInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
if let Some(client) = extract_pool!(&connections, &pool_key, Agent) {
|
||||
|
|
@ -5396,7 +5692,7 @@ pub async fn list_functions_core(
|
|||
schema: &str,
|
||||
) -> Result<Vec<db::FunctionInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
||||
|
|
@ -5416,7 +5712,7 @@ pub async fn list_sequences_core(
|
|||
with_last_values: bool,
|
||||
) -> Result<Vec<db::SequenceInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
|
@ -5439,7 +5735,7 @@ pub async fn list_rules_core(
|
|||
schema: &str,
|
||||
) -> Result<Vec<db::RuleInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
||||
|
|
@ -5458,7 +5754,7 @@ pub async fn list_extensions_core(
|
|||
schema: Option<&str>,
|
||||
) -> Result<Vec<db::ExtensionInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) {
|
||||
let connections = state.connections.read().await;
|
||||
|
|
@ -5486,7 +5782,7 @@ pub async fn list_available_extensions_core(
|
|||
database: &str,
|
||||
) -> Result<Vec<db::ExtensionInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
if db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Kingbase) {
|
||||
let connections = state.connections.read().await;
|
||||
|
|
@ -5519,7 +5815,7 @@ pub async fn list_owners_core(
|
|||
schema: &str,
|
||||
) -> Result<Vec<db::OwnerInfo>, String> {
|
||||
retry_metadata_connection(state, connection_id, Some(database), || async {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let connections = state.connections.read().await;
|
||||
let pool = connections.get(&pool_key).ok_or("Pool not found")?;
|
||||
|
||||
|
|
@ -5600,7 +5896,21 @@ async fn get_table_ddl_core_with_options(
|
|||
return Ok(source.source);
|
||||
}
|
||||
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
retry_metadata_connection(state, connection_id, Some(database), || {
|
||||
get_table_ddl_once(state, connection_id, database, schema, table, include_postgres_access)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn get_table_ddl_once(
|
||||
state: &AppState,
|
||||
connection_id: &str,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
include_postgres_access: bool,
|
||||
) -> Result<String, String> {
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
|
||||
{
|
||||
|
|
@ -6362,7 +6672,7 @@ async fn get_object_source_once(
|
|||
signature: Option<&str>,
|
||||
relation_name: Option<&str>,
|
||||
) -> Result<db::ObjectSource, String> {
|
||||
let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?;
|
||||
let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?;
|
||||
let db_config = connection_config(state, connection_id).await;
|
||||
let source = {
|
||||
let connections = state.connections.read().await;
|
||||
|
|
@ -6525,7 +6835,7 @@ pub fn oracle_list_objects_sql(schema: &str) -> String {
|
|||
}
|
||||
|
||||
async fn oracle_agent_list_objects(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
timeout_duration: Option<Duration>,
|
||||
|
|
@ -6565,7 +6875,7 @@ async fn oracle_agent_list_objects(
|
|||
}
|
||||
|
||||
async fn oracle_agent_object_source(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
name: &str,
|
||||
|
|
@ -6585,7 +6895,7 @@ async fn oracle_agent_object_source(
|
|||
}
|
||||
|
||||
async fn oracle_agent_table_ddl(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
|
|
@ -6666,7 +6976,7 @@ fn append_oracle_comments_to_ddl(
|
|||
}
|
||||
|
||||
async fn db2_agent_table_ddl(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: &str,
|
||||
table: &str,
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ fn list_available_extensions_sql(catalog: ExtensionCatalog) -> String {
|
|||
}
|
||||
|
||||
async fn query_result(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
sql: &str,
|
||||
max_rows: usize,
|
||||
|
|
@ -88,7 +88,7 @@ async fn query_result(
|
|||
}
|
||||
|
||||
async fn query_result_with_catalog_fallback(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
sys_sql: String,
|
||||
pg_sql: String,
|
||||
|
|
@ -104,7 +104,7 @@ async fn query_result_with_catalog_fallback(
|
|||
}
|
||||
|
||||
pub(super) async fn list_extensions(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
schema: Option<&str>,
|
||||
timeout_duration: Option<Duration>,
|
||||
|
|
@ -122,7 +122,7 @@ pub(super) async fn list_extensions(
|
|||
}
|
||||
|
||||
pub(super) async fn list_available_extensions(
|
||||
client: Arc<tokio::sync::Mutex<db::agent_driver::AgentDriverClient>>,
|
||||
client: Arc<db::agent_driver::PooledAgentClient>,
|
||||
database: &str,
|
||||
timeout_duration: Option<Duration>,
|
||||
) -> Result<Vec<db::ExtensionInfo>, String> {
|
||||
|
|
|
|||
|
|
@ -162,7 +162,7 @@ async fn connect_agent_pool(
|
|||
}
|
||||
}
|
||||
|
||||
Ok(PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(client))))
|
||||
Ok(PoolKind::agent(client))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -1216,7 +1216,7 @@ pub async fn connect_db(
|
|||
.await
|
||||
.map_err(|err| mongo_legacy_error_with_auth_hint(&err))?;
|
||||
state.ensure_current_connection_attempt(&id, Some(attempt)).await?;
|
||||
PoolKind::Agent(std::sync::Arc::new(tokio::sync::Mutex::new(client)))
|
||||
PoolKind::agent(client)
|
||||
} else {
|
||||
let native_err = match db::mongo_driver::connect(&url, connect_timeout, idle_timeout).await {
|
||||
Ok(client) => {
|
||||
|
|
@ -1266,7 +1266,7 @@ pub async fn connect_db(
|
|||
mark_mongo_legacy_driver(&mut connected_config);
|
||||
connected_db_config = metadata_connection_config(&connected_config);
|
||||
persist_mongo_legacy_driver_profile(state.inner(), &connected_config).await?;
|
||||
PoolKind::Agent(std::sync::Arc::new(tokio::sync::Mutex::new(client)))
|
||||
PoolKind::agent(client)
|
||||
} else {
|
||||
return Err(native_err);
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue