/
query.rs
120 lines (104 loc) · 3.38 KB
/
query.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
use std::sync::Arc;
use async_trait::async_trait;
use futures::stream::AbortHandle;
use parking_lot::Mutex;
use spi::query::dispatcher::{QueryInfo, QueryStatus};
use spi::query::execution::{Output, QueryExecution, QueryStateMachineRef};
use spi::query::logical_planner::QueryPlan;
use spi::query::optimizer::Optimizer;
use spi::query::scheduler::SchedulerRef;
use spi::{QueryError, Result};
use trace::debug;
pub struct SqlQueryExecution {
query_state_machine: QueryStateMachineRef,
plan: QueryPlan,
optimizer: Arc<dyn Optimizer + Send + Sync>,
scheduler: SchedulerRef,
abort_handle: Mutex<Option<AbortHandle>>,
}
impl SqlQueryExecution {
pub fn new(
query_state_machine: QueryStateMachineRef,
plan: QueryPlan,
optimizer: Arc<dyn Optimizer + Send + Sync>,
scheduler: SchedulerRef,
) -> Self {
Self {
query_state_machine,
plan,
optimizer,
scheduler,
abort_handle: Mutex::new(None),
}
}
async fn start(&self) -> Result<Output> {
// begin optimize
self.query_state_machine.begin_optimize();
let physical_plan = self
.optimizer
.optimize(&self.plan.df_plan, &self.query_state_machine.session)
.await?;
self.query_state_machine.end_optimize();
// begin schedule
self.query_state_machine.begin_schedule();
let stream = self
.scheduler
.schedule(
physical_plan.clone(),
self.query_state_machine.session.inner().task_ctx(),
)
.await?
.stream();
debug!("Success build result stream.");
self.query_state_machine.end_schedule();
Ok(Output::StreamData(stream))
}
}
#[async_trait]
impl QueryExecution for SqlQueryExecution {
async fn start(&self) -> Result<Output> {
let (task, abort_handle) = futures::future::abortable(self.start());
{
*self.abort_handle.lock() = Some(abort_handle);
}
task.await.map_err(|_| QueryError::Cancel)?
}
fn cancel(&self) -> Result<()> {
debug!(
"cancel sql query execution: query_id: {:?}, sql: {}, state: {:?}",
&self.query_state_machine.query_id,
self.query_state_machine.query.content(),
self.query_state_machine.state()
);
// change state
self.query_state_machine.cancel();
// stop future task
if let Some(e) = self.abort_handle.lock().as_ref() {
e.abort()
};
debug!(
"canceled sql query execution: query_id: {:?}, sql: {}, state: {:?}",
&self.query_state_machine.query_id,
self.query_state_machine.query.content(),
self.query_state_machine.state()
);
Ok(())
}
fn info(&self) -> QueryInfo {
let qsm = &self.query_state_machine;
QueryInfo::new(
qsm.query_id,
qsm.query.content().to_string(),
*qsm.session.tenant_id(),
qsm.session.tenant().to_string(),
qsm.session.default_database().to_string(),
qsm.session.user().clone(),
)
}
fn status(&self) -> QueryStatus {
QueryStatus::new(
self.query_state_machine.state().clone(),
self.query_state_machine.duration(),
)
}
}