catalog/system_schema/information_schema/
procedure_info.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 arrow_schema::SchemaRef as ArrowSchemaRef;
18use common_catalog::consts::INFORMATION_SCHEMA_PROCEDURE_INFO_TABLE_ID;
19use common_error::ext::BoxedError;
20use common_procedure::ProcedureInfo;
21use common_recordbatch::adapter::RecordBatchStreamAdapter;
22use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
23use common_time::timestamp::Timestamp;
24use datafusion::execution::TaskContext;
25use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
26use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
27use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
28use datatypes::prelude::{ConcreteDataType, ScalarVectorBuilder, VectorRef};
29use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
30use datatypes::timestamp::TimestampMillisecond;
31use datatypes::value::Value;
32use datatypes::vectors::{StringVectorBuilder, TimestampMillisecondVectorBuilder};
33use snafu::ResultExt;
34use store_api::storage::{ScanRequest, TableId};
35
36use crate::CatalogManager;
37use crate::error::{CreateRecordBatchSnafu, InternalSnafu, Result};
38use crate::system_schema::SystemTable;
39use crate::system_schema::information_schema::{
40    InformationSchemaTableScope, InformationTable, PROCEDURE_INFO, Predicates,
41};
42use crate::system_schema::utils;
43
44const PROCEDURE_ID: &str = "procedure_id";
45const PROCEDURE_TYPE: &str = "procedure_type";
46const START_TIME: &str = "start_time";
47const END_TIME: &str = "end_time";
48const STATUS: &str = "status";
49const LOCK_KEYS: &str = "lock_keys";
50
51const INIT_CAPACITY: usize = 42;
52
53/// The `PROCEDURE_INFO` table provides information about the current procedure information of the cluster.
54///
55/// - `procedure_id`: the unique identifier of the procedure.
56/// - `procedure_name`: the name of the procedure.
57/// - `start_time`: the starting execution time of the procedure.
58/// - `end_time`: the ending execution time of the procedure.
59/// - `status`: the status of the procedure.
60/// - `lock_keys`: the lock keys of the procedure.
61#[derive(Debug)]
62pub(super) struct InformationSchemaProcedureInfo {
63    schema: SchemaRef,
64    catalog_manager: Weak<dyn CatalogManager>,
65}
66
67impl InformationSchemaProcedureInfo {
68    pub(super) fn new(catalog_manager: Weak<dyn CatalogManager>) -> Self {
69        Self {
70            schema: Self::schema(),
71            catalog_manager,
72        }
73    }
74
75    pub(crate) fn schema() -> SchemaRef {
76        Arc::new(Schema::new(vec![
77            ColumnSchema::new(PROCEDURE_ID, ConcreteDataType::string_datatype(), false),
78            ColumnSchema::new(PROCEDURE_TYPE, ConcreteDataType::string_datatype(), false),
79            ColumnSchema::new(
80                START_TIME,
81                ConcreteDataType::timestamp_millisecond_datatype(),
82                true,
83            ),
84            ColumnSchema::new(
85                END_TIME,
86                ConcreteDataType::timestamp_millisecond_datatype(),
87                true,
88            ),
89            ColumnSchema::new(STATUS, ConcreteDataType::string_datatype(), false),
90            ColumnSchema::new(LOCK_KEYS, ConcreteDataType::string_datatype(), true),
91        ]))
92    }
93
94    fn builder(&self) -> InformationSchemaProcedureInfoBuilder {
95        InformationSchemaProcedureInfoBuilder::new(
96            self.schema.clone(),
97            self.catalog_manager.clone(),
98        )
99    }
100}
101
102impl InformationTable for InformationSchemaProcedureInfo {
103    fn scope(&self) -> InformationSchemaTableScope {
104        InformationSchemaTableScope::Global
105    }
106}
107
108impl SystemTable for InformationSchemaProcedureInfo {
109    fn table_id(&self) -> TableId {
110        INFORMATION_SCHEMA_PROCEDURE_INFO_TABLE_ID
111    }
112
113    fn table_name(&self) -> &'static str {
114        PROCEDURE_INFO
115    }
116
117    fn schema(&self) -> SchemaRef {
118        self.schema.clone()
119    }
120
121    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
122        let schema = self.schema.arrow_schema().clone();
123        let mut builder = self.builder();
124        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
125            schema,
126            futures::stream::once(async move {
127                builder
128                    .make_procedure_info(Some(request))
129                    .await
130                    .map(|x| x.into_df_record_batch())
131                    .map_err(Into::into)
132            }),
133        ));
134        Ok(Box::pin(
135            RecordBatchStreamAdapter::try_new(stream)
136                .map_err(BoxedError::new)
137                .context(InternalSnafu)?,
138        ))
139    }
140}
141
142struct InformationSchemaProcedureInfoBuilder {
143    schema: SchemaRef,
144    catalog_manager: Weak<dyn CatalogManager>,
145
146    procedure_ids: StringVectorBuilder,
147    procedure_types: StringVectorBuilder,
148    start_times: TimestampMillisecondVectorBuilder,
149    end_times: TimestampMillisecondVectorBuilder,
150    statuses: StringVectorBuilder,
151    lock_keys: StringVectorBuilder,
152}
153
154impl InformationSchemaProcedureInfoBuilder {
155    fn new(schema: SchemaRef, catalog_manager: Weak<dyn CatalogManager>) -> Self {
156        Self {
157            schema,
158            catalog_manager,
159            procedure_ids: StringVectorBuilder::with_capacity(INIT_CAPACITY),
160            procedure_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
161            start_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
162            end_times: TimestampMillisecondVectorBuilder::with_capacity(INIT_CAPACITY),
163            statuses: StringVectorBuilder::with_capacity(INIT_CAPACITY),
164            lock_keys: StringVectorBuilder::with_capacity(INIT_CAPACITY),
165        }
166    }
167
168    /// Construct the `information_schema.procedure_info` virtual table
169    async fn make_procedure_info(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
170        let predicates = Predicates::from_scan_request(&request);
171        let information_extension = utils::information_extension(&self.catalog_manager)?;
172        let procedures = information_extension.procedures().await?;
173        for (status, procedure_info) in procedures {
174            self.add_procedure(&predicates, status, procedure_info);
175        }
176        self.finish()
177    }
178
179    fn add_procedure(
180        &mut self,
181        predicates: &Predicates,
182        status: String,
183        procedure_info: ProcedureInfo,
184    ) {
185        let ProcedureInfo {
186            id,
187            type_name,
188            start_time_ms,
189            end_time_ms,
190            lock_keys,
191            ..
192        } = procedure_info;
193        let pid = id.to_string();
194        let start_time = TimestampMillisecond(Timestamp::new_millisecond(start_time_ms));
195        let end_time = TimestampMillisecond(Timestamp::new_millisecond(end_time_ms));
196        let lock_keys = lock_keys.join(",");
197
198        let row = [
199            (PROCEDURE_ID, &Value::from(pid.clone())),
200            (PROCEDURE_TYPE, &Value::from(type_name.clone())),
201            (START_TIME, &Value::from(start_time)),
202            (END_TIME, &Value::from(end_time)),
203            (STATUS, &Value::from(status.clone())),
204            (LOCK_KEYS, &Value::from(lock_keys.clone())),
205        ];
206        if !predicates.eval(&row) {
207            return;
208        }
209        self.procedure_ids.push(Some(&pid));
210        self.procedure_types.push(Some(&type_name));
211        self.start_times.push(Some(start_time));
212        self.end_times.push(Some(end_time));
213        self.statuses.push(Some(&status));
214        self.lock_keys.push(Some(&lock_keys));
215    }
216
217    fn finish(&mut self) -> Result<RecordBatch> {
218        let columns: Vec<VectorRef> = vec![
219            Arc::new(self.procedure_ids.finish()),
220            Arc::new(self.procedure_types.finish()),
221            Arc::new(self.start_times.finish()),
222            Arc::new(self.end_times.finish()),
223            Arc::new(self.statuses.finish()),
224            Arc::new(self.lock_keys.finish()),
225        ];
226        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
227    }
228}
229
230impl DfPartitionStream for InformationSchemaProcedureInfo {
231    fn schema(&self) -> &ArrowSchemaRef {
232        self.schema.arrow_schema()
233    }
234
235    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
236        let schema = self.schema.arrow_schema().clone();
237        let mut builder = self.builder();
238        Box::pin(DfRecordBatchStreamAdapter::new(
239            schema,
240            futures::stream::once(async move {
241                builder
242                    .make_procedure_info(None)
243                    .await
244                    .map(|x| x.into_df_record_batch())
245                    .map_err(Into::into)
246            }),
247        ))
248    }
249}
catalog/system_schema/information_schema/procedure_info.rs

catalog/system_schema/information_schema/
procedure_info.rs
