catalog/system_schema/information_schema/
key_column_usage.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_KEY_COLUMN_USAGE_TABLE_ID;
19use common_error::ext::BoxedError;
20use common_recordbatch::adapter::RecordBatchStreamAdapter;
21use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
22use datafusion::execution::TaskContext;
23use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
24use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
25use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
26use datatypes::prelude::{ConcreteDataType, MutableVector, ScalarVectorBuilder, VectorRef};
27use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
28use datatypes::value::Value;
29use datatypes::vectors::{StringVector, StringVectorBuilder, UInt32VectorBuilder};
30use futures_util::TryStreamExt;
31use snafu::{OptionExt, ResultExt};
32use store_api::storage::{ScanRequest, TableId};
33
34use crate::CatalogManager;
35use crate::error::{
36    CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu,
37};
38use crate::system_schema::SystemTable;
39use crate::system_schema::information_schema::{
40    InformationSchemaTableScope, InformationTable, KEY_COLUMN_USAGE, Predicates,
41    primary_key_encoding_index_type,
42};
43
44pub const CONSTRAINT_SCHEMA: &str = "constraint_schema";
45pub const CONSTRAINT_NAME: &str = "constraint_name";
46// It's always `def` in MySQL
47pub const TABLE_CATALOG: &str = "table_catalog";
48// The real catalog name for this key column.
49pub const REAL_TABLE_CATALOG: &str = "real_table_catalog";
50pub const TABLE_SCHEMA: &str = "table_schema";
51pub const TABLE_NAME: &str = "table_name";
52pub const COLUMN_NAME: &str = "column_name";
53pub const ORDINAL_POSITION: &str = "ordinal_position";
54/// The type of the index.
55pub const GREPTIME_INDEX_TYPE: &str = "greptime_index_type";
56const INIT_CAPACITY: usize = 42;
57
58/// Time index constraint name
59pub(crate) const CONSTRAINT_NAME_TIME_INDEX: &str = "TIME INDEX";
60
61/// Primary key constraint name
62pub(crate) const CONSTRAINT_NAME_PRI: &str = "PRIMARY";
63
64/// Inverted index name
65pub(crate) const CONSTRAINT_NAME_INVERTED_INDEX: &str = "INVERTED INDEX";
66
67/// Fulltext index name
68pub(crate) const CONSTRAINT_NAME_FULLTEXT_INDEX: &str = "FULLTEXT INDEX";
69
70/// Skipping index name
71pub(crate) const CONSTRAINT_NAME_SKIPPING_INDEX: &str = "SKIPPING INDEX";
72
73/// The virtual table implementation for `information_schema.KEY_COLUMN_USAGE`.
74///
75/// Provides an extra column `greptime_index_type` for the index type of the key column.
76#[derive(Debug)]
77pub(super) struct InformationSchemaKeyColumnUsage {
78    schema: SchemaRef,
79    catalog_name: String,
80    catalog_manager: Weak<dyn CatalogManager>,
81}
82
83impl InformationSchemaKeyColumnUsage {
84    pub(super) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
85        Self {
86            schema: Self::schema(),
87            catalog_name,
88            catalog_manager,
89        }
90    }
91
92    pub(crate) fn schema() -> SchemaRef {
93        Arc::new(Schema::new(vec![
94            ColumnSchema::new(
95                "constraint_catalog",
96                ConcreteDataType::string_datatype(),
97                false,
98            ),
99            ColumnSchema::new(
100                CONSTRAINT_SCHEMA,
101                ConcreteDataType::string_datatype(),
102                false,
103            ),
104            ColumnSchema::new(CONSTRAINT_NAME, ConcreteDataType::string_datatype(), false),
105            ColumnSchema::new(TABLE_CATALOG, ConcreteDataType::string_datatype(), false),
106            ColumnSchema::new(
107                REAL_TABLE_CATALOG,
108                ConcreteDataType::string_datatype(),
109                false,
110            ),
111            ColumnSchema::new(TABLE_SCHEMA, ConcreteDataType::string_datatype(), false),
112            ColumnSchema::new(TABLE_NAME, ConcreteDataType::string_datatype(), false),
113            ColumnSchema::new(COLUMN_NAME, ConcreteDataType::string_datatype(), false),
114            ColumnSchema::new(ORDINAL_POSITION, ConcreteDataType::uint32_datatype(), false),
115            ColumnSchema::new(
116                "position_in_unique_constraint",
117                ConcreteDataType::uint32_datatype(),
118                true,
119            ),
120            ColumnSchema::new(
121                "referenced_table_schema",
122                ConcreteDataType::string_datatype(),
123                true,
124            ),
125            ColumnSchema::new(
126                "referenced_table_name",
127                ConcreteDataType::string_datatype(),
128                true,
129            ),
130            ColumnSchema::new(
131                "referenced_column_name",
132                ConcreteDataType::string_datatype(),
133                true,
134            ),
135            ColumnSchema::new(
136                GREPTIME_INDEX_TYPE,
137                ConcreteDataType::string_datatype(),
138                true,
139            ),
140        ]))
141    }
142
143    fn builder(&self) -> InformationSchemaKeyColumnUsageBuilder {
144        InformationSchemaKeyColumnUsageBuilder::new(
145            self.schema.clone(),
146            self.catalog_name.clone(),
147            self.catalog_manager.clone(),
148        )
149    }
150}
151
152impl InformationTable for InformationSchemaKeyColumnUsage {
153    fn scope(&self) -> InformationSchemaTableScope {
154        InformationSchemaTableScope::Catalog
155    }
156}
157
158impl SystemTable for InformationSchemaKeyColumnUsage {
159    fn table_id(&self) -> TableId {
160        INFORMATION_SCHEMA_KEY_COLUMN_USAGE_TABLE_ID
161    }
162
163    fn table_name(&self) -> &'static str {
164        KEY_COLUMN_USAGE
165    }
166
167    fn schema(&self) -> SchemaRef {
168        self.schema.clone()
169    }
170
171    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
172        let schema = self.schema.arrow_schema().clone();
173        let mut builder = self.builder();
174        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
175            schema,
176            futures::stream::once(async move {
177                builder
178                    .make_key_column_usage(Some(request))
179                    .await
180                    .map(|x| x.into_df_record_batch())
181                    .map_err(Into::into)
182            }),
183        ));
184        Ok(Box::pin(
185            RecordBatchStreamAdapter::try_new(stream)
186                .map_err(BoxedError::new)
187                .context(InternalSnafu)?,
188        ))
189    }
190}
191
192/// Builds the `information_schema.KEY_COLUMN_USAGE` table row by row
193///
194/// Columns are based on <https://dev.mysql.com/doc/refman/8.2/en/information-schema-key-column-usage-table.html>
195struct InformationSchemaKeyColumnUsageBuilder {
196    schema: SchemaRef,
197    catalog_name: String,
198    catalog_manager: Weak<dyn CatalogManager>,
199
200    constraint_catalog: StringVectorBuilder,
201    constraint_schema: StringVectorBuilder,
202    constraint_name: StringVectorBuilder,
203    table_catalog: StringVectorBuilder,
204    real_table_catalog: StringVectorBuilder,
205    table_schema: StringVectorBuilder,
206    table_name: StringVectorBuilder,
207    column_name: StringVectorBuilder,
208    ordinal_position: UInt32VectorBuilder,
209    position_in_unique_constraint: UInt32VectorBuilder,
210    greptime_index_type: StringVectorBuilder,
211}
212
213impl InformationSchemaKeyColumnUsageBuilder {
214    fn new(
215        schema: SchemaRef,
216        catalog_name: String,
217        catalog_manager: Weak<dyn CatalogManager>,
218    ) -> Self {
219        Self {
220            schema,
221            catalog_name,
222            catalog_manager,
223            constraint_catalog: StringVectorBuilder::with_capacity(INIT_CAPACITY),
224            constraint_schema: StringVectorBuilder::with_capacity(INIT_CAPACITY),
225            constraint_name: StringVectorBuilder::with_capacity(INIT_CAPACITY),
226            table_catalog: StringVectorBuilder::with_capacity(INIT_CAPACITY),
227            real_table_catalog: StringVectorBuilder::with_capacity(INIT_CAPACITY),
228            table_schema: StringVectorBuilder::with_capacity(INIT_CAPACITY),
229            table_name: StringVectorBuilder::with_capacity(INIT_CAPACITY),
230            column_name: StringVectorBuilder::with_capacity(INIT_CAPACITY),
231            ordinal_position: UInt32VectorBuilder::with_capacity(INIT_CAPACITY),
232            position_in_unique_constraint: UInt32VectorBuilder::with_capacity(INIT_CAPACITY),
233            greptime_index_type: StringVectorBuilder::with_capacity(INIT_CAPACITY),
234        }
235    }
236
237    /// Construct the `information_schema.KEY_COLUMN_USAGE` virtual table
238    async fn make_key_column_usage(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
239        let catalog_name = self.catalog_name.clone();
240        let catalog_manager = self
241            .catalog_manager
242            .upgrade()
243            .context(UpgradeWeakCatalogManagerRefSnafu)?;
244        let predicates = Predicates::from_scan_request(&request);
245
246        for schema_name in catalog_manager.schema_names(&catalog_name, None).await? {
247            let mut stream = catalog_manager.tables(&catalog_name, &schema_name, None);
248
249            while let Some(table) = stream.try_next().await? {
250                let table_info = table.table_info();
251                let table_name = &table_info.name;
252                let keys = &table_info.meta.primary_key_indices;
253                let schema = table.schema();
254                let primary_key_encoding =
255                    primary_key_encoding_index_type(&table_info.meta.options.extra_options);
256
257                for (idx, column) in schema.column_schemas().iter().enumerate() {
258                    if column.is_time_index() {
259                        self.add_key_column_usage(
260                            &predicates,
261                            &schema_name,
262                            CONSTRAINT_NAME_TIME_INDEX,
263                            &catalog_name,
264                            &schema_name,
265                            table_name,
266                            &column.name,
267                            1, //always 1 for time index
268                            "",
269                        );
270                    }
271                    if let Some(pk_seq) = keys.iter().position(|k| *k == idx) {
272                        self.add_key_column_usage(
273                            &predicates,
274                            &schema_name,
275                            CONSTRAINT_NAME_PRI,
276                            &catalog_name,
277                            &schema_name,
278                            table_name,
279                            &column.name,
280                            pk_seq as u32 + 1,
281                            primary_key_encoding,
282                        );
283                    }
284                }
285            }
286        }
287
288        self.finish()
289    }
290
291    // TODO(dimbtp): Foreign key constraint has not `None` value for last 4
292    // fields, but it is not supported yet.
293    #[allow(clippy::too_many_arguments)]
294    fn add_key_column_usage(
295        &mut self,
296        predicates: &Predicates,
297        constraint_schema: &str,
298        constraint_name: &str,
299        table_catalog: &str,
300        table_schema: &str,
301        table_name: &str,
302        column_name: &str,
303        ordinal_position: u32,
304        index_types: &str,
305    ) {
306        let row = [
307            (CONSTRAINT_SCHEMA, &Value::from(constraint_schema)),
308            (CONSTRAINT_NAME, &Value::from(constraint_name)),
309            (REAL_TABLE_CATALOG, &Value::from(table_catalog)),
310            (TABLE_SCHEMA, &Value::from(table_schema)),
311            (TABLE_NAME, &Value::from(table_name)),
312            (COLUMN_NAME, &Value::from(column_name)),
313            (ORDINAL_POSITION, &Value::from(ordinal_position)),
314            (GREPTIME_INDEX_TYPE, &Value::from(index_types)),
315        ];
316
317        if !predicates.eval(&row) {
318            return;
319        }
320
321        self.constraint_catalog.push(Some("def"));
322        self.constraint_schema.push(Some(constraint_schema));
323        self.constraint_name.push(Some(constraint_name));
324        self.table_catalog.push(Some("def"));
325        self.real_table_catalog.push(Some(table_catalog));
326        self.table_schema.push(Some(table_schema));
327        self.table_name.push(Some(table_name));
328        self.column_name.push(Some(column_name));
329        self.ordinal_position.push(Some(ordinal_position));
330        self.position_in_unique_constraint.push(None);
331        self.greptime_index_type.push(Some(index_types));
332    }
333
334    fn finish(&mut self) -> Result<RecordBatch> {
335        let rows_num = self.table_catalog.len();
336
337        let null_string_vector = Arc::new(StringVector::from(vec![None as Option<&str>; rows_num]));
338        let columns: Vec<VectorRef> = vec![
339            Arc::new(self.constraint_catalog.finish()),
340            Arc::new(self.constraint_schema.finish()),
341            Arc::new(self.constraint_name.finish()),
342            Arc::new(self.table_catalog.finish()),
343            Arc::new(self.real_table_catalog.finish()),
344            Arc::new(self.table_schema.finish()),
345            Arc::new(self.table_name.finish()),
346            Arc::new(self.column_name.finish()),
347            Arc::new(self.ordinal_position.finish()),
348            Arc::new(self.position_in_unique_constraint.finish()),
349            null_string_vector.clone(),
350            null_string_vector.clone(),
351            null_string_vector,
352            Arc::new(self.greptime_index_type.finish()),
353        ];
354        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
355    }
356}
357
358impl DfPartitionStream for InformationSchemaKeyColumnUsage {
359    fn schema(&self) -> &ArrowSchemaRef {
360        self.schema.arrow_schema()
361    }
362
363    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
364        let schema = self.schema.arrow_schema().clone();
365        let mut builder = self.builder();
366        Box::pin(DfRecordBatchStreamAdapter::new(
367            schema,
368            futures::stream::once(async move {
369                builder
370                    .make_key_column_usage(None)
371                    .await
372                    .map(|x| x.into_df_record_batch())
373                    .map_err(Into::into)
374            }),
375        ))
376    }
377}
catalog/system_schema/information_schema/key_column_usage.rs

catalog/system_schema/information_schema/
key_column_usage.rs
