github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 06:00:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(jdbc): honor query execution settings

Browse Source
This commit is contained in:
t8y2
2026-06-02 20:07:54 +08:00
parent 8d26602c91
commit 2f765fa69c
7 changed files with 247 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/agent_driver.rs
+77 -28
View File
@@ -286,6 +286,17 @@ impl AgentDriverClient {
&mut self,
method: &str,
params: Value,
) -> Result<T, String> {
self.call_with_timeout(method, params, Some(Duration::from_secs(RPC_TIMEOUT_SECS))).await
}
/// Send a JSON-RPC 2.0 request and wait for the response.
/// `None` disables the client-side RPC timeout for long-running query calls.
pub async fn call_with_timeout<T: DeserializeOwned + Send + 'static>(
&mut self,
method: &str,
params: Value,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.next_id += 1;
let id = self.next_id;
@@ -317,37 +328,42 @@ impl AgentDriverClient {
// Read response from stdout (blocking, with timeout)
let mut reader = self.stdout.take().ok_or("Agent stdout not available")?;
let (returned_reader, result) = tokio::time::timeout(
Duration::from_secs(RPC_TIMEOUT_SECS),
tokio::task::spawn_blocking(move || {
let line = match read_agent_line(&mut reader, "response") {
Ok(line) => line,
Err(e) => return (reader, Err(e)),
};
let response_task = tokio::task::spawn_blocking(move || {
let line = match read_agent_line(&mut reader, "response") {
Ok(line) => line,
Err(e) => return (reader, Err(e)),
};
let resp: Value = match serde_json::from_str(line.trim()) {
Ok(v) => v,
Err(e) => {
return (reader, Err(format!("Invalid JSON response from agent: {e}")));
}
};
let resp: Value = match serde_json::from_str(line.trim()) {
Ok(v) => v,
Err(e) => {
return (reader, Err(format!("Invalid JSON response from agent: {e}")));
}
};
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}"))
} 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}"))
} else {
Err(format!("Agent response missing both 'result' and 'error': {line}"))
};
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}"))
} 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}"))
} else {
Err(format!("Agent response missing both 'result' and 'error': {line}"))
};
(reader, result)
}),
)
.await
.map_err(|_| format!("Agent RPC call timed out ({RPC_TIMEOUT_SECS}s)"))?
(reader, result)
});
let (returned_reader, result) = match timeout_duration {
Some(duration) => match tokio::time::timeout(duration, response_task).await {
Ok(result) => result,
Err(_) => {
self.kill();
return Err(format!("Agent RPC call timed out ({}s)", duration.as_secs()));
}
},
None => response_task.await,
}
.map_err(|e| format!("Agent RPC task failed: {e}"))?;
let _ = self.stdout.insert(returned_reader);
@@ -362,6 +378,15 @@ impl AgentDriverClient {
self.call(method.as_str(), params).await
}
pub async fn call_method_with_timeout<T: DeserializeOwned + Send + 'static>(
&mut self,
method: AgentMethod,
params: Value,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.call_with_timeout(method.as_str(), params, timeout_duration).await
}
pub async fn connect(&mut self, params: Value) -> Result<Value, String> {
self.call_method(AgentMethod::Connect, params).await
}
@@ -458,6 +483,14 @@ impl AgentDriverClient {
self.call_method(AgentMethod::ExecuteQuery, params).await
}
pub async fn execute_query_with_timeout<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.call_method_with_timeout(AgentMethod::ExecuteQuery, params, timeout_duration).await
}
pub async fn execute_query_page<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
@@ -465,10 +498,26 @@ impl AgentDriverClient {
self.call_method(AgentMethod::ExecuteQueryPage, params).await
}
pub async fn execute_query_page_with_timeout<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.call_method_with_timeout(AgentMethod::ExecuteQueryPage, params, timeout_duration).await
}
pub async fn fetch_query_page<T: DeserializeOwned + Send + 'static>(&mut self, params: Value) -> Result<T, String> {
self.call_method(AgentMethod::FetchQueryPage, params).await
}
pub async fn fetch_query_page_with_timeout<T: DeserializeOwned + Send + 'static>(
&mut self,
params: Value,
timeout_duration: Option<Duration>,
) -> Result<T, String> {
self.call_method_with_timeout(AgentMethod::FetchQueryPage, params, timeout_duration).await
}
pub async fn close_query_session<T: DeserializeOwned + Send + 'static>(
&mut self,
session_id: &str,
crates/dbx-core/src/plugins.rs
+23 -14
View File
@@ -242,26 +242,35 @@ impl PluginDriverSession {
}
pub async fn invoke<T>(&self, method: &str, params: serde_json::Value) -> Result<T, String>
where
T: DeserializeOwned,
{
self.invoke_with_timeout(method, params, Some(PLUGIN_REQUEST_TIMEOUT)).await
}
pub async fn invoke_with_timeout<T>(
&self,
method: &str,
params: serde_json::Value,
timeout_duration: Option<Duration>,
) -> Result<T, String>
where
T: DeserializeOwned,
{
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
let result = timeout(PLUGIN_REQUEST_TIMEOUT, async {
let invoke = async {
let mut process = self.process.lock().await;
self.invoke_locked(&mut process, request_id, method, params).await
})
.await;
match result {
Ok(result) => result,
Err(_) => {
self.kill().await;
Err(format!(
"Plugin '{}' timed out after {} seconds",
self.plugin.manifest.id,
PLUGIN_REQUEST_TIMEOUT.as_secs()
))
}
};
match timeout_duration {
Some(duration) => match timeout(duration, invoke).await {
Ok(result) => result,
Err(_) => {
self.kill().await;
Err(format!("Plugin '{}' timed out after {} seconds", self.plugin.manifest.id, duration.as_secs()))
}
},
None => invoke.await,
}
}
crates/dbx-core/src/query.rs
+55 -10
View File
@@ -316,6 +316,9 @@ pub fn agent_execute_query_params(
if let Some(fetch_size) = options.fetch_size {
params["fetchSize"] = serde_json::json!(fetch_size);
}
if let Some(timeout_secs) = options.timeout_secs {
params["timeoutSecs"] = serde_json::json!(timeout_secs);
}
params
}
@@ -339,6 +342,9 @@ pub fn agent_execute_query_page_params(
if let Some(fetch_size) = options.fetch_size {
params["fetchSize"] = serde_json::json!(fetch_size);
}
if let Some(timeout_secs) = options.timeout_secs {
params["timeoutSecs"] = serde_json::json!(timeout_secs);
}
params
}
@@ -580,18 +586,19 @@ pub async fn do_execute(
let database = database.map(|s| s.to_string());
let schema = schema.map(|s| s.to_string());
let max_rows = options.max_rows;
let rpc_timeout = query_timeout;
drop(connections);
wait_for_query_opt(cancel_token, query_timeout, async move {
let mut client = client.lock().await;
if let Some(session_id) = options.result_session_id.as_deref() {
let params = agent_fetch_query_page_params(session_id, options.page_size.unwrap_or(MAX_ROWS));
client.fetch_query_page(params).await
client.fetch_query_page_with_timeout(params, rpc_timeout).await
} else if options.page_size.is_some() {
let params = agent_execute_query_page_params(&sql, database.as_deref(), schema.as_deref(), options);
client.execute_query_page(params).await
client.execute_query_page_with_timeout(params, rpc_timeout).await
} else {
let params = agent_execute_query_params(&sql, database.as_deref(), schema.as_deref(), options);
client.execute_query(params).await
client.execute_query_with_timeout(params, rpc_timeout).await
}
})
.await
@@ -621,10 +628,12 @@ pub async fn do_execute(
let schema = schema.map(str::to_string);
let database = config.effective_database().unwrap_or("").to_string();
let max_rows = options.max_rows;
let plugin_timeout = query_timeout;
drop(connections);
wait_for_query_opt(cancel_token, query_timeout, async move {
let params = external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref());
session.invoke::<db::QueryResult>("executeQuery", params).await
let params =
external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options);
session.invoke_with_timeout::<db::QueryResult>("executeQuery", params, plugin_timeout).await
})
.await
.map(|result| normalize_query_result_for_js(truncate_result_with_max_rows(result, max_rows)))
@@ -637,13 +646,22 @@ fn external_driver_query_params(
sql: &str,
database: &str,
schema: Option<&str>,
options: &QueryExecutionOptions,
) -> serde_json::Value {
serde_json::json!({
let mut params = serde_json::json!({
"connection": config,
"sql": sql,
"database": database,
"schema": schema,
})
"maxRows": options.max_rows.unwrap_or(MAX_ROWS),
});
if let Some(fetch_size) = options.fetch_size {
params["fetchSize"] = serde_json::json!(fetch_size);
}
if let Some(timeout_secs) = options.timeout_secs {
params["timeoutSecs"] = serde_json::json!(timeout_secs);
}
params
}
pub async fn execute_sql_statement(
@@ -1534,12 +1552,26 @@ mod tests {
one_time: false,
};
let params = external_driver_query_params(&config, "SELECT * FROM events", "analytics", Some("app"));
let params = external_driver_query_params(
&config,
"SELECT * FROM events",
"analytics",
Some("app"),
&QueryExecutionOptions {
max_rows: Some(500),
fetch_size: Some(250),
timeout_secs: Some(600),
..Default::default()
},
);
assert_eq!(params["connection"]["id"], "jdbc-1");
assert_eq!(params["sql"], "SELECT * FROM events");
assert_eq!(params["database"], "analytics");
assert_eq!(params["schema"], "app");
assert_eq!(params["maxRows"], 500);
assert_eq!(params["fetchSize"], 250);
assert_eq!(params["timeoutSecs"], 600);
}
#[test]
@@ -1548,7 +1580,12 @@ mod tests {
"SELECT * FROM events",
Some("analytics"),
Some("app"),
QueryExecutionOptions { max_rows: Some(500), fetch_size: Some(250), ..Default::default() },
QueryExecutionOptions {
max_rows: Some(500),
fetch_size: Some(250),
timeout_secs: Some(600),
..Default::default()
},
);
assert_eq!(params["sql"], "SELECT * FROM events");
@@ -1556,6 +1593,7 @@ mod tests {
assert_eq!(params["schema"], "app");
assert_eq!(params["maxRows"], 500);
assert_eq!(params["fetchSize"], 250);
assert_eq!(params["timeoutSecs"], 600);
}
#[test]
@@ -1567,6 +1605,7 @@ mod tests {
assert!(params.get("schema").is_none());
assert_eq!(params["maxRows"], MAX_ROWS);
assert!(params.get("fetchSize").is_none());
assert!(params.get("timeoutSecs").is_none());
}
#[test]
@@ -1575,7 +1614,12 @@ mod tests {
"SELECT * FROM events",
Some("analytics"),
Some("app"),
QueryExecutionOptions { page_size: Some(500), fetch_size: Some(250), ..Default::default() },
QueryExecutionOptions {
page_size: Some(500),
fetch_size: Some(250),
timeout_secs: Some(600),
..Default::default()
},
);
assert_eq!(params["sql"], "SELECT * FROM events");
@@ -1583,6 +1627,7 @@ mod tests {
assert_eq!(params["schema"], "app");
assert_eq!(params["pageSize"], 500);
assert_eq!(params["fetchSize"], 250);
assert_eq!(params["timeoutSecs"], 600);
assert_eq!(params["maxRows"], MAX_ROWS);
}
plugins/jdbc/bin/dbx-jdbc-plugin
+5 -1
View File
@@ -17,4 +17,8 @@ else
exit 127
fi
exec "$JAVA_BIN" -jar "$DIR/../lib/dbx-jdbc-plugin.jar"
exec "$JAVA_BIN" \
-Dfile.encoding=UTF-8 \
-Dsun.stdout.encoding=UTF-8 \
-Dsun.stderr.encoding=UTF-8 \
-jar "$DIR/../lib/dbx-jdbc-plugin.jar"
plugins/jdbc/bin/dbx-jdbc-plugin.bat
+1 -1
View File
@@ -28,4 +28,4 @@ echo Java runtime not found. Install Java or the optional DBX JDBC runtime. >&2
exit /b 127
:run
"%JAVA_BIN%" -jar "%~dp0..\lib\dbx-jdbc-plugin.jar"
"%JAVA_BIN%" -Dfile.encoding=UTF-8 -Dsun.stdout.encoding=UTF-8 -Dsun.stderr.encoding=UTF-8 -jar "%~dp0..\lib\dbx-jdbc-plugin.jar"
plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java
+56 -7
View File
@@ -6,7 +6,9 @@ import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.net.URLEncoder;
import java.math.BigDecimal;
import java.net.URL;
@@ -62,15 +64,19 @@ public final class DbxJdbcPlugin {
}
public static void main(String[] args) throws Exception {
try (BufferedReader reader = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8))) {
try (
BufferedReader reader = new BufferedReader(new InputStreamReader(System.in, StandardCharsets.UTF_8));
BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(System.out, StandardCharsets.UTF_8))
) {
String line;
while ((line = reader.readLine()) != null) {
if (line.isBlank()) {
continue;
}
ObjectNode response = handleLine(line);
System.out.println(MAPPER.writeValueAsString(response));
System.out.flush();
writer.write(MAPPER.writeValueAsString(response));
writer.newLine();
writer.flush();
if (response.path("_dbx_close").asBoolean(false)) {
break;
}
@@ -120,7 +126,10 @@ public final class DbxJdbcPlugin {
connection,
requireText(params, "sql"),
optionalText(params, "database"),
optionalText(params, "schema")
optionalText(params, "schema"),
positiveInt(params, "maxRows", MAX_ROWS),
nonNegativeInt(params, "fetchSize", 0),
nonNegativeInt(params, "timeoutSecs", -1)
);
case "listDatabases" -> listDatabases(connection);
case "listSchemas" -> listSchemas(connection, optionalText(params, "database"));
@@ -218,12 +227,20 @@ public final class DbxJdbcPlugin {
return sharedConnection;
}
private static JsonNode executeQuery(JsonNode connection, String sql, String database, String schema) throws SQLException {
private static JsonNode executeQuery(
JsonNode connection,
String sql,
String database,
String schema,
int maxRows,
int fetchSize,
int timeoutSecs
) throws SQLException {
long start = System.nanoTime();
Connection conn = openConnection(connection);
applyExecutionContext(connection, conn, database, schema);
try (Statement statement = conn.createStatement()) {
statement.setMaxRows(MAX_ROWS + 1);
applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs);
boolean hasResultSet = statement.execute(trimStatementSql(sql));
ObjectNode result = MAPPER.createObjectNode();
ArrayNode columns = MAPPER.createArrayNode();
@@ -239,7 +256,7 @@ public final class DbxJdbcPlugin {
columns.add(label == null || label.isBlank() ? meta.getColumnName(i) : label);
}
while (rs.next()) {
if (rows.size() >= MAX_ROWS) {
if (rows.size() >= maxRows) {
truncated = true;
break;
}
@@ -261,6 +278,23 @@ public final class DbxJdbcPlugin {
}
}
private static void applyStatementOptions(Statement statement, int maxRows, int fetchSize, int timeoutSecs)
throws SQLException {
statement.setMaxRows((int) Math.min(Integer.MAX_VALUE, (long) maxRows + 1L));
if (fetchSize > 0) {
try {
statement.setFetchSize(fetchSize);
} catch (SQLFeatureNotSupportedException ignored) {
}
}
if (timeoutSecs >= 0) {
try {
statement.setQueryTimeout(timeoutSecs);
} catch (SQLFeatureNotSupportedException ignored) {
}
}
}
private static String trimStatementSql(String sql) {
return sql == null ? "" : sql.trim().replaceFirst(";\\s*$", "");
}
@@ -927,6 +961,21 @@ public final class DbxJdbcPlugin {
return text.isEmpty() ? null : text;
}
private static int positiveInt(JsonNode node, String field, int defaultValue) {
return Math.max(1, nonNegativeInt(node, field, defaultValue));
}
private static int nonNegativeInt(JsonNode node, String field, int defaultValue) {
JsonNode value = node.path(field);
if (value.isMissingNode() || value.isNull()) {
return defaultValue;
}
if (!value.canConvertToInt()) {
return defaultValue;
}
return Math.max(0, value.asInt(defaultValue));
}
private static String emptyToNull(String value) {
return value == null || value.isBlank() ? null : value;
}
plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java
+30
View File
@@ -75,6 +75,36 @@ final class DbxJdbcPluginTest {
assertEquals("0x0001abff", response.path("result").path("rows").path(0).path(0).asText());
}
@Test
void executeQueryPreservesChineseTextValues() throws Exception {
JsonNode response = request("executeQuery", """
{
"connection": %s,
"sql": "SELECT '中文测试' AS label"
}
""".formatted(CONNECTION));
assertFalse(response.has("error"), response.toString());
assertEquals("中文测试", response.path("result").path("rows").path(0).path(0).asText());
}
@Test
void executeQueryHonorsMaxRowsAndAcceptsExecutionOptions() throws Exception {
JsonNode response = request("executeQuery", """
{
"connection": %s,
"sql": "SELECT * FROM (VALUES (1), (2)) AS t(n)",
"maxRows": 1,
"fetchSize": 1,
"timeoutSecs": 60
}
""".formatted(CONNECTION));
assertFalse(response.has("error"), response.toString());
assertEquals(1, response.path("result").path("rows").size());
assertEquals(true, response.path("result").path("truncated").asBoolean());
}
@Test
void driverQuirksDetectYashanJdbcUrl() throws Exception {
JsonNode yashan = MAPPER.readTree("""