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

catalog/system_schema/information_schema/
key_column_usage.rs
