Skip to main content

catalog/system_schema/information_schema/
process_list.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_catalog::consts::INFORMATION_SCHEMA_PROCESS_LIST_TABLE_ID;
18use common_error::ext::BoxedError;
19use common_frontend::DisplayProcessId;
20use common_recordbatch::adapter::RecordBatchStreamAdapter;
21use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
22use common_time::util::current_time_millis;
23use common_time::{Duration, Timestamp};
24use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
25use datatypes::prelude::ConcreteDataType as CDT;
26use datatypes::scalars::ScalarVectorBuilder;
27use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
28use datatypes::value::Value;
29use datatypes::vectors::{
30    DurationMillisecondVectorBuilder, StringVectorBuilder, TimestampMillisecondVectorBuilder,
31    VectorRef,
32};
33use snafu::ResultExt;
34use store_api::storage::{ScanRequest, TableId};
35
36use crate::error::{self, InternalSnafu};
37use crate::information_schema::Predicates;
38use crate::process_manager::ProcessManagerRef;
39use crate::system_schema::SystemTable;
40use crate::system_schema::information_schema::{InformationSchemaTableScope, InformationTable};
41
42/// Column names of `information_schema.process_list`
43pub const ID: &str = "id";
44pub const CATALOG: &str = "catalog";
45pub const SCHEMAS: &str = "schemas";
46pub const QUERY: &str = "query";
47pub const CLIENT: &str = "client";
48pub const FRONTEND: &str = "frontend";
49pub const START_TIMESTAMP: &str = "start_timestamp";
50pub const ELAPSED_TIME: &str = "elapsed_time";
51
52/// `information_schema.process_list` table implementation that tracks running
53/// queries in current cluster.
54pub struct InformationSchemaProcessList {
55    schema: SchemaRef,
56    process_manager: ProcessManagerRef,
57}
58
59impl InformationSchemaProcessList {
60    pub fn new(process_manager: ProcessManagerRef) -> Self {
61        Self {
62            schema: Self::schema(),
63            process_manager,
64        }
65    }
66
67    fn schema() -> SchemaRef {
68        Arc::new(Schema::new(vec![
69            ColumnSchema::new(ID, CDT::string_datatype(), false),
70            ColumnSchema::new(CATALOG, CDT::string_datatype(), false),
71            ColumnSchema::new(SCHEMAS, CDT::string_datatype(), false),
72            ColumnSchema::new(QUERY, CDT::string_datatype(), false),
73            ColumnSchema::new(CLIENT, CDT::string_datatype(), false),
74            ColumnSchema::new(FRONTEND, CDT::string_datatype(), false),
75            ColumnSchema::new(
76                START_TIMESTAMP,
77                CDT::timestamp_millisecond_datatype(),
78                false,
79            ),
80            ColumnSchema::new(ELAPSED_TIME, CDT::duration_millisecond_datatype(), false),
81        ]))
82    }
83}
84
85impl InformationTable for InformationSchemaProcessList {
86    fn scope(&self) -> InformationSchemaTableScope {
87        InformationSchemaTableScope::Global
88    }
89}
90
91impl SystemTable for InformationSchemaProcessList {
92    fn table_id(&self) -> TableId {
93        INFORMATION_SCHEMA_PROCESS_LIST_TABLE_ID
94    }
95
96    fn table_name(&self) -> &'static str {
97        "process_list"
98    }
99
100    fn schema(&self) -> SchemaRef {
101        self.schema.clone()
102    }
103
104    fn to_stream(&self, request: ScanRequest) -> error::Result<SendableRecordBatchStream> {
105        let process_manager = self.process_manager.clone();
106        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
107            self.schema.arrow_schema().clone(),
108            futures::stream::once(async move {
109                make_process_list(process_manager, request)
110                    .await
111                    .map(RecordBatch::into_df_record_batch)
112                    .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))
113            }),
114        ));
115
116        Ok(Box::pin(
117            RecordBatchStreamAdapter::try_new(stream)
118                .map_err(BoxedError::new)
119                .context(InternalSnafu)?,
120        ))
121    }
122}
123
124/// Build running process list.
125async fn make_process_list(
126    process_manager: ProcessManagerRef,
127    request: ScanRequest,
128) -> error::Result<RecordBatch> {
129    let predicates = Predicates::from_scan_request(&Some(request));
130    let current_time = current_time_millis();
131    // todo(hl): find a way to extract user catalog to filter queries from other users.
132    let queries = process_manager.list_all_processes(None).await?;
133
134    let mut id_builder = StringVectorBuilder::with_capacity(queries.len());
135    let mut catalog_builder = StringVectorBuilder::with_capacity(queries.len());
136    let mut schemas_builder = StringVectorBuilder::with_capacity(queries.len());
137    let mut query_builder = StringVectorBuilder::with_capacity(queries.len());
138    let mut client_builder = StringVectorBuilder::with_capacity(queries.len());
139    let mut frontend_builder = StringVectorBuilder::with_capacity(queries.len());
140    let mut start_time_builder = TimestampMillisecondVectorBuilder::with_capacity(queries.len());
141    let mut elapsed_time_builder = DurationMillisecondVectorBuilder::with_capacity(queries.len());
142
143    for process in queries {
144        let display_id = DisplayProcessId {
145            server_addr: process.frontend.clone(),
146            id: process.id,
147        }
148        .to_string();
149        let schemas = process.schemas.join(",");
150        let id = Value::from(display_id);
151        let catalog = Value::from(process.catalog);
152        let schemas = Value::from(schemas);
153        let query = Value::from(process.query);
154        let client = Value::from(process.client);
155        let frontend = Value::from(process.frontend);
156        let start_timestamp = Value::from(Timestamp::new_millisecond(process.start_timestamp));
157        let elapsed_time = Value::from(Duration::new_millisecond(
158            current_time - process.start_timestamp,
159        ));
160        let row = [
161            (ID, &id),
162            (CATALOG, &catalog),
163            (SCHEMAS, &schemas),
164            (QUERY, &query),
165            (CLIENT, &client),
166            (FRONTEND, &frontend),
167            (START_TIMESTAMP, &start_timestamp),
168            (ELAPSED_TIME, &elapsed_time),
169        ];
170        if predicates.eval(&row) {
171            id_builder.push(id.as_string().as_deref());
172            catalog_builder.push(catalog.as_string().as_deref());
173            schemas_builder.push(schemas.as_string().as_deref());
174            query_builder.push(query.as_string().as_deref());
175            client_builder.push(client.as_string().as_deref());
176            frontend_builder.push(frontend.as_string().as_deref());
177            start_time_builder.push(start_timestamp.as_timestamp().map(|t| t.value().into()));
178            elapsed_time_builder.push(elapsed_time.as_duration().map(|d| d.value().into()));
179        }
180    }
181
182    RecordBatch::new(
183        InformationSchemaProcessList::schema(),
184        vec![
185            Arc::new(id_builder.finish()) as VectorRef,
186            Arc::new(catalog_builder.finish()) as VectorRef,
187            Arc::new(schemas_builder.finish()) as VectorRef,
188            Arc::new(query_builder.finish()) as VectorRef,
189            Arc::new(client_builder.finish()) as VectorRef,
190            Arc::new(frontend_builder.finish()) as VectorRef,
191            Arc::new(start_time_builder.finish()) as VectorRef,
192            Arc::new(elapsed_time_builder.finish()) as VectorRef,
193        ],
194    )
195    .context(error::CreateRecordBatchSnafu)
196}