catalog/
information_extension.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 api::v1::meta::ProcedureStatus;
18use common_error::ext::BoxedError;
19use common_meta::cluster::{ClusterInfo, NodeInfo, Role};
20use common_meta::datanode::RegionStat;
21use common_meta::key::flow::flow_state::FlowStat;
22use common_meta::node_manager::DatanodeManagerRef;
23use common_meta::procedure_executor::{ExecutorContext, ProcedureExecutor};
24use common_meta::rpc::procedure;
25use common_procedure::{ProcedureInfo, ProcedureState};
26use common_query::request::QueryRequest;
27use common_recordbatch::adapter::{AsyncRecordBatchStreamAdapter, DfRecordBatchStreamAdapter};
28use common_recordbatch::util::{ChainedRecordBatchStream, LimitedRecordBatchStream};
29use common_recordbatch::{DfSendableRecordBatchStream, SendableRecordBatchStream};
30use datafusion::execution::TaskContext;
31use datafusion::physical_expr::{EquivalenceProperties, Partitioning};
32use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
33use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties};
34use datatypes::arrow::datatypes::SchemaRef as ArrowSchemaRef;
35use datatypes::schema::SchemaRef;
36use futures_util::future::try_join_all;
37use meta_client::MetaClientRef;
38use snafu::ResultExt;
39use store_api::storage::RegionId;
40
41use crate::error;
42use crate::information_schema::{DatanodeInspectRequest, InformationExtension};
43
44pub struct DistributedInformationExtension {
45    meta_client: MetaClientRef,
46    datanode_manager: DatanodeManagerRef,
47}
48
49impl DistributedInformationExtension {
50    pub fn new(meta_client: MetaClientRef, datanode_manager: DatanodeManagerRef) -> Self {
51        Self {
52            meta_client,
53            datanode_manager,
54        }
55    }
56
57    async fn inspect_datanode_stream(
58        meta_client: MetaClientRef,
59        datanode_manager: DatanodeManagerRef,
60        request: DatanodeInspectRequest,
61    ) -> std::result::Result<SendableRecordBatchStream, crate::error::Error> {
62        let nodes = meta_client
63            .list_nodes(Some(Role::Datanode))
64            .await
65            .map_err(BoxedError::new)
66            .context(crate::error::ListNodesSnafu)?;
67
68        let limit = request.scan.limit;
69        let plan = request
70            .build_plan()
71            .context(crate::error::DatafusionSnafu)?;
72
73        let streams = try_join_all(nodes.into_iter().map(|node| {
74            let datanode_manager = datanode_manager.clone();
75            let plan = plan.clone();
76            async move {
77                let client = datanode_manager.datanode(&node.peer).await;
78                client
79                    .handle_query(QueryRequest {
80                        plan,
81                        region_id: RegionId::default(),
82                        header: None,
83                    })
84                    .await
85                    .context(crate::error::HandleQuerySnafu)
86            }
87        }))
88        .await?;
89
90        let chained =
91            ChainedRecordBatchStream::new(streams).context(crate::error::CreateRecordBatchSnafu)?;
92        match limit {
93            Some(limit) => Ok(Box::pin(LimitedRecordBatchStream::new(
94                Box::pin(chained),
95                limit,
96            ))),
97            None => Ok(Box::pin(chained)),
98        }
99    }
100}
101
102struct DistributedInspectExec {
103    request: DatanodeInspectRequest,
104    meta_client: MetaClientRef,
105    datanode_manager: DatanodeManagerRef,
106    schema: SchemaRef,
107    arrow_schema: ArrowSchemaRef,
108    properties: Arc<PlanProperties>,
109}
110
111impl std::fmt::Debug for DistributedInspectExec {
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        f.debug_struct("DistributedInspectExec")
114            .field("request", &self.request)
115            .field("schema", &self.schema)
116            .finish_non_exhaustive()
117    }
118}
119
120impl DistributedInspectExec {
121    fn new(
122        request: DatanodeInspectRequest,
123        schema: SchemaRef,
124        meta_client: MetaClientRef,
125        datanode_manager: DatanodeManagerRef,
126    ) -> Self {
127        let arrow_schema = schema.arrow_schema().clone();
128        let properties = Arc::new(PlanProperties::new(
129            EquivalenceProperties::new(arrow_schema.clone()),
130            Partitioning::UnknownPartitioning(1),
131            EmissionType::Incremental,
132            Boundedness::Bounded,
133        ));
134        Self {
135            request,
136            meta_client,
137            datanode_manager,
138            schema,
139            arrow_schema,
140            properties,
141        }
142    }
143}
144
145impl DisplayAs for DistributedInspectExec {
146    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
147        write!(
148            f,
149            "DistributedInspectExec: kind={:?}, scan={}, schema={:?}",
150            self.request.kind, self.request.scan, self.arrow_schema
151        )
152    }
153}
154
155impl ExecutionPlan for DistributedInspectExec {
156    fn name(&self) -> &str {
157        "DistributedInspectExec"
158    }
159
160    fn as_any(&self) -> &dyn std::any::Any {
161        self
162    }
163
164    fn schema(&self) -> ArrowSchemaRef {
165        self.arrow_schema.clone()
166    }
167
168    fn properties(&self) -> &Arc<PlanProperties> {
169        &self.properties
170    }
171
172    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
173        vec![]
174    }
175
176    fn with_new_children(
177        self: Arc<Self>,
178        children: Vec<Arc<dyn ExecutionPlan>>,
179    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
180        if children.is_empty() {
181            Ok(self)
182        } else {
183            Err(datafusion::error::DataFusionError::Internal(
184                "DistributedInspectExec is a leaf execution plan and cannot accept children"
185                    .to_string(),
186            ))
187        }
188    }
189
190    fn execute(
191        &self,
192        partition: usize,
193        _context: Arc<TaskContext>,
194    ) -> datafusion::error::Result<DfSendableRecordBatchStream> {
195        if partition != 0 {
196            return Err(datafusion::error::DataFusionError::Execution(format!(
197                "DistributedInspectExec invalid partition. Expected 0, got {partition}"
198            )));
199        }
200
201        let meta_client = self.meta_client.clone();
202        let datanode_manager = self.datanode_manager.clone();
203        let request = self.request.clone();
204        let schema = self.schema.clone();
205        let future = async move {
206            DistributedInformationExtension::inspect_datanode_stream(
207                meta_client,
208                datanode_manager,
209                request,
210            )
211            .await
212            .map_err(BoxedError::new)
213            .context(common_recordbatch::error::ExternalSnafu)
214        };
215        let stream = AsyncRecordBatchStreamAdapter::new(schema, Box::pin(future));
216        Ok(Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))))
217    }
218}
219
220#[async_trait::async_trait]
221impl InformationExtension for DistributedInformationExtension {
222    type Error = crate::error::Error;
223
224    async fn nodes(&self) -> std::result::Result<Vec<NodeInfo>, Self::Error> {
225        self.meta_client
226            .list_nodes(None)
227            .await
228            .map_err(BoxedError::new)
229            .context(error::ListNodesSnafu)
230    }
231
232    async fn procedures(&self) -> std::result::Result<Vec<(String, ProcedureInfo)>, Self::Error> {
233        let procedures = self
234            .meta_client
235            .list_procedures(&ExecutorContext::default())
236            .await
237            .map_err(BoxedError::new)
238            .context(error::ListProceduresSnafu)?
239            .procedures;
240        let mut result = Vec::with_capacity(procedures.len());
241        for procedure in procedures {
242            let pid = match procedure.id {
243                Some(pid) => pid,
244                None => return error::ProcedureIdNotFoundSnafu {}.fail(),
245            };
246            let pid = procedure::pb_pid_to_pid(&pid)
247                .map_err(BoxedError::new)
248                .context(error::ConvertProtoDataSnafu)?;
249            let status = ProcedureStatus::try_from(procedure.status)
250                .map(|v| v.as_str_name())
251                .unwrap_or("Unknown")
252                .to_string();
253            let procedure_info = ProcedureInfo {
254                id: pid,
255                type_name: procedure.type_name,
256                start_time_ms: procedure.start_time_ms,
257                end_time_ms: procedure.end_time_ms,
258                state: ProcedureState::Running,
259                lock_keys: procedure.lock_keys,
260            };
261            result.push((status, procedure_info));
262        }
263
264        Ok(result)
265    }
266
267    async fn region_stats(&self) -> std::result::Result<Vec<RegionStat>, Self::Error> {
268        self.meta_client
269            .list_region_stats()
270            .await
271            .map_err(BoxedError::new)
272            .context(error::ListRegionStatsSnafu)
273    }
274
275    async fn flow_stats(&self) -> std::result::Result<Option<FlowStat>, Self::Error> {
276        self.meta_client
277            .list_flow_stats()
278            .await
279            .map_err(BoxedError::new)
280            .context(crate::error::ListFlowStatsSnafu)
281    }
282
283    async fn inspect_datanode(
284        &self,
285        request: DatanodeInspectRequest,
286    ) -> std::result::Result<SendableRecordBatchStream, Self::Error> {
287        Self::inspect_datanode_stream(
288            self.meta_client.clone(),
289            self.datanode_manager.clone(),
290            request,
291        )
292        .await
293    }
294
295    fn inspect_datanode_plan(
296        &self,
297        request: DatanodeInspectRequest,
298        schema: SchemaRef,
299    ) -> std::result::Result<Option<Arc<dyn ExecutionPlan>>, Self::Error> {
300        Ok(Some(Arc::new(DistributedInspectExec::new(
301            request,
302            schema,
303            self.meta_client.clone(),
304            self.datanode_manager.clone(),
305        ))))
306    }
307}
catalog/information_extension.rs

catalog/
information_extension.rs
