catalog/system_schema/information_schema/
flow_statistics.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, Weak};
16
17use common_catalog::consts::INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID;
18use common_error::ext::BoxedError;
19use common_meta::key::FlowId;
20use common_meta::key::flow::FlowMetadataManager;
21use common_meta::key::flow::flow_state::FlowStat;
22use common_recordbatch::adapter::RecordBatchStreamAdapter;
23use common_recordbatch::{DfSendableRecordBatchStream, RecordBatch, SendableRecordBatchStream};
24use common_time::util::current_time_millis;
25use datafusion::execution::TaskContext;
26use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
27use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
28use datatypes::prelude::ConcreteDataType as CDT;
29use datatypes::scalars::ScalarVectorBuilder;
30use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
31use datatypes::timestamp::TimestampMillisecond;
32use datatypes::value::Value;
33use datatypes::vectors::{
34    Int64VectorBuilder, StringVectorBuilder, TimestampMillisecondVectorBuilder,
35    UInt32VectorBuilder, UInt64VectorBuilder, VectorRef,
36};
37use futures::TryStreamExt;
38use snafu::ResultExt;
39use store_api::storage::{ScanRequest, TableId};
40
41use crate::CatalogManager;
42use crate::error::{CreateRecordBatchSnafu, InternalSnafu, ListFlowsSnafu, Result};
43use crate::information_schema::{FLOW_STATISTICS, Predicates};
44use crate::system_schema::SystemTable;
45use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable};
46use crate::system_schema::utils;
47
48const INIT_CAPACITY: usize = 42;
49
50// rows of information_schema.flow_statistics
51pub const FLOW_ID: &str = "flow_id";
52pub const FLOW_NAME: &str = "flow_name";
53pub const START_TIME: &str = "start_time";
54pub const LAST_EXECUTION_TIME: &str = "last_execution_time";
55pub const UPTIME_SECONDS: &str = "uptime_seconds";
56pub const STATE_SIZE: &str = "state_size";
57
58/// The `information_schema.flow_statistics` provides runtime statistics about flows.
59#[derive(Debug)]
60pub(super) struct InformationSchemaFlowStatistics {
61    schema: SchemaRef,
62    catalog_name: String,
63    catalog_manager: Weak<dyn CatalogManager>,
64    flow_metadata_manager: Arc<FlowMetadataManager>,
65}
66
67impl InformationSchemaFlowStatistics {
68    pub(super) fn new(
69        catalog_name: String,
70        catalog_manager: Weak<dyn CatalogManager>,
71        flow_metadata_manager: Arc<FlowMetadataManager>,
72    ) -> Self {
73        Self {
74            schema: Self::schema(),
75            catalog_name,
76            catalog_manager,
77            flow_metadata_manager,
78        }
79    }
80
81    pub(crate) fn schema() -> SchemaRef {
82        Arc::new(Schema::new(
83            vec![
84                (FLOW_ID, CDT::uint32_datatype(), false),
85                (FLOW_NAME, CDT::string_datatype(), false),
86                (START_TIME, CDT::timestamp_millisecond_datatype(), true),
87                (
88                    LAST_EXECUTION_TIME,
89                    CDT::timestamp_millisecond_datatype(),
90                    true,
91                ),
92                (UPTIME_SECONDS, CDT::int64_datatype(), true),
93                (STATE_SIZE, CDT::uint64_datatype(), true),
94            ]
95            .into_iter()
96            .map(|(name, ty, nullable)| ColumnSchema::new(name, ty, nullable))
97            .collect(),
98        ))
99    }
100
101    fn builder(&self) -> InformationSchemaFlowStatisticsBuilder {
102        InformationSchemaFlowStatisticsBuilder::new(
103            self.schema.clone(),
104            self.catalog_name.clone(),
105            self.catalog_manager.clone(),
106            &self.flow_metadata_manager,
107        )
108    }
109}
110
111impl InformationTable for InformationSchemaFlowStatistics {
112    fn scope(&self) -> InformationSchemaTableScope {
113        InformationSchemaTableScope::Catalog
114    }
115}
116
117impl SystemTable for InformationSchemaFlowStatistics {
118    fn table_id(&self) -> TableId {
119        INFORMATION_SCHEMA_FLOW_STATISTICS_TABLE_ID
120    }
121
122    fn table_name(&self) -> &'static str {
123        FLOW_STATISTICS
124    }
125
126    fn schema(&self) -> SchemaRef {
127        self.schema.clone()
128    }
129
130    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
131        let schema = self.schema.arrow_schema().clone();
132        let mut builder = self.builder();
133        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
134            schema,
135            futures::stream::once(async move {
136                builder
137                    .make_flow_statistics(Some(request))
138                    .await
139                    .map(|x| x.into_df_record_batch())
140                    .map_err(|err| datafusion::error::DataFusionError::External(Box::new(err)))
141            }),
142        ));
143        Ok(Box::pin(
144            RecordBatchStreamAdapter::try_new(stream)
145                .map_err(BoxedError::new)
146                .context(InternalSnafu)?,
147        ))
148    }
149}
150
151/// Builds the `information_schema.flow_statistics` table row by row.
152struct InformationSchemaFlowStatisticsBuilder {
153    schema: SchemaRef,
154    catalog_name: String,
155    catalog_manager: Weak<dyn CatalogManager>,
156    flow_metadata_manager: Arc<FlowMetadataManager>,
157
158    flow_ids: UInt32VectorBuilder,
159    flow_names: StringVectorBuilder,
160    start_times: TimestampMillisecondVectorBuilder,
161    last_execution_times: TimestampMillisecondVectorBuilder,
162    uptime_seconds: Int64VectorBuilder,
163    state_sizes: UInt64VectorBuilder,
164}
165
166impl InformationSchemaFlowStatisticsBuilder {
167    fn new(
168        schema: SchemaRef,
169        catalog_name: String,
170        catalog_manager: Weak<dyn CatalogManager>,
171        flow_metadata_manager: &Arc<FlowMetadataManager>,
172    ) -> Self {
173        Self {
174            schema,
175            catalog_name,
176            catalog_manager,
177            flow_metadata_manager: flow_metadata_manager.clone(),
178
179            flow_ids: UInt32VectorBuilder::with_capacity(INIT_CAPACITY),
180            flow_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
181            start_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
182            last_execution_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
183            uptime_seconds: Int64VectorBuilder::with_capacity(INIT_CAPACITY),
184            state_sizes: UInt64VectorBuilder::with_capacity(INIT_CAPACITY),
185        }
186    }
187
188    /// Construct the `information_schema.flow_statistics` virtual table.
189    async fn make_flow_statistics(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
190        let catalog_name = self.catalog_name.clone();
191        let predicates = Predicates::from_scan_request(&request);
192
193        let flow_info_manager = self.flow_metadata_manager.clone();
194
195        let mut stream = flow_info_manager
196            .flow_name_manager()
197            .flow_names(&catalog_name)
198            .await;
199
200        let flow_stat = {
201            let information_extension = utils::information_extension(&self.catalog_manager)?;
202            information_extension.flow_stats().await?
203        };
204
205        let now = current_time_millis();
206
207        while let Some((flow_name, flow_id)) = stream
208            .try_next()
209            .await
210            .map_err(BoxedError::new)
211            .context(ListFlowsSnafu {
212                catalog: &catalog_name,
213            })?
214        {
215            self.add_flow_statistic(&predicates, flow_id.flow_id(), &flow_name, &flow_stat, now);
216        }
217
218        self.finish()
219    }
220
221    fn add_flow_statistic(
222        &mut self,
223        predicates: &Predicates,
224        flow_id: FlowId,
225        flow_name: &str,
226        flow_stat: &Option<FlowStat>,
227        now: i64,
228    ) {
229        let row = [
230            (FLOW_ID, &Value::from(flow_id)),
231            (FLOW_NAME, &Value::from(flow_name.to_string())),
232        ];
233        if !predicates.eval(&row) {
234            return;
235        }
236
237        let start_time = flow_stat
238            .as_ref()
239            .and_then(|stat| stat.start_time_map.get(&flow_id).copied());
240
241        self.flow_ids.push(Some(flow_id));
242        self.flow_names.push(Some(flow_name));
243        self.start_times
244            .push(start_time.map(TimestampMillisecond::new));
245        self.last_execution_times
246            .push(flow_stat.as_ref().and_then(|stat| {
247                stat.last_exec_time_map
248                    .get(&flow_id)
249                    .map(|v| TimestampMillisecond::new(*v))
250            }));
251        self.uptime_seconds
252            .push(start_time.map(|start| ((now - start) / 1000).max(0)));
253        self.state_sizes.push(
254            flow_stat
255                .as_ref()
256                .and_then(|stat| stat.state_size.get(&flow_id).map(|v| *v as u64)),
257        );
258    }
259
260    fn finish(&mut self) -> Result<RecordBatch> {
261        let columns: Vec<VectorRef> = vec![
262            Arc::new(self.flow_ids.finish()),
263            Arc::new(self.flow_names.finish()),
264            Arc::new(self.start_times.finish()),
265            Arc::new(self.last_execution_times.finish()),
266            Arc::new(self.uptime_seconds.finish()),
267            Arc::new(self.state_sizes.finish()),
268        ];
269        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
270    }
271}
272
273impl DfPartitionStream for InformationSchemaFlowStatistics {
274    fn schema(&self) -> &arrow_schema::SchemaRef {
275        self.schema.arrow_schema()
276    }
277
278    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
279        let schema: Arc<arrow_schema::Schema> = self.schema.arrow_schema().clone();
280        let mut builder = self.builder();
281        Box::pin(DfRecordBatchStreamAdapter::new(
282            schema,
283            futures::stream::once(async move {
284                builder
285                    .make_flow_statistics(None)
286                    .await
287                    .map(|x| x.into_df_record_batch())
288                    .map_err(Into::into)
289            }),
290        ))
291    }
292}
catalog/system_schema/information_schema/flow_statistics.rs

catalog/system_schema/information_schema/
flow_statistics.rs
