diff --git a/crates/dbx-core/src/db/agent_driver.rs b/crates/dbx-core/src/db/agent_driver.rs index bf03982835..1ee10bc8c5 100644 --- a/crates/dbx-core/src/db/agent_driver.rs +++ b/crates/dbx-core/src/db/agent_driver.rs @@ -286,6 +286,17 @@ impl AgentDriverClient { &mut self, method: &str, params: Value, + ) -> Result { + 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( + &mut self, + method: &str, + params: Value, + timeout_duration: Option, ) -> Result { 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::(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::(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( + &mut self, + method: AgentMethod, + params: Value, + timeout_duration: Option, + ) -> Result { + self.call_with_timeout(method.as_str(), params, timeout_duration).await + } + pub async fn connect(&mut self, params: Value) -> Result { 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( + &mut self, + params: Value, + timeout_duration: Option, + ) -> Result { + self.call_method_with_timeout(AgentMethod::ExecuteQuery, params, timeout_duration).await + } + pub async fn execute_query_page( &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( + &mut self, + params: Value, + timeout_duration: Option, + ) -> Result { + self.call_method_with_timeout(AgentMethod::ExecuteQueryPage, params, timeout_duration).await + } + pub async fn fetch_query_page(&mut self, params: Value) -> Result { self.call_method(AgentMethod::FetchQueryPage, params).await } + pub async fn fetch_query_page_with_timeout( + &mut self, + params: Value, + timeout_duration: Option, + ) -> Result { + self.call_method_with_timeout(AgentMethod::FetchQueryPage, params, timeout_duration).await + } + pub async fn close_query_session( &mut self, session_id: &str, diff --git a/crates/dbx-core/src/plugins.rs b/crates/dbx-core/src/plugins.rs index 95e4d1f71b..dd5aab0e1b 100644 --- a/crates/dbx-core/src/plugins.rs +++ b/crates/dbx-core/src/plugins.rs @@ -242,26 +242,35 @@ impl PluginDriverSession { } pub async fn invoke(&self, method: &str, params: serde_json::Value) -> Result + where + T: DeserializeOwned, + { + self.invoke_with_timeout(method, params, Some(PLUGIN_REQUEST_TIMEOUT)).await + } + + pub async fn invoke_with_timeout( + &self, + method: &str, + params: serde_json::Value, + timeout_duration: Option, + ) -> Result 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, } } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 7f86c6d824..e3254c9d99 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -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::("executeQuery", params).await + let params = + external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options); + session.invoke_with_timeout::("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); } diff --git a/plugins/jdbc/bin/dbx-jdbc-plugin b/plugins/jdbc/bin/dbx-jdbc-plugin index a2e1f672c0..112d9815fa 100755 --- a/plugins/jdbc/bin/dbx-jdbc-plugin +++ b/plugins/jdbc/bin/dbx-jdbc-plugin @@ -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" diff --git a/plugins/jdbc/bin/dbx-jdbc-plugin.bat b/plugins/jdbc/bin/dbx-jdbc-plugin.bat index c9672a4663..4a4b98f04e 100644 --- a/plugins/jdbc/bin/dbx-jdbc-plugin.bat +++ b/plugins/jdbc/bin/dbx-jdbc-plugin.bat @@ -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" diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 75b61117a1..6979afbf93 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -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; } diff --git a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java index aa17042079..120a39c3d0 100644 --- a/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java +++ b/plugins/jdbc/src/test/java/app/dbx/jdbc/DbxJdbcPluginTest.java @@ -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("""