xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(agent): track execution context activity

This commit is contained in:
t8y2 2026-06-22 23:17:11 +08:00
parent ee29694776
commit ac3d3a7948
6 changed files with 336 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
apps/desktop/src/stores/queryStore.ts
View File

@ -1211,6 +1211,7 @@ export const useQueryStore = defineStore("query", () => {
}
const countSql = options.countSql;
const clientSessionId = tabClientSessionId({ id: options.tabId }, "count");
const countExecutionId = `${options.executionId}:count`;
if (typeof options.pageLimit === "number" && resultRowCount < options.pageLimit) {
setQueryTotalRowCountIfCurrent(options.tabId, options.executionId, options.result, (options.pageOffset ?? 0) + resultRowCount);
@ -1220,7 +1221,7 @@ export const useQueryStore = defineStore("query", () => {
void (async () => {
try {
console.info("[DBX][executeTabSql:count:start]", { traceId: options.traceId, elapsed: options.elapsed() });
const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, undefined, {
const countResult = await api.executeQuery(options.connectionId, options.database, countSql, options.schema, countExecutionId, {
clientSessionId,
timeoutSecs: options.timeoutSecs,
});
@ -2040,6 +2041,7 @@ export const useQueryStore = defineStore("query", () => {
let executionTimeMs = 0;
let offset = 0;
const clientSessionId = tabClientSessionId(tab, "export");
const exportExecutionId = uuid();
try {
while (true) {
@ -2055,7 +2057,7 @@ export const useQueryStore = defineStore("query", () => {
limit: pageLimit,
offset,
});
const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, undefined, {
const results = await api.executeMulti(tab.connectionId, tab.database, sql, undefined, exportExecutionId, {
maxRows: pageLimit,
fetchSize: pageLimit,
clientSessionId,
@ -2106,6 +2108,7 @@ export const useQueryStore = defineStore("query", () => {
let offset = 0;
let sessionId: string | undefined;
const clientSessionId = tabClientSessionId(tab, "export");
const exportExecutionId = uuid();
try {
while (true) {
@ -2126,7 +2129,7 @@ export const useQueryStore = defineStore("query", () => {
timeoutSecs: queryTimeoutSecs,
}
: { maxRows: plan.pageLimit, fetchSize: plan.pageLimit, clientSessionId, timeoutSecs: queryTimeoutSecs };
const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, undefined, executionOptions);
const results = await api.executeMulti(tab.connectionId, tab.database, plan.sqlToExecute, tab.schema, exportExecutionId, executionOptions);
const result = results[0];
if (!result) break;
if (columns.length === 0) columns = result.columns;

195
crates/dbx-core/src/connection.rs
View File

@ -1,7 +1,7 @@
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tokio::task::JoinHandle;
@ -84,6 +84,7 @@ pub enum PoolKind {
pub struct AppState {
pub connections: Arc<RwLock<HashMap<String, PoolKind>>>,
keepalive_tasks: Arc<RwLock<HashMap<String, JoinHandle<()>>>>,
pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
pub configs: RwLock<HashMap<String, ConnectionConfig>>,
pub running_queries: RunningQueries,
pub tunnels: TunnelManager,
@ -95,6 +96,40 @@ pub struct AppState {
pub mq_registry: crate::mq::MqAdminRegistry,
}
#[derive(Clone, Copy)]
struct PoolActivity {
last_used_at: Instant,
}
impl PoolActivity {
fn now() -> Self {
Self { last_used_at: Instant::now() }
}
}
pub struct PoolActivityTouch {
pool_key: String,
connections: Arc<RwLock<HashMap<String, PoolKind>>>,
pool_activity: Arc<RwLock<HashMap<String, PoolActivity>>>,
}
impl Drop for PoolActivityTouch {
fn drop(&mut self) {
let pool_key = self.pool_key.clone();
let connections = self.connections.clone();
let pool_activity = self.pool_activity.clone();
let Ok(handle) = tokio::runtime::Handle::try_current() else {
return;
};
handle.spawn(async move {
if !connections.read().await.contains_key(&pool_key) {
return;
}
pool_activity.write().await.insert(pool_key, PoolActivity::now());
});
}
}
pub fn metadata_connection_config(config: &ConnectionConfig) -> ConnectionConfig {
let mut db_config = config.canonicalized();
if database_capabilities::is_metadata_connection_scoped(&db_config.db_type) {
@ -263,6 +298,7 @@ impl AppState {
Self {
connections: Arc::new(RwLock::new(HashMap::new())),
keepalive_tasks: Arc::new(RwLock::new(HashMap::new())),
pool_activity: Arc::new(RwLock::new(HashMap::new())),
configs: RwLock::new(HashMap::new()),
running_queries: RunningQueries::default(),
tunnels: TunnelManager::new(),
@ -326,6 +362,7 @@ impl AppState {
pub async fn insert_connection_pool(&self, pool_key: String, pool: PoolKind, config: &ConnectionConfig) {
self.stop_keepalive_task(&pool_key).await;
self.pool_activity.write().await.insert(pool_key.clone(), PoolActivity::now());
self.start_keepalive_task(&pool_key, &pool, config).await;
let previous = self.connections.write().await.insert(pool_key, pool);
if let Some(pool) = previous {
@ -335,37 +372,77 @@ impl AppState {
async fn start_keepalive_task(&self, pool_key: &str, pool: &PoolKind, config: &ConnectionConfig) {
let interval_secs = config.keepalive_interval_secs;
if interval_secs == 0 {
let idle_timeout_secs = config.idle_timeout_secs;
let idle_cleanup_enabled = is_session_scoped_pool_key(pool_key) && idle_timeout_secs > 0;
let mut target = keepalive_target_from_pool(pool, config);
if interval_secs == 0 && !idle_cleanup_enabled {
return;
}
let Some(mut target) = keepalive_target_from_pool(pool, config) else {
if interval_secs > 0 && target.is_none() {
log::debug!(
"Connection keepalive requested for '{pool_key}', but this database driver does not keep a pingable client handle."
);
return;
if !idle_cleanup_enabled {
return;
}
};
let key = pool_key.to_string();
let interval = Duration::from_secs(interval_secs.max(1));
let interval = if interval_secs > 0 {
Duration::from_secs(interval_secs.max(1))
} else {
Duration::from_secs(idle_timeout_secs.min(60).max(1))
};
let timeout = Duration::from_secs(config.effective_connect_timeout_secs().max(1));
let connections = self.connections.clone();
let keepalive_tasks = self.keepalive_tasks.clone();
let pool_activity = self.pool_activity.clone();
let running_queries = self.running_queries.clone();
let idle_timeout = Duration::from_secs(idle_timeout_secs.max(1));
let handle = tokio::spawn(async move {
loop {
tokio::time::sleep(interval).await;
let result = tokio::time::timeout(timeout, ping_keepalive_target(&mut target, timeout)).await;
match result {
Ok(Ok(())) => {}
Ok(Err(err)) => {
log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool");
if running_queries.is_pool_active(&key) {
continue;
}
if idle_cleanup_enabled {
let idle_for = {
let activity = pool_activity.read().await;
activity.get(&key).map(|activity| activity.last_used_at.elapsed())
};
if idle_for.is_some_and(|elapsed| elapsed >= idle_timeout) {
log::info!(
"Closing idle session-scoped connection pool '{key}' after {}s",
idle_timeout.as_secs()
);
keepalive_tasks.write().await.remove(&key);
pool_activity.write().await.remove(&key);
let removed = connections.write().await.remove(&key);
if let Some(pool) = removed {
close_pool_kind_with_timeout(key, pool).await;
}
break;
}
Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()),
}
if let Some(target) = target.as_mut() {
let result = tokio::time::timeout(timeout, ping_keepalive_target(target, timeout)).await;
match result {
Ok(Ok(())) => {}
Ok(Err(err)) => {
log::warn!("Connection keepalive failed for '{key}': {err}; invalidating pool");
keepalive_tasks.write().await.remove(&key);
pool_activity.write().await.remove(&key);
let removed = connections.write().await.remove(&key);
if let Some(pool) = removed {
close_pool_kind_with_timeout(key, pool).await;
}
break;
}
Err(_) => log::warn!("Connection keepalive timed out for '{key}' after {}s", timeout.as_secs()),
}
}
}
});
@ -391,6 +468,18 @@ impl AppState {
}
}
pub async fn touch_pool_activity(&self, pool_key: &str) {
self.pool_activity.write().await.insert(pool_key.to_string(), PoolActivity::now());
}
pub fn pool_activity_touch(&self, pool_key: &str) -> PoolActivityTouch {
PoolActivityTouch {
pool_key: pool_key.to_string(),
connections: self.connections.clone(),
pool_activity: self.pool_activity.clone(),
}
}
pub async fn get_or_create_pool(&self, connection_id: &str, database: Option<&str>) -> Result<String, String> {
self.get_or_create_pool_for_session(connection_id, database, None).await
}
@ -413,6 +502,7 @@ impl AppState {
if conns.contains_key(&pool_key) {
drop(conns);
if !self.remove_stale_connection_pool(&pool_key).await {
self.touch_pool_activity(&pool_key).await;
return Ok(pool_key);
}
} else {
@ -904,6 +994,10 @@ impl AppState {
}
async fn remove_stale_connection_pool(&self, pool_key: &str) -> bool {
if self.running_queries.is_pool_active(pool_key) {
return false;
}
let stale = {
let connections = self.connections.read().await;
let Some(pool) = connections.get(pool_key) else {
@ -1085,6 +1179,7 @@ impl AppState {
}
self.stop_keepalive_task(pool_key).await;
self.pool_activity.write().await.remove(pool_key);
let removed = self.connections.write().await.remove(pool_key);
if let Some(pool) = removed {
close_pool_kind(pool).await;
@ -1115,6 +1210,7 @@ impl AppState {
self.reset_connection_transport(connection_id).await;
} else {
self.stop_keepalive_task(&pool_key).await;
self.pool_activity.write().await.remove(&pool_key);
let removed = self.connections.write().await.remove(&pool_key);
if let Some(pool) = removed {
close_pool_kind(pool).await;
@ -1143,6 +1239,7 @@ impl AppState {
return Ok(false);
}
self.stop_keepalive_task(&pool_key).await;
self.pool_activity.write().await.remove(&pool_key);
let removed = self.connections.write().await.remove(&pool_key);
if let Some(pool) = removed {
close_pool_kind(pool).await;
@ -1154,6 +1251,7 @@ impl AppState {
pub async fn remove_pool_by_key(&self, pool_key: &str) -> bool {
self.stop_keepalive_task(pool_key).await;
self.pool_activity.write().await.remove(pool_key);
let removed = self.connections.write().await.remove(pool_key);
if let Some(pool) = removed {
close_pool_kind(pool).await;
@ -1182,6 +1280,12 @@ impl AppState {
.cloned()
.collect();
self.stop_keepalive_tasks(&keys_to_remove).await;
{
let mut activity = self.pool_activity.write().await;
for key in &keys_to_remove {
activity.remove(key);
}
}
let mut conns = self.connections.write().await;
let mut removed = Vec::with_capacity(keys_to_remove.len());
for key in keys_to_remove {
@ -1472,6 +1576,12 @@ impl AppState {
// Remove dead pools
if !dead_keys.is_empty() {
self.stop_keepalive_tasks(&dead_keys).await;
{
let mut activity = self.pool_activity.write().await;
for key in &dead_keys {
activity.remove(key);
}
}
let mut conns = self.connections.write().await;
for key in &dead_keys {
if let Some(pool) = conns.remove(key) {
@ -1517,6 +1627,12 @@ impl AppState {
.cloned()
.collect();
self.stop_keepalive_tasks(&keys_to_remove).await;
{
let mut activity = self.pool_activity.write().await;
for key in &keys_to_remove {
activity.remove(key);
}
}
let mut conns = self.connections.write().await;
let mut removed = Vec::with_capacity(keys_to_remove.len());
for key in keys_to_remove {
@ -1545,6 +1661,12 @@ impl AppState {
})
.collect();
self.stop_keepalive_tasks(&keys_to_remove).await;
{
let mut activity = self.pool_activity.write().await;
for key in &keys_to_remove {
activity.remove(key);
}
}
let mut conns = self.connections.write().await;
let mut removed = Vec::with_capacity(keys_to_remove.len());
for key in keys_to_remove {
@ -1697,6 +1819,10 @@ fn session_scoped_pool_key(base_pool_key: String, client_session_id: Option<&str
.unwrap_or(base_pool_key)
}
fn is_session_scoped_pool_key(pool_key: &str) -> bool {
pool_key.contains(":session:")
}
pub(crate) fn config_for_pool_key<'a>(
pool_key: &str,
configs: &'a HashMap<String, ConnectionConfig>,
@ -2763,6 +2889,23 @@ mod tests {
);
}
#[test]
fn session_scoped_pool_keys_are_sanitized_and_detected() {
let key = super::session_scoped_pool_key_for(
Some(DatabaseType::Mysql),
"mysql-conn:analytics".to_string(),
Some("tab-1:count"),
);
assert_eq!(key, "mysql-conn:analytics:session:tab-1_count");
assert!(super::is_session_scoped_pool_key(&key));
assert!(!super::is_session_scoped_pool_key("mysql-conn:analytics"));
assert_eq!(
super::session_scoped_pool_key_for(Some(DatabaseType::DuckDb), "duckdb-conn".to_string(), Some("tab-1")),
"duckdb-conn"
);
}
#[test]
fn mysql_direct_connections_skip_tcp_probe() {
let mut config = mysql_config(Some("app"));
@ -2884,6 +3027,36 @@ mod tests {
let _ = std::fs::remove_dir_all(dir);
}
#[tokio::test]
async fn pool_activity_touch_updates_existing_pool_only() {
let (state, dir) = test_app_state().await;
let pool = crate::db::sqlite::connect_path(":memory:").await.unwrap();
let pool_key = "conn:session:tab-1";
state.connections.write().await.insert(pool_key.to_string(), PoolKind::Sqlite(pool));
state.pool_activity.write().await.insert(
pool_key.to_string(),
super::PoolActivity { last_used_at: std::time::Instant::now() - std::time::Duration::from_secs(10) },
);
{
let _touch = state.pool_activity_touch(pool_key);
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
let elapsed = state.pool_activity.read().await.get(pool_key).unwrap().last_used_at.elapsed();
assert!(elapsed < std::time::Duration::from_secs(10));
{
let _touch = state.pool_activity_touch(pool_key);
state.connections.write().await.remove(pool_key);
state.pool_activity.write().await.remove(pool_key);
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
assert!(!state.pool_activity.read().await.contains_key(pool_key));
let _ = std::fs::remove_dir_all(dir);
}
#[cfg(feature = "duckdb-bundled")]
#[tokio::test]
async fn duckdb_client_session_reuses_base_pool_to_avoid_file_locks() {

16
crates/dbx-core/src/query.rs
View File

@ -805,6 +805,12 @@ pub async fn do_execute(
cancel_token: Option<CancellationToken>,
options: QueryExecutionOptions,
) -> Result<db::QueryResult, String> {
if let Some(execution_id) = options.execution_id.as_deref() {
state.running_queries.set_pool_key(execution_id, pool_key.to_string());
}
state.touch_pool_activity(pool_key).await;
let _activity_touch = state.pool_activity_touch(pool_key);
let query_timeout = resolve_query_timeout(options.timeout_secs);
let (_duckdb_attached_names, conn_name_if_readonly) = {
let configs = state.configs.read().await;
@ -1257,6 +1263,11 @@ async fn execute_postgres_drop_database(
let pool_key = state
.get_or_create_pool_for_session(connection_id, Some(admin_database), options.client_session_id.as_deref())
.await?;
if let Some(execution_id) = options.execution_id.as_deref() {
state.running_queries.set_pool_key(execution_id, pool_key.clone());
}
state.touch_pool_activity(&pool_key).await;
let _activity_touch = state.pool_activity_touch(pool_key.as_str());
if is_canceled(&cancel_token) {
return Err(canceled_error());
@ -1390,6 +1401,11 @@ pub async fn execute_multi_core_with_options(
.get_or_create_pool_for_session(connection_id, Some(database), options.client_session_id.as_deref())
.await?
};
if let Some(execution_id) = options.execution_id.as_deref() {
state.running_queries.set_pool_key(execution_id, pool_key.clone());
}
state.touch_pool_activity(&pool_key).await;
let _activity_touch = state.pool_activity_touch(pool_key.as_str());
let is_sqlserver = {
let connections = state.connections.read().await;

107
crates/dbx-core/src/query_cancel.rs
View File

@ -4,16 +4,64 @@ use tokio_util::sync::CancellationToken;
type InterruptFn = Box<dyn Fn() + Send + 'static>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum RunningTaskKind {
Query,
Count,
Explain,
Export,
Unknown,
}
impl Default for RunningTaskKind {
fn default() -> Self {
Self::Unknown
}
}
#[derive(Clone, Debug, Default)]
pub struct RunningTaskMetadata {
pub kind: RunningTaskKind,
pub connection_id: Option<String>,
pub database: Option<String>,
pub client_session_id: Option<String>,
}
impl RunningTaskMetadata {
pub fn query(
connection_id: impl Into<String>,
database: impl Into<String>,
client_session_id: Option<String>,
) -> Self {
let kind = task_kind_from_client_session_id(client_session_id.as_deref());
Self { kind, connection_id: Some(connection_id.into()), database: Some(database.into()), client_session_id }
}
}
#[derive(Clone)]
struct RunningTask {
token: CancellationToken,
metadata: RunningTaskMetadata,
pool_key: Option<String>,
}
#[derive(Clone, Default)]
pub struct RunningQueries {
inner: Arc<Mutex<HashMap<String, CancellationToken>>>,
inner: Arc<Mutex<HashMap<String, RunningTask>>>,
interrupts: Arc<Mutex<HashMap<String, InterruptFn>>>,
}
impl RunningQueries {
pub fn register(&self, execution_id: String) -> RegisteredQuery {
self.register_task(execution_id, RunningTaskMetadata::default())
}
pub fn register_task(&self, execution_id: String, metadata: RunningTaskMetadata) -> RegisteredQuery {
let token = CancellationToken::new();
self.inner.lock().unwrap_or_else(|e| e.into_inner()).insert(execution_id.clone(), token.clone());
self.inner
.lock()
.unwrap_or_else(|e| e.into_inner())
.insert(execution_id.clone(), RunningTask { token: token.clone(), metadata, pool_key: None });
RegisteredQuery { execution_id, token, running_queries: self.clone() }
}
@ -23,7 +71,8 @@ impl RunningQueries {
}
pub fn cancel(&self, execution_id: &str) -> bool {
let token = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).cloned();
let token =
self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.token.clone());
let interrupt = self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
if let Some(interrupt) = interrupt {
@ -37,17 +86,51 @@ impl RunningQueries {
}
}
pub fn set_pool_key(&self, execution_id: &str, pool_key: impl Into<String>) {
if let Some(task) = self.inner.lock().unwrap_or_else(|e| e.into_inner()).get_mut(execution_id) {
task.pool_key = Some(pool_key.into());
}
}
pub fn is_pool_active(&self, pool_key: &str) -> bool {
self.inner.lock().unwrap_or_else(|e| e.into_inner()).values().any(|task| {
let _kind = task.metadata.kind;
task.pool_key.as_deref() == Some(pool_key)
})
}
#[cfg(test)]
pub fn has(&self, execution_id: &str) -> bool {
self.inner.lock().unwrap_or_else(|e| e.into_inner()).contains_key(execution_id)
}
#[cfg(test)]
pub fn task_kind(&self, execution_id: &str) -> Option<RunningTaskKind> {
self.inner.lock().unwrap_or_else(|e| e.into_inner()).get(execution_id).map(|task| task.metadata.kind)
}
fn remove(&self, execution_id: &str) {
self.inner.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
self.interrupts.lock().unwrap_or_else(|e| e.into_inner()).remove(execution_id);
}
}
fn task_kind_from_client_session_id(client_session_id: Option<&str>) -> RunningTaskKind {
let Some(session_id) = client_session_id else {
return RunningTaskKind::Query;
};
let session_id = session_id.trim().to_ascii_lowercase();
if session_id.ends_with(":count") {
RunningTaskKind::Count
} else if session_id.ends_with(":explain") {
RunningTaskKind::Explain
} else if session_id.ends_with(":export") {
RunningTaskKind::Export
} else {
RunningTaskKind::Query
}
}
pub struct RegisteredQuery {
execution_id: String,
token: CancellationToken,
@ -68,7 +151,7 @@ impl Drop for RegisteredQuery {
#[cfg(test)]
mod tests {
use super::RunningQueries;
use super::{RunningQueries, RunningTaskKind, RunningTaskMetadata};
#[test]
fn cancel_marks_registered_query_as_cancelled() {
@ -103,4 +186,20 @@ mod tests {
assert!(!running.has("exec-1"));
}
#[test]
fn register_task_tracks_kind_and_pool_activity() {
let running = RunningQueries::default();
let registered = running.register_task(
"exec-1".to_string(),
RunningTaskMetadata::query("conn-1", "main", Some("tab-1:export".to_string())),
);
running.set_pool_key("exec-1", "conn-1:main:session:tab-1_export");
assert_eq!(running.task_kind("exec-1"), Some(RunningTaskKind::Export));
assert!(running.is_pool_active("conn-1:main:session:tab-1_export"));
drop(registered);
assert!(!running.is_pool_active("conn-1:main:session:tab-1_export"));
}
}

11
crates/dbx-web/src/routes/query.rs
View File

@ -6,6 +6,7 @@ use serde::Deserialize;
use crate::error::AppError;
use crate::state::WebState;
use dbx_core::query_cancel::RunningTaskMetadata;
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
@ -269,7 +270,10 @@ pub async fn execute_query(
) -> Result<Json<serde_json::Value>, AppError> {
let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let registered = state.app.running_queries.register(execution_id.clone());
let registered = state.app.running_queries.register_task(
execution_id.clone(),
RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()),
);
let cancel_token = registered.token();
let result = dbx_core::query::execute_sql_statement_with_options(
@ -302,7 +306,10 @@ pub async fn execute_multi(
) -> Result<Json<serde_json::Value>, AppError> {
let execution_id = req.execution_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let registered = state.app.running_queries.register(execution_id.clone());
let registered = state.app.running_queries.register_task(
execution_id.clone(),
RunningTaskMetadata::query(req.connection_id.clone(), req.database.clone(), req.client_session_id.clone()),
);
let cancel_token = registered.token();
let result = dbx_core::query::execute_multi_core_with_options(

25
src-tauri/src/commands/query.rs
View File

@ -5,6 +5,7 @@ use tauri::State;
use crate::commands::connection::AppState;
use dbx_core::db;
use dbx_core::models::connection::DatabaseType;
use dbx_core::query_cancel::RunningTaskMetadata;
use dbx_core::sql::split_sql_statements;
#[tauri::command]
@ -23,8 +24,13 @@ pub async fn execute_query(
client_session_id: Option<String>,
timeout_secs: Option<u64>,
) -> Result<db::QueryResult, String> {
let registered_query =
execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone()));
let execution_id = execution_id.filter(|id| !id.trim().is_empty());
let registered_query = execution_id.as_ref().map(|id| {
state.running_queries.register_task(
id.clone(),
RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()),
)
});
let cancel_token = registered_query.as_ref().map(|query| query.token());
dbx_core::query::execute_sql_statement_with_options(
@ -41,7 +47,7 @@ pub async fn execute_query(
result_session_id,
client_session_id,
timeout_secs,
execution_id: execution_id.filter(|id| !id.trim().is_empty()),
execution_id,
},
)
.await
@ -63,10 +69,15 @@ pub async fn execute_multi(
client_session_id: Option<String>,
timeout_secs: Option<u64>,
) -> Result<Vec<db::QueryResult>, String> {
let registered_query =
execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone()));
let execution_id = execution_id.filter(|id| !id.trim().is_empty());
let registered_query = execution_id.as_ref().map(|id| {
state.running_queries.register_task(
id.clone(),
RunningTaskMetadata::query(connection_id.clone(), database.clone(), client_session_id.clone()),
)
});
let cancel_token = registered_query.as_ref().map(|query| query.token());
let trace_id = execution_id.clone().as_deref().unwrap_or("no-execution-id").to_string();
let trace_id = execution_id.as_deref().unwrap_or("no-execution-id").to_string();
let started_at = Instant::now();
log::info!(
"[query][execute_multi:start] trace_id={} connection_id={} database={} schema={:?} sql={}",
@ -91,7 +102,7 @@ pub async fn execute_multi(
result_session_id,
client_session_id,
timeout_secs,
execution_id: execution_id.filter(|id| !id.trim().is_empty()),
execution_id,
},
)
.await;