fix(jdbc): apply query execution context

This commit is contained in:
t8y2 2026-05-16 10:49:54 +08:00
parent 2024171c8d
commit 3be6f4486a
6 changed files with 165 additions and 10 deletions

View File

@ -248,13 +248,10 @@ pub async fn do_execute(
let session = session.clone();
let sql = sql.to_string();
let schema = schema.map(str::to_string);
let database = config.effective_database().unwrap_or("").to_string();
drop(connections);
wait_for_query(cancel_token, async move {
let params = serde_json::json!({
"connection": config,
"sql": sql,
"schema": schema,
});
let params = external_driver_query_params(&config, &sql, &database, schema.as_deref());
session.invoke::<db::QueryResult>("executeQuery", params).await
})
.await
@ -263,6 +260,20 @@ pub async fn do_execute(
}
}
fn external_driver_query_params(
config: &crate::models::connection::ConnectionConfig,
sql: &str,
database: &str,
schema: Option<&str>,
) -> serde_json::Value {
serde_json::json!({
"connection": config,
"sql": sql,
"database": database,
"schema": schema,
})
}
pub async fn execute_sql_statement(
state: &AppState,
connection_id: &str,
@ -719,6 +730,7 @@ async fn exec_tx_none_inner(
#[cfg(test)]
mod tests {
use super::*;
use crate::models::connection::{ConnectionConfig, DatabaseType, ProxyType};
#[tokio::test]
async fn wait_for_query_returns_cancelled_when_token_is_cancelled() {
@ -789,4 +801,51 @@ mod tests {
assert!(!is_connection_error("syntax error at position 5"));
assert!(!is_connection_error("os error 13"));
}
#[test]
fn external_driver_query_params_include_database_and_schema_context() {
let config = ConnectionConfig {
id: "jdbc-1".to_string(),
name: "JDBC".to_string(),
db_type: DatabaseType::Jdbc,
driver_profile: None,
driver_label: None,
url_params: None,
host: "localhost".to_string(),
port: 0,
username: String::new(),
password: String::new(),
database: None,
visible_databases: None,
color: None,
ssh_enabled: false,
ssh_host: String::new(),
ssh_port: 22,
ssh_user: String::new(),
ssh_password: String::new(),
ssh_key_path: String::new(),
ssh_key_passphrase: String::new(),
ssh_expose_lan: false,
ssh_connect_timeout_secs: 5,
proxy_enabled: false,
proxy_type: ProxyType::Socks5,
proxy_host: String::new(),
proxy_port: 1080,
proxy_username: String::new(),
proxy_password: String::new(),
ssl: false,
sysdba: false,
connection_string: Some("jdbc:h2:mem:test".to_string()),
external_config: None,
jdbc_driver_class: None,
jdbc_driver_paths: Vec::new(),
};
let params = external_driver_query_params(&config, "SELECT * FROM events", "analytics", Some("app"));
assert_eq!(params["connection"]["id"], "jdbc-1");
assert_eq!(params["sql"], "SELECT * FROM events");
assert_eq!(params["database"], "analytics");
assert_eq!(params["schema"], "app");
}
}

View File

@ -1,7 +1,7 @@
{
"id": "jdbc",
"name": "DBX JDBC Plugin",
"version": "0.1.0",
"version": "0.1.1",
"protocol_version": 1,
"description": "Adds optional JDBC driver support to DBX.",
"executable": "bin/dbx-jdbc-plugin",

View File

@ -18,10 +18,27 @@
<artifactId>jackson-databind</artifactId>
<version>2.17.2</version>
</dependency>
<dependency>
<groupId>com.h2database</groupId>
<artifactId>h2</artifactId>
<version>2.2.224</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.junit.jupiter</groupId>
<artifactId>junit-jupiter</artifactId>
<version>5.10.3</version>
<scope>test</scope>
</dependency>
</dependencies>
<build>
<plugins>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-surefire-plugin</artifactId>
<version>3.2.5</version>
</plugin>
<plugin>
<groupId>org.apache.maven.plugins</groupId>
<artifactId>maven-shade-plugin</artifactId>

View File

@ -100,7 +100,12 @@ public final class DbxJdbcPlugin {
result.put("ok", true);
yield result;
}
case "executeQuery" -> executeQuery(connection, requireText(params, "sql"));
case "executeQuery" -> executeQuery(
connection,
requireText(params, "sql"),
optionalText(params, "database"),
optionalText(params, "schema")
);
case "listDatabases" -> listDatabases(connection);
case "listSchemas" -> listSchemas(connection, optionalText(params, "database"));
case "listTables" -> listTables(connection, optionalText(params, "database"), optionalText(params, "schema"));
@ -180,9 +185,10 @@ public final class DbxJdbcPlugin {
return sharedConnection;
}
private static JsonNode executeQuery(JsonNode connection, String sql) throws SQLException {
private static JsonNode executeQuery(JsonNode connection, String sql, String database, String schema) throws SQLException {
long start = System.nanoTime();
Connection conn = openConnection(connection);
applyExecutionContext(conn, database, schema);
try (Statement statement = conn.createStatement()) {
statement.setMaxRows(MAX_ROWS + 1);
boolean hasResultSet = statement.execute(sql);
@ -222,6 +228,21 @@ public final class DbxJdbcPlugin {
}
}
private static void applyExecutionContext(Connection conn, String database, String schema) throws SQLException {
if (database != null) {
try {
conn.setCatalog(database);
} catch (SQLFeatureNotSupportedException | AbstractMethodError ignored) {
}
}
if (schema != null) {
try {
conn.setSchema(schema);
} catch (SQLFeatureNotSupportedException | AbstractMethodError ignored) {
}
}
}
private static JsonNode listDatabases(JsonNode connection) throws SQLException {
ArrayNode result = MAPPER.createArrayNode();
Connection conn = openConnection(connection);

View File

@ -0,0 +1,58 @@
package app.dbx.jdbc;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import java.lang.reflect.Method;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
final class DbxJdbcPluginTest {
private static final ObjectMapper MAPPER = new ObjectMapper();
private static final String CONNECTION = """
{
"connection_string": "jdbc:h2:mem:dbx_ctx;DB_CLOSE_DELAY=-1",
"username": "sa"
}
""";
@AfterEach
void closeConnection() throws Exception {
request("close", """
{ "connection": %s }
""".formatted(CONNECTION));
}
@Test
void executeQueryAppliesSchemaContext() throws Exception {
request("executeQuery", """
{
"connection": %s,
"sql": "CREATE SCHEMA IF NOT EXISTS app"
}
""".formatted(CONNECTION));
JsonNode response = request("executeQuery", """
{
"connection": %s,
"schema": "APP",
"sql": "SELECT SCHEMA() AS schema_name"
}
""".formatted(CONNECTION));
assertFalse(response.has("error"), response.toString());
assertEquals("APP", response.path("result").path("rows").path(0).path(0).asText());
}
private static JsonNode request(String method, String params) throws Exception {
Method handleLine = DbxJdbcPlugin.class.getDeclaredMethod("handleLine", String.class);
handleLine.setAccessible(true);
String line = """
{ "id": 1, "method": "%s", "params": %s }
""".formatted(method, params);
return MAPPER.valueToTree(handleLine.invoke(null, line));
}
}

View File

@ -6,8 +6,8 @@ use serde::Serialize;
use super::connection::AppState;
const JDBC_PLUGIN_DOWNLOAD_URL: &str = "https://github.com/t8y2/dbx/releases/latest/download/dbx-jdbc-plugin-0.1.0.zip";
const JDBC_PLUGIN_R2_PATH: &str = "releases/latest/dbx-jdbc-plugin-0.1.0.zip";
const JDBC_PLUGIN_DOWNLOAD_URL: &str = "https://github.com/t8y2/dbx/releases/latest/download/dbx-jdbc-plugin-0.1.1.zip";
const JDBC_PLUGIN_R2_PATH: &str = "releases/latest/dbx-jdbc-plugin-0.1.1.zip";
#[tauri::command]
pub async fn list_plugins(state: State<'_, Arc<AppState>>) -> Result<Vec<InstalledPlugin>, String> {