fix(jdbc): keep query pagination cursor sessions
This commit is contained in:
1 parent
73f4762b8d
commit
593635a94d
8 files changed
+462
-10
No files matched your search
@@ -1268,7 +1268,7 @@ export const useQueryStore = defineStore("query", () => {
|
|||||||
await connStore.ensureConnected(tab.connectionId);
|
await connStore.ensureConnected(tab.connectionId);
|
||||||
const conn = connStore.getConfig(tab.connectionId);
|
const conn = connStore.getConfig(tab.connectionId);
|
||||||
const effectiveDbType = effectiveDatabaseTypeForConnection(conn);
|
const effectiveDbType = effectiveDatabaseTypeForConnection(conn);
|
||||||
const useAgentCursor = supportsDatabaseFeature(conn?.db_type, "driverManagement");
|
const useAgentCursor = conn?.db_type === "jdbc" || supportsDatabaseFeature(conn?.db_type, "driverManagement");
|
||||||
const queryTimeoutSecs = queryTimeoutSecsForConnection(conn);
|
const queryTimeoutSecs = queryTimeoutSecsForConnection(conn);
|
||||||
const settingsStore = useSettingsStore();
|
const settingsStore = useSettingsStore();
|
||||||
await previousResultSessionClose;
|
await previousResultSessionClose;
|
||||||
@@ -1520,7 +1520,6 @@ export const useQueryStore = defineStore("query", () => {
|
|||||||
...(typeof pageLimit === "number"
|
...(typeof pageLimit === "number"
|
||||||
? useAgentResultSession
|
? useAgentResultSession
|
||||||
? {
|
? {
|
||||||
maxRows: pageLimit,
|
|
||||||
fetchSize: pageLimit,
|
fetchSize: pageLimit,
|
||||||
pageSize: pageLimit,
|
pageSize: pageLimit,
|
||||||
resultSessionId: options?.pagination?.sessionId,
|
resultSessionId: options?.pagination?.sessionId,
|
||||||
@@ -1999,7 +1998,7 @@ export const useQueryStore = defineStore("query", () => {
|
|||||||
const conn = connStore.getConfig(tab.connectionId);
|
const conn = connStore.getConfig(tab.connectionId);
|
||||||
const effectiveDbType = effectiveDatabaseTypeForConnection(conn);
|
const effectiveDbType = effectiveDatabaseTypeForConnection(conn);
|
||||||
const queryTimeoutSecs = queryTimeoutSecsForConnection(conn);
|
const queryTimeoutSecs = queryTimeoutSecsForConnection(conn);
|
||||||
const useAgentCursor = supportsDatabaseFeature(conn?.db_type, "driverManagement");
|
const useAgentCursor = conn?.db_type === "jdbc" || supportsDatabaseFeature(conn?.db_type, "driverManagement");
|
||||||
const queryBaseSql = tab.resultBaseSql ?? sql;
|
const queryBaseSql = tab.resultBaseSql ?? sql;
|
||||||
const pageLimit = Math.max(tab.resultPageLimit ?? 0, TABLE_DATA_EXPORT_PAGE_SIZE);
|
const pageLimit = Math.max(tab.resultPageLimit ?? 0, TABLE_DATA_EXPORT_PAGE_SIZE);
|
||||||
const rows: QueryResult["rows"] = [];
|
const rows: QueryResult["rows"] = [];
|
||||||
@@ -2021,7 +2020,6 @@ export const useQueryStore = defineStore("query", () => {
|
|||||||
if (typeof plan.pageLimit !== "number" || typeof plan.pageOffset !== "number") return tab.result;
|
if (typeof plan.pageLimit !== "number" || typeof plan.pageOffset !== "number") return tab.result;
|
||||||
const executionOptions = plan.useAgentResultSession
|
const executionOptions = plan.useAgentResultSession
|
||||||
? {
|
? {
|
||||||
maxRows: plan.pageLimit,
|
|
||||||
fetchSize: plan.pageLimit,
|
fetchSize: plan.pageLimit,
|
||||||
pageSize: plan.pageLimit,
|
pageSize: plan.pageLimit,
|
||||||
resultSessionId: sessionId,
|
resultSessionId: sessionId,
|
||||||
|
|||||||
@@ -1223,6 +1223,11 @@ impl AppState {
|
|||||||
close_removed_pools_in_background(removed);
|
close_removed_pools_in_background(removed);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn remove_external_driver_pools(&self, driver_id: &str) {
|
||||||
|
let removed = self.drain_external_driver_pools(driver_id).await;
|
||||||
|
close_removed_pools(removed).await;
|
||||||
|
}
|
||||||
|
|
||||||
async fn drain_connection_pools(&self, connection_id: &str) -> Vec<(String, PoolKind)> {
|
async fn drain_connection_pools(&self, connection_id: &str) -> Vec<(String, PoolKind)> {
|
||||||
let pool_prefix = format!("{connection_id}:");
|
let pool_prefix = format!("{connection_id}:");
|
||||||
let keys_to_remove: Vec<String> = self
|
let keys_to_remove: Vec<String> = self
|
||||||
@@ -1248,6 +1253,30 @@ impl AppState {
|
|||||||
removed
|
removed
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn drain_external_driver_pools(&self, driver_id: &str) -> Vec<(String, PoolKind)> {
|
||||||
|
let keys_to_remove: Vec<String> = self
|
||||||
|
.connections
|
||||||
|
.read()
|
||||||
|
.await
|
||||||
|
.iter()
|
||||||
|
.filter_map(|(key, pool)| match pool {
|
||||||
|
PoolKind::ExternalDriver { driver_id: pool_driver_id, .. } if pool_driver_id == driver_id => {
|
||||||
|
Some(key.clone())
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||||||
|
let mut conns = self.connections.write().await;
|
||||||
|
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||||||
|
for key in keys_to_remove {
|
||||||
|
if let Some(pool) = conns.remove(&key) {
|
||||||
|
removed.push((key, pool));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
removed
|
||||||
|
}
|
||||||
|
|
||||||
async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
|
async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
|
||||||
let configs = self.configs.read().await;
|
let configs = self.configs.read().await;
|
||||||
configs.get(connection_id).is_some_and(|config| config.has_effective_transport_layers())
|
configs.get(connection_id).is_some_and(|config| config.has_effective_transport_layers())
|
||||||
@@ -1429,7 +1458,9 @@ pub async fn close_pool_kind(pool: PoolKind) {
|
|||||||
let _ = client.disconnect().await;
|
let _ = client.disconnect().await;
|
||||||
}
|
}
|
||||||
PoolKind::ExternalTabular(_) => {}
|
PoolKind::ExternalTabular(_) => {}
|
||||||
PoolKind::ExternalDriver { .. } => {}
|
PoolKind::ExternalDriver { session, .. } => {
|
||||||
|
session.shutdown().await;
|
||||||
|
}
|
||||||
PoolKind::MessageQueue => {}
|
PoolKind::MessageQueue => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -330,6 +330,10 @@ impl PluginDriverSession {
|
|||||||
let _ = process.child.kill().await;
|
let _ = process.child.kill().await;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn shutdown(&self) {
|
||||||
|
self.kill().await;
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn pid(&self) -> Option<u32> {
|
pub async fn pid(&self) -> Option<u32> {
|
||||||
let process = self.process.lock().await;
|
let process = self.process.lock().await;
|
||||||
process.child.id()
|
process.child.id()
|
||||||
@@ -509,7 +513,10 @@ fn resolve_plugin_executable(plugin_dir: &Path, executable: &str) -> PathBuf {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{read_decoded_plugin_response, read_plugin_line, InstalledPlugin, PluginManifest};
|
use super::{
|
||||||
|
read_decoded_plugin_response, read_plugin_line, InstalledPlugin, PluginDriverManifest, PluginDriverSession,
|
||||||
|
PluginManifest, PluginRuntimeEnv,
|
||||||
|
};
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
use tokio::io::BufReader;
|
use tokio::io::BufReader;
|
||||||
|
|
||||||
@@ -547,4 +554,59 @@ mod tests {
|
|||||||
|
|
||||||
assert_eq!(result["ok"], true);
|
assert_eq!(result["ok"], true);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
#[tokio::test]
|
||||||
|
async fn shutdown_kills_plugin_child_process() {
|
||||||
|
let dir = std::env::temp_dir().join(format!("dbx-plugin-shutdown-test-{}", uuid::Uuid::new_v4()));
|
||||||
|
std::fs::create_dir_all(&dir).unwrap();
|
||||||
|
let executable = dir.join("plugin.sh");
|
||||||
|
std::fs::write(&executable, "#!/bin/sh\nsleep 30\n").unwrap();
|
||||||
|
{
|
||||||
|
use std::os::unix::fs::PermissionsExt;
|
||||||
|
|
||||||
|
let mut permissions = std::fs::metadata(&executable).unwrap().permissions();
|
||||||
|
permissions.set_mode(0o755);
|
||||||
|
std::fs::set_permissions(&executable, permissions).unwrap();
|
||||||
|
}
|
||||||
|
let plugin = InstalledPlugin {
|
||||||
|
manifest: PluginManifest {
|
||||||
|
id: "jdbc".to_string(),
|
||||||
|
name: "JDBC".to_string(),
|
||||||
|
version: "test".to_string(),
|
||||||
|
protocol_version: 1,
|
||||||
|
description: String::new(),
|
||||||
|
executable: Some("plugin.sh".to_string()),
|
||||||
|
drivers: vec![PluginDriverManifest {
|
||||||
|
id: "jdbc".to_string(),
|
||||||
|
label: "JDBC".to_string(),
|
||||||
|
kind: "external".to_string(),
|
||||||
|
database_type: Some("jdbc".to_string()),
|
||||||
|
}],
|
||||||
|
},
|
||||||
|
path: dir.clone(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let session = PluginDriverSession::start(plugin, "jdbc".to_string(), PluginRuntimeEnv::default())
|
||||||
|
.await
|
||||||
|
.expect("session should start");
|
||||||
|
let pid = session.pid().await.expect("child should have a pid");
|
||||||
|
|
||||||
|
session.shutdown().await;
|
||||||
|
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||||
|
|
||||||
|
assert!(!process_exists(pid));
|
||||||
|
let _ = std::fs::remove_dir_all(dir);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(unix)]
|
||||||
|
fn process_exists(pid: u32) -> bool {
|
||||||
|
std::process::Command::new("kill")
|
||||||
|
.arg("-0")
|
||||||
|
.arg(pid.to_string())
|
||||||
|
.stderr(std::process::Stdio::null())
|
||||||
|
.status()
|
||||||
|
.map(|status| status.success())
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
@@ -1033,9 +1033,22 @@ pub async fn do_execute(
|
|||||||
let plugin_timeout = query_timeout;
|
let plugin_timeout = query_timeout;
|
||||||
drop(connections);
|
drop(connections);
|
||||||
wait_for_query_opt(cancel_token, query_timeout, async move {
|
wait_for_query_opt(cancel_token, query_timeout, async move {
|
||||||
let params =
|
if let Some(session_id) = options.result_session_id.as_deref() {
|
||||||
external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options);
|
let params = external_driver_fetch_query_page_params(
|
||||||
session.invoke_with_timeout::<db::QueryResult>("executeQuery", params, plugin_timeout).await
|
config.as_ref(),
|
||||||
|
session_id,
|
||||||
|
options.page_size.unwrap_or(MAX_ROWS),
|
||||||
|
);
|
||||||
|
session.invoke_with_timeout::<db::QueryResult>("fetchQueryPage", params, plugin_timeout).await
|
||||||
|
} else if options.page_size.is_some() {
|
||||||
|
let params =
|
||||||
|
external_driver_query_params(config.as_ref(), &sql, &database, schema.as_deref(), &options);
|
||||||
|
session.invoke_with_timeout::<db::QueryResult>("executeQueryPage", params, plugin_timeout).await
|
||||||
|
} else {
|
||||||
|
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
|
.await
|
||||||
.map(|result| normalize_query_result_for_js(truncate_result_with_max_rows(result, max_rows)))
|
.map(|result| normalize_query_result_for_js(truncate_result_with_max_rows(result, max_rows)))
|
||||||
@@ -1063,9 +1076,24 @@ fn external_driver_query_params(
|
|||||||
if let Some(timeout_secs) = options.timeout_secs {
|
if let Some(timeout_secs) = options.timeout_secs {
|
||||||
params["timeoutSecs"] = serde_json::json!(timeout_secs);
|
params["timeoutSecs"] = serde_json::json!(timeout_secs);
|
||||||
}
|
}
|
||||||
|
if let Some(page_size) = options.page_size {
|
||||||
|
params["pageSize"] = serde_json::json!(page_size);
|
||||||
|
}
|
||||||
params
|
params
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn external_driver_fetch_query_page_params(
|
||||||
|
config: &crate::models::connection::ConnectionConfig,
|
||||||
|
session_id: &str,
|
||||||
|
page_size: usize,
|
||||||
|
) -> serde_json::Value {
|
||||||
|
serde_json::json!({
|
||||||
|
"connection": config,
|
||||||
|
"sessionId": session_id,
|
||||||
|
"pageSize": page_size,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn execute_sql_statement(
|
pub async fn execute_sql_statement(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
connection_id: &str,
|
connection_id: &str,
|
||||||
@@ -1242,6 +1270,16 @@ pub async fn close_query_session(
|
|||||||
let mut client = client.lock().await;
|
let mut client = client.lock().await;
|
||||||
client.close_query_session(session_id).await
|
client.close_query_session(session_id).await
|
||||||
}
|
}
|
||||||
|
PoolKind::ExternalDriver { config, session, .. } => {
|
||||||
|
let config = config.clone();
|
||||||
|
let session = session.clone();
|
||||||
|
drop(connections);
|
||||||
|
let params = external_driver_fetch_query_page_params(config.as_ref(), session_id, 1);
|
||||||
|
session
|
||||||
|
.invoke::<serde_json::Value>("closeQuerySession", params)
|
||||||
|
.await
|
||||||
|
.map(|value| value.get("ok").and_then(|ok| ok.as_bool()).unwrap_or(false))
|
||||||
|
}
|
||||||
_ => Ok(false),
|
_ => Ok(false),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1320,6 +1320,80 @@ test("query result export fetches every paginated page", async () => {
|
|||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
test("jdbc query pagination uses result sessions without capping max rows to one page", async () => {
|
||||||
|
const restoreStorage = installMemoryStorage();
|
||||||
|
setActivePinia(createPinia());
|
||||||
|
const connectionStore = useConnectionStore();
|
||||||
|
const store = useQueryStore();
|
||||||
|
const originalFetch = globalThis.fetch;
|
||||||
|
let prepareBody: any;
|
||||||
|
let executeBody: any;
|
||||||
|
|
||||||
|
connectionStore.addEphemeralConnection({
|
||||||
|
...conn("jdbc-1"),
|
||||||
|
db_type: "jdbc",
|
||||||
|
connection_string: "jdbc:Cache://127.0.0.1:1972/USER",
|
||||||
|
jdbc_driver_class: "com.intersys.jdbc.CacheDriver",
|
||||||
|
});
|
||||||
|
const tabId = store.createTab("jdbc-1", "", "Query", "query", "SQLUser");
|
||||||
|
const tab = store.tabs.find((item) => item.id === tabId);
|
||||||
|
assert.ok(tab);
|
||||||
|
|
||||||
|
globalThis.fetch = (async (input, init) => {
|
||||||
|
const url = String(input);
|
||||||
|
if (url === "/api/query/prepare-pagination-plan") {
|
||||||
|
prepareBody = JSON.parse(String(init?.body ?? "{}"));
|
||||||
|
return new Response(
|
||||||
|
JSON.stringify({
|
||||||
|
sqlToExecute: "SELECT * FROM CT_Loc",
|
||||||
|
pageLimit: 100,
|
||||||
|
pageOffset: 0,
|
||||||
|
useAgentResultSession: true,
|
||||||
|
}),
|
||||||
|
{ status: 200, headers: { "Content-Type": "application/json" } },
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (url === "/api/query/execute-multi") {
|
||||||
|
executeBody = JSON.parse(String(init?.body ?? "{}"));
|
||||||
|
return new Response(
|
||||||
|
JSON.stringify([
|
||||||
|
{
|
||||||
|
columns: ["id"],
|
||||||
|
rows: Array.from({ length: 100 }, (_, index) => [index + 1]),
|
||||||
|
affected_rows: 0,
|
||||||
|
execution_time_ms: 1,
|
||||||
|
session_id: "session-1",
|
||||||
|
has_more: true,
|
||||||
|
},
|
||||||
|
]),
|
||||||
|
{ status: 200, headers: { "Content-Type": "application/json" } },
|
||||||
|
);
|
||||||
|
}
|
||||||
|
if (url === "/api/query/analyze-editability") {
|
||||||
|
return new Response(JSON.stringify({ editable: false, reason: "complex-source" }), {
|
||||||
|
status: 200,
|
||||||
|
headers: { "Content-Type": "application/json" },
|
||||||
|
});
|
||||||
|
}
|
||||||
|
return new Response("unexpected request", { status: 500 });
|
||||||
|
}) as typeof fetch;
|
||||||
|
|
||||||
|
try {
|
||||||
|
await store.executeTabSql(tabId, "SELECT * FROM CT_Loc");
|
||||||
|
|
||||||
|
assert.equal(prepareBody.options.useAgentCursor, true);
|
||||||
|
assert.equal(executeBody.pageSize, 100);
|
||||||
|
assert.equal(executeBody.fetchSize, 100);
|
||||||
|
assert.equal(executeBody.maxRows, undefined);
|
||||||
|
assert.equal(executeBody.clientSessionId, tabId);
|
||||||
|
assert.equal(tab.resultSessionId, "session-1");
|
||||||
|
assert.equal(tab.result?.has_more, true);
|
||||||
|
} finally {
|
||||||
|
globalThis.fetch = originalFetch;
|
||||||
|
restoreStorage();
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
test("table data export fetches every filtered page", async () => {
|
test("table data export fetches every filtered page", async () => {
|
||||||
const restoreStorage = installMemoryStorage();
|
const restoreStorage = installMemoryStorage();
|
||||||
setActivePinia(createPinia());
|
setActivePinia(createPinia());
|
||||||
|
|||||||
@@ -34,12 +34,15 @@ import java.sql.Types;
|
|||||||
import java.time.temporal.TemporalAccessor;
|
import java.time.temporal.TemporalAccessor;
|
||||||
import java.util.ArrayList;
|
import java.util.ArrayList;
|
||||||
import java.util.Collections;
|
import java.util.Collections;
|
||||||
|
import java.util.HashMap;
|
||||||
import java.util.HashSet;
|
import java.util.HashSet;
|
||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.Locale;
|
import java.util.Locale;
|
||||||
|
import java.util.Map;
|
||||||
import java.util.Properties;
|
import java.util.Properties;
|
||||||
import java.util.ServiceLoader;
|
import java.util.ServiceLoader;
|
||||||
import java.util.Set;
|
import java.util.Set;
|
||||||
|
import java.util.UUID;
|
||||||
import java.util.logging.Logger;
|
import java.util.logging.Logger;
|
||||||
|
|
||||||
public final class DbxJdbcPlugin {
|
public final class DbxJdbcPlugin {
|
||||||
@@ -109,6 +112,7 @@ public final class DbxJdbcPlugin {
|
|||||||
private static String registeredDriverKey = "";
|
private static String registeredDriverKey = "";
|
||||||
private static String sharedConnectionKey = "";
|
private static String sharedConnectionKey = "";
|
||||||
private static Connection sharedConnection;
|
private static Connection sharedConnection;
|
||||||
|
private static final Map<String, QuerySession> QUERY_SESSIONS = new HashMap<>();
|
||||||
|
|
||||||
record JdbcDriverQuirks(
|
record JdbcDriverQuirks(
|
||||||
boolean skipExecutionContext,
|
boolean skipExecutionContext,
|
||||||
@@ -235,6 +239,21 @@ public final class DbxJdbcPlugin {
|
|||||||
nonNegativeInt(params, "fetchSize", 0),
|
nonNegativeInt(params, "fetchSize", 0),
|
||||||
nonNegativeInt(params, "timeoutSecs", -1)
|
nonNegativeInt(params, "timeoutSecs", -1)
|
||||||
);
|
);
|
||||||
|
case "executeQueryPage", "execute_query_page" -> executeQueryPage(
|
||||||
|
connection,
|
||||||
|
requireText(params, "sql"),
|
||||||
|
optionalText(params, "database"),
|
||||||
|
optionalText(params, "schema"),
|
||||||
|
positiveInt(params, "pageSize", 100),
|
||||||
|
positiveInt(params, "maxRows", MAX_ROWS),
|
||||||
|
nonNegativeInt(params, "fetchSize", 0),
|
||||||
|
nonNegativeInt(params, "timeoutSecs", -1)
|
||||||
|
);
|
||||||
|
case "fetchQueryPage", "fetch_query_page" -> fetchQueryPage(
|
||||||
|
requireText(params, "sessionId"),
|
||||||
|
positiveInt(params, "pageSize", 100)
|
||||||
|
);
|
||||||
|
case "closeQuerySession", "close_query_session" -> closeQuerySessionResult(requireText(params, "sessionId"));
|
||||||
case "listDatabases" -> listDatabases(connection);
|
case "listDatabases" -> listDatabases(connection);
|
||||||
case "listSchemas" -> listSchemas(connection, optionalText(params, "database"));
|
case "listSchemas" -> listSchemas(connection, optionalText(params, "database"));
|
||||||
case "listTables" -> listTables(connection, optionalText(params, "database"), optionalText(params, "schema"));
|
case "listTables" -> listTables(connection, optionalText(params, "database"), optionalText(params, "schema"));
|
||||||
@@ -273,6 +292,7 @@ public final class DbxJdbcPlugin {
|
|||||||
if (driverKey.equals(registeredDriverKey)) {
|
if (driverKey.equals(registeredDriverKey)) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
closeSharedConnection();
|
||||||
List<URL> urls = new ArrayList<>();
|
List<URL> urls = new ArrayList<>();
|
||||||
JsonNode paths = connection.path("jdbc_driver_paths");
|
JsonNode paths = connection.path("jdbc_driver_paths");
|
||||||
if (paths.isArray()) {
|
if (paths.isArray()) {
|
||||||
@@ -363,7 +383,7 @@ public final class DbxJdbcPlugin {
|
|||||||
int maxRows,
|
int maxRows,
|
||||||
int fetchSize,
|
int fetchSize,
|
||||||
int timeoutSecs
|
int timeoutSecs
|
||||||
) throws SQLException {
|
) throws Exception {
|
||||||
long start = System.nanoTime();
|
long start = System.nanoTime();
|
||||||
Connection conn = openConnection(connection);
|
Connection conn = openConnection(connection);
|
||||||
applyExecutionContext(connection, conn, database, schema);
|
applyExecutionContext(connection, conn, database, schema);
|
||||||
@@ -411,6 +431,187 @@ public final class DbxJdbcPlugin {
|
|||||||
private record ExecutedStatement(ResultSet resultSet, int updateCount) {
|
private record ExecutedStatement(ResultSet resultSet, int updateCount) {
|
||||||
}
|
}
|
||||||
|
|
||||||
|
private static final class QuerySession {
|
||||||
|
private final String id;
|
||||||
|
private final Statement statement;
|
||||||
|
private final ResultSet resultSet;
|
||||||
|
private final ResultSetMetaData meta;
|
||||||
|
private final ArrayNode columns;
|
||||||
|
private final int maxRows;
|
||||||
|
private final long startNanos;
|
||||||
|
private int rowsReturned;
|
||||||
|
private ArrayNode pendingRow;
|
||||||
|
|
||||||
|
private QuerySession(
|
||||||
|
String id,
|
||||||
|
Statement statement,
|
||||||
|
ResultSet resultSet,
|
||||||
|
ResultSetMetaData meta,
|
||||||
|
ArrayNode columns,
|
||||||
|
int maxRows,
|
||||||
|
long startNanos
|
||||||
|
) {
|
||||||
|
this.id = id;
|
||||||
|
this.statement = statement;
|
||||||
|
this.resultSet = resultSet;
|
||||||
|
this.meta = meta;
|
||||||
|
this.columns = columns;
|
||||||
|
this.maxRows = Math.max(1, maxRows);
|
||||||
|
this.startNanos = startNanos;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static JsonNode executeQueryPage(
|
||||||
|
JsonNode connection,
|
||||||
|
String sql,
|
||||||
|
String database,
|
||||||
|
String schema,
|
||||||
|
int pageSize,
|
||||||
|
int maxRows,
|
||||||
|
int fetchSize,
|
||||||
|
int timeoutSecs
|
||||||
|
) throws Exception {
|
||||||
|
long start = System.nanoTime();
|
||||||
|
Connection conn = openConnection(connection);
|
||||||
|
applyExecutionContext(connection, conn, database, schema);
|
||||||
|
JdbcDriverQuirks quirks = driverQuirks(connection);
|
||||||
|
Statement statement = conn.createStatement();
|
||||||
|
try {
|
||||||
|
applyStatementOptions(statement, maxRows, fetchSize, timeoutSecs, quirks);
|
||||||
|
String trimmedSql = trimStatementSql(sql);
|
||||||
|
ExecutedStatement executed = executeStatementForResult(statement, trimmedSql, quirks);
|
||||||
|
ResultSet rs = executed.resultSet();
|
||||||
|
if (rs == null) {
|
||||||
|
ObjectNode result = MAPPER.createObjectNode();
|
||||||
|
result.set("columns", MAPPER.createArrayNode());
|
||||||
|
result.set("rows", MAPPER.createArrayNode());
|
||||||
|
result.put("affected_rows", Math.max(executed.updateCount(), 0));
|
||||||
|
result.put("execution_time_ms", (System.nanoTime() - start) / 1_000_000);
|
||||||
|
result.put("truncated", false);
|
||||||
|
result.putNull("session_id");
|
||||||
|
result.put("has_more", false);
|
||||||
|
statement.close();
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
ResultSetMetaData meta = rs.getMetaData();
|
||||||
|
ArrayNode columns = MAPPER.createArrayNode();
|
||||||
|
int columnCount = meta.getColumnCount();
|
||||||
|
for (int i = 1; i <= columnCount; i++) {
|
||||||
|
String label = meta.getColumnLabel(i);
|
||||||
|
columns.add(label == null || label.isBlank() ? meta.getColumnName(i) : label);
|
||||||
|
}
|
||||||
|
String sessionId = UUID.randomUUID().toString();
|
||||||
|
QuerySession session = new QuerySession(sessionId, statement, rs, meta, columns, maxRows, start);
|
||||||
|
QUERY_SESSIONS.put(sessionId, session);
|
||||||
|
return readQuerySessionPage(session, pageSize);
|
||||||
|
} catch (Exception error) {
|
||||||
|
try {
|
||||||
|
statement.close();
|
||||||
|
} catch (Exception ignored) {
|
||||||
|
}
|
||||||
|
throw error;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static JsonNode fetchQueryPage(String sessionId, int pageSize) throws SQLException {
|
||||||
|
QuerySession session = QUERY_SESSIONS.get(sessionId);
|
||||||
|
if (session == null) {
|
||||||
|
throw new IllegalArgumentException("Unknown query session: " + sessionId);
|
||||||
|
}
|
||||||
|
return readQuerySessionPage(session, pageSize);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static JsonNode readQuerySessionPage(QuerySession session, int pageSize) throws SQLException {
|
||||||
|
int effectivePageSize = Math.max(1, pageSize);
|
||||||
|
ArrayNode rows = MAPPER.createArrayNode();
|
||||||
|
boolean truncated = false;
|
||||||
|
|
||||||
|
while (rows.size() < effectivePageSize && session.rowsReturned < session.maxRows) {
|
||||||
|
ArrayNode row;
|
||||||
|
if (session.pendingRow != null) {
|
||||||
|
row = session.pendingRow;
|
||||||
|
session.pendingRow = null;
|
||||||
|
} else {
|
||||||
|
if (!session.resultSet.next()) {
|
||||||
|
closeQuerySession(session.id);
|
||||||
|
return queryPageResult(session, rows, false, false);
|
||||||
|
}
|
||||||
|
row = readRow(session.resultSet, session.meta);
|
||||||
|
}
|
||||||
|
rows.add(row);
|
||||||
|
session.rowsReturned++;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (session.rowsReturned >= session.maxRows) {
|
||||||
|
truncated = session.pendingRow != null || session.resultSet.next();
|
||||||
|
closeQuerySession(session.id);
|
||||||
|
return queryPageResult(session, rows, truncated, false);
|
||||||
|
}
|
||||||
|
|
||||||
|
boolean hasMore = session.resultSet.next();
|
||||||
|
if (!hasMore) {
|
||||||
|
closeQuerySession(session.id);
|
||||||
|
return queryPageResult(session, rows, false, false);
|
||||||
|
}
|
||||||
|
|
||||||
|
session.pendingRow = readRow(session.resultSet, session.meta);
|
||||||
|
return queryPageResult(session, rows, false, true);
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ObjectNode queryPageResult(QuerySession session, ArrayNode rows, boolean truncated, boolean hasMore) {
|
||||||
|
ObjectNode result = MAPPER.createObjectNode();
|
||||||
|
result.set("columns", session.columns.deepCopy());
|
||||||
|
result.set("rows", rows);
|
||||||
|
result.put("affected_rows", 0);
|
||||||
|
result.put("execution_time_ms", (System.nanoTime() - session.startNanos) / 1_000_000);
|
||||||
|
result.put("truncated", truncated);
|
||||||
|
if (hasMore) {
|
||||||
|
result.put("session_id", session.id);
|
||||||
|
} else {
|
||||||
|
result.putNull("session_id");
|
||||||
|
}
|
||||||
|
result.put("has_more", hasMore);
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ObjectNode closeQuerySessionResult(String sessionId) {
|
||||||
|
ObjectNode result = MAPPER.createObjectNode();
|
||||||
|
result.put("ok", closeQuerySession(sessionId));
|
||||||
|
return result;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static boolean closeQuerySession(String sessionId) {
|
||||||
|
QuerySession session = QUERY_SESSIONS.remove(sessionId);
|
||||||
|
if (session == null) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
session.resultSet.close();
|
||||||
|
} catch (Exception ignored) {
|
||||||
|
}
|
||||||
|
try {
|
||||||
|
session.statement.close();
|
||||||
|
} catch (Exception ignored) {
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
private static void closeAllQuerySessions() {
|
||||||
|
List<String> sessionIds = new ArrayList<>(QUERY_SESSIONS.keySet());
|
||||||
|
for (String sessionId : sessionIds) {
|
||||||
|
closeQuerySession(sessionId);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
private static ArrayNode readRow(ResultSet rs, ResultSetMetaData meta) throws SQLException {
|
||||||
|
ArrayNode row = MAPPER.createArrayNode();
|
||||||
|
for (int i = 1; i <= meta.getColumnCount(); i++) {
|
||||||
|
row.add(MAPPER.valueToTree(readValue(rs, meta, i)));
|
||||||
|
}
|
||||||
|
return row;
|
||||||
|
}
|
||||||
|
|
||||||
private static ExecutedStatement executeStatementForResult(
|
private static ExecutedStatement executeStatementForResult(
|
||||||
Statement statement,
|
Statement statement,
|
||||||
String sql,
|
String sql,
|
||||||
@@ -1194,6 +1395,7 @@ public final class DbxJdbcPlugin {
|
|||||||
}
|
}
|
||||||
|
|
||||||
private static void closeSharedConnection() {
|
private static void closeSharedConnection() {
|
||||||
|
closeAllQuerySessions();
|
||||||
if (sharedConnection != null) {
|
if (sharedConnection != null) {
|
||||||
try {
|
try {
|
||||||
sharedConnection.close();
|
sharedConnection.close();
|
||||||
|
|||||||
@@ -105,6 +105,50 @@ final class DbxJdbcPluginTest {
|
|||||||
assertEquals("中文测试", response.path("result").path("rows").path(0).path(0).asText());
|
assertEquals("中文测试", response.path("result").path("rows").path(0).path(0).asText());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Test
|
||||||
|
void executeQueryPageKeepsCursorForNextPages() throws Exception {
|
||||||
|
JsonNode first = request("executeQueryPage", """
|
||||||
|
{
|
||||||
|
"connection": %s,
|
||||||
|
"sql": "SELECT X FROM SYSTEM_RANGE(1, 5)",
|
||||||
|
"pageSize": 2,
|
||||||
|
"maxRows": 10
|
||||||
|
}
|
||||||
|
""".formatted(CONNECTION));
|
||||||
|
|
||||||
|
assertFalse(first.has("error"), first.toString());
|
||||||
|
assertEquals(1, first.path("result").path("rows").path(0).path(0).asInt());
|
||||||
|
assertEquals(2, first.path("result").path("rows").path(1).path(0).asInt());
|
||||||
|
assertEquals(true, first.path("result").path("has_more").asBoolean());
|
||||||
|
String sessionId = first.path("result").path("session_id").asText();
|
||||||
|
|
||||||
|
JsonNode second = request("fetchQueryPage", """
|
||||||
|
{
|
||||||
|
"connection": %s,
|
||||||
|
"sessionId": "%s",
|
||||||
|
"pageSize": 2
|
||||||
|
}
|
||||||
|
""".formatted(CONNECTION, sessionId));
|
||||||
|
|
||||||
|
assertFalse(second.has("error"), second.toString());
|
||||||
|
assertEquals(3, second.path("result").path("rows").path(0).path(0).asInt());
|
||||||
|
assertEquals(4, second.path("result").path("rows").path(1).path(0).asInt());
|
||||||
|
assertEquals(true, second.path("result").path("has_more").asBoolean());
|
||||||
|
|
||||||
|
JsonNode third = request("fetch_query_page", """
|
||||||
|
{
|
||||||
|
"connection": %s,
|
||||||
|
"sessionId": "%s",
|
||||||
|
"pageSize": 2
|
||||||
|
}
|
||||||
|
""".formatted(CONNECTION, second.path("result").path("session_id").asText()));
|
||||||
|
|
||||||
|
assertFalse(third.has("error"), third.toString());
|
||||||
|
assertEquals(5, third.path("result").path("rows").path(0).path(0).asInt());
|
||||||
|
assertEquals(false, third.path("result").path("has_more").asBoolean());
|
||||||
|
assertEquals(true, third.path("result").path("session_id").isNull());
|
||||||
|
}
|
||||||
|
|
||||||
@Test
|
@Test
|
||||||
void readValueFormatsDateColumnsWithoutMidnightTime() throws Exception {
|
void readValueFormatsDateColumnsWithoutMidnightTime() throws Exception {
|
||||||
Method method = DbxJdbcPlugin.class.getDeclaredMethod("readValue", ResultSet.class, ResultSetMetaData.class, int.class);
|
Method method = DbxJdbcPlugin.class.getDeclaredMethod("readValue", ResultSet.class, ResultSetMetaData.class, int.class);
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ pub async fn install_jdbc_plugin(
|
|||||||
state: State<'_, Arc<AppState>>,
|
state: State<'_, Arc<AppState>>,
|
||||||
) -> Result<JdbcPluginStatus, String> {
|
) -> Result<JdbcPluginStatus, String> {
|
||||||
let app_handle = app.clone();
|
let app_handle = app.clone();
|
||||||
|
state.remove_external_driver_pools("jdbc").await;
|
||||||
jdbc::install_jdbc_plugin_with_progress(state.plugins.root_dir(), move |event| {
|
jdbc::install_jdbc_plugin_with_progress(state.plugins.root_dir(), move |event| {
|
||||||
emit_agent_progress(&app_handle, event);
|
emit_agent_progress(&app_handle, event);
|
||||||
})
|
})
|
||||||
@@ -37,11 +38,13 @@ pub async fn install_jdbc_plugin_local(
|
|||||||
state: State<'_, Arc<AppState>>,
|
state: State<'_, Arc<AppState>>,
|
||||||
path: String,
|
path: String,
|
||||||
) -> Result<JdbcPluginStatus, String> {
|
) -> Result<JdbcPluginStatus, String> {
|
||||||
|
state.remove_external_driver_pools("jdbc").await;
|
||||||
jdbc::install_jdbc_plugin_from_file(state.plugins.root_dir(), &path).await
|
jdbc::install_jdbc_plugin_from_file(state.plugins.root_dir(), &path).await
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tauri::command]
|
#[tauri::command]
|
||||||
pub async fn uninstall_jdbc_plugin(state: State<'_, Arc<AppState>>) -> Result<JdbcPluginStatus, String> {
|
pub async fn uninstall_jdbc_plugin(state: State<'_, Arc<AppState>>) -> Result<JdbcPluginStatus, String> {
|
||||||
|
state.remove_external_driver_pools("jdbc").await;
|
||||||
let root_dir = state.plugins.root_dir().to_path_buf();
|
let root_dir = state.plugins.root_dir().to_path_buf();
|
||||||
tauri::async_runtime::spawn_blocking(move || jdbc::uninstall_jdbc_plugin(&root_dir))
|
tauri::async_runtime::spawn_blocking(move || jdbc::uninstall_jdbc_plugin(&root_dir))
|
||||||
.await
|
.await
|
||||||
|
|||||||
Reference in new issue
Block a user