query/query_engine/
context.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Arc;
16
17use common_query::logical_plan::SubstraitPlanDecoderRef;
18use common_telemetry::tracing_context::TracingContext;
19use datafusion::execution::context::{SessionState, TaskContext};
20use session::context::QueryContextRef;
21
22use crate::query_engine::default_serializer::DefaultPlanDecoder;
23
24#[derive(Debug)]
25pub struct QueryEngineContext {
26    state: SessionState,
27    query_ctx: QueryContextRef,
28}
29
30impl QueryEngineContext {
31    pub fn new(state: SessionState, query_ctx: QueryContextRef) -> Self {
32        Self { state, query_ctx }
33    }
34
35    #[inline]
36    pub fn state(&self) -> &SessionState {
37        &self.state
38    }
39
40    #[inline]
41    pub fn query_ctx(&self) -> QueryContextRef {
42        self.query_ctx.clone()
43    }
44
45    pub fn build_task_ctx(&self) -> Arc<TaskContext> {
46        let dbname = self.query_ctx.get_db_string();
47        let state = &self.state;
48        let tracing_context = TracingContext::from_current_span();
49
50        // pass tracing context in session_id
51        let session_id = tracing_context.to_json();
52
53        Arc::new(TaskContext::new(
54            Some(dbname),
55            session_id,
56            state.config().clone(),
57            state.scalar_functions().clone(),
58            state.higher_order_functions().clone(),
59            state.aggregate_functions().clone(),
60            state.window_functions().clone(),
61            state.runtime_env().clone(),
62        ))
63    }
64
65    /// Creates a [`LogicalPlan`] decoder
66    pub fn new_plan_decoder(&self) -> crate::error::Result<SubstraitPlanDecoderRef> {
67        Ok(Arc::new(DefaultPlanDecoder::new(
68            self.state.clone(),
69            &self.query_ctx,
70        )?))
71    }
72
73    /// Mock an engine context for unit tests.
74    #[cfg(test)]
75    pub fn mock() -> Self {
76        use common_base::Plugins;
77        use session::context::QueryContext;
78
79        use crate::options::QueryOptions;
80        use crate::query_engine::QueryEngineState;
81
82        let state = Arc::new(QueryEngineState::new(
83            catalog::memory::new_memory_catalog_manager().unwrap(),
84            None,
85            None,
86            None,
87            None,
88            None,
89            false,
90            Plugins::default(),
91            QueryOptions::default(),
92        ));
93
94        QueryEngineContext::new(state.session_state(), QueryContext::arc())
95    }
96}
query/query_engine/context.rs

query/query_engine/
context.rs
