catalog/system_schema/information_schema/
table_constraints.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_TABLE_CONSTRAINTS_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};
27use datatypes::scalars::ScalarVectorBuilder;
28use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
29use datatypes::value::Value;
30use datatypes::vectors::{StringVector, StringVectorBuilder, VectorRef};
31use futures::TryStreamExt;
32use snafu::{OptionExt, ResultExt};
33use store_api::storage::{ScanRequest, TableId};
34
35use crate::CatalogManager;
36use crate::error::{
37    CreateRecordBatchSnafu, InternalSnafu, Result, UpgradeWeakCatalogManagerRefSnafu,
38};
39use crate::information_schema::Predicates;
40use crate::information_schema::key_column_usage::{
41    CONSTRAINT_NAME_PRI, CONSTRAINT_NAME_TIME_INDEX,
42};
43use crate::system_schema::SystemTable;
44use crate::system_schema::information_schema::{
45    InformationSchemaTableScope, InformationTable, TABLE_CONSTRAINTS,
46};
47
48/// The `TABLE_CONSTRAINTS` table describes which tables have constraints.
49#[derive(Debug)]
50pub(super) struct InformationSchemaTableConstraints {
51    schema: SchemaRef,
52    catalog_name: String,
53    catalog_manager: Weak<dyn CatalogManager>,
54}
55
56const CONSTRAINT_CATALOG: &str = "constraint_catalog";
57const CONSTRAINT_SCHEMA: &str = "constraint_schema";
58const CONSTRAINT_NAME: &str = "constraint_name";
59const TABLE_SCHEMA: &str = "table_schema";
60const TABLE_NAME: &str = "table_name";
61const CONSTRAINT_TYPE: &str = "constraint_type";
62const ENFORCED: &str = "enforced";
63
64const INIT_CAPACITY: usize = 42;
65
66const TIME_INDEX_CONSTRAINT_TYPE: &str = "TIME INDEX";
67const PRI_KEY_CONSTRAINT_TYPE: &str = "PRIMARY KEY";
68
69impl InformationSchemaTableConstraints {
70    pub(super) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
71        Self {
72            schema: Self::schema(),
73            catalog_name,
74            catalog_manager,
75        }
76    }
77
78    fn schema() -> SchemaRef {
79        Arc::new(Schema::new(vec![
80            ColumnSchema::new(
81                CONSTRAINT_CATALOG,
82                ConcreteDataType::string_datatype(),
83                false,
84            ),
85            ColumnSchema::new(
86                CONSTRAINT_SCHEMA,
87                ConcreteDataType::string_datatype(),
88                false,
89            ),
90            ColumnSchema::new(CONSTRAINT_NAME, ConcreteDataType::string_datatype(), false),
91            ColumnSchema::new(TABLE_SCHEMA, ConcreteDataType::string_datatype(), false),
92            ColumnSchema::new(TABLE_NAME, ConcreteDataType::string_datatype(), false),
93            ColumnSchema::new(CONSTRAINT_TYPE, ConcreteDataType::string_datatype(), false),
94            ColumnSchema::new(ENFORCED, ConcreteDataType::string_datatype(), false),
95        ]))
96    }
97
98    fn builder(&self) -> InformationSchemaTableConstraintsBuilder {
99        InformationSchemaTableConstraintsBuilder::new(
100            self.schema.clone(),
101            self.catalog_name.clone(),
102            self.catalog_manager.clone(),
103        )
104    }
105}
106
107impl InformationTable for InformationSchemaTableConstraints {
108    fn scope(&self) -> InformationSchemaTableScope {
109        InformationSchemaTableScope::Catalog
110    }
111}
112
113impl SystemTable for InformationSchemaTableConstraints {
114    fn table_id(&self) -> TableId {
115        INFORMATION_SCHEMA_TABLE_CONSTRAINTS_TABLE_ID
116    }
117
118    fn table_name(&self) -> &'static str {
119        TABLE_CONSTRAINTS
120    }
121
122    fn schema(&self) -> SchemaRef {
123        self.schema.clone()
124    }
125
126    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
127        let schema = self.schema.arrow_schema().clone();
128        let mut builder = self.builder();
129        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
130            schema,
131            futures::stream::once(async move {
132                builder
133                    .make_table_constraints(Some(request))
134                    .await
135                    .map(|x| x.into_df_record_batch())
136                    .map_err(Into::into)
137            }),
138        ));
139        Ok(Box::pin(
140            RecordBatchStreamAdapter::try_new(stream)
141                .map_err(BoxedError::new)
142                .context(InternalSnafu)?,
143        ))
144    }
145}
146
147struct InformationSchemaTableConstraintsBuilder {
148    schema: SchemaRef,
149    catalog_name: String,
150    catalog_manager: Weak<dyn CatalogManager>,
151
152    constraint_schemas: StringVectorBuilder,
153    constraint_names: StringVectorBuilder,
154    table_schemas: StringVectorBuilder,
155    table_names: StringVectorBuilder,
156    constraint_types: StringVectorBuilder,
157}
158
159impl InformationSchemaTableConstraintsBuilder {
160    fn new(
161        schema: SchemaRef,
162        catalog_name: String,
163        catalog_manager: Weak<dyn CatalogManager>,
164    ) -> Self {
165        Self {
166            schema,
167            catalog_name,
168            catalog_manager,
169            constraint_schemas: StringVectorBuilder::with_capacity(INIT_CAPACITY),
170            constraint_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
171            table_schemas: StringVectorBuilder::with_capacity(INIT_CAPACITY),
172            table_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
173            constraint_types: StringVectorBuilder::with_capacity(INIT_CAPACITY),
174        }
175    }
176
177    /// Construct the `information_schema.table_constraints` virtual table
178    async fn make_table_constraints(
179        &mut self,
180        request: Option<ScanRequest>,
181    ) -> Result<RecordBatch> {
182        let catalog_name = self.catalog_name.clone();
183        let catalog_manager = self
184            .catalog_manager
185            .upgrade()
186            .context(UpgradeWeakCatalogManagerRefSnafu)?;
187        let predicates = Predicates::from_scan_request(&request);
188
189        for schema_name in catalog_manager.schema_names(&catalog_name, None).await? {
190            let mut stream = catalog_manager.tables(&catalog_name, &schema_name, None);
191
192            while let Some(table) = stream.try_next().await? {
193                let keys = &table.table_info().meta.primary_key_indices;
194                let schema = table.schema();
195
196                if schema.timestamp_index().is_some() {
197                    self.add_table_constraint(
198                        &predicates,
199                        &schema_name,
200                        CONSTRAINT_NAME_TIME_INDEX,
201                        &schema_name,
202                        &table.table_info().name,
203                        TIME_INDEX_CONSTRAINT_TYPE,
204                    );
205                }
206
207                if !keys.is_empty() {
208                    self.add_table_constraint(
209                        &predicates,
210                        &schema_name,
211                        CONSTRAINT_NAME_PRI,
212                        &schema_name,
213                        &table.table_info().name,
214                        PRI_KEY_CONSTRAINT_TYPE,
215                    );
216                }
217            }
218        }
219
220        self.finish()
221    }
222
223    fn add_table_constraint(
224        &mut self,
225        predicates: &Predicates,
226        constraint_schema: &str,
227        constraint_name: &str,
228        table_schema: &str,
229        table_name: &str,
230        constraint_type: &str,
231    ) {
232        let row = [
233            (CONSTRAINT_SCHEMA, &Value::from(constraint_schema)),
234            (CONSTRAINT_NAME, &Value::from(constraint_name)),
235            (TABLE_SCHEMA, &Value::from(table_schema)),
236            (TABLE_NAME, &Value::from(table_name)),
237            (CONSTRAINT_TYPE, &Value::from(constraint_type)),
238        ];
239
240        if !predicates.eval(&row) {
241            return;
242        }
243
244        self.constraint_schemas.push(Some(constraint_schema));
245        self.constraint_names.push(Some(constraint_name));
246        self.table_schemas.push(Some(table_schema));
247        self.table_names.push(Some(table_name));
248        self.constraint_types.push(Some(constraint_type));
249    }
250
251    fn finish(&mut self) -> Result<RecordBatch> {
252        let rows_num = self.constraint_names.len();
253
254        let constraint_catalogs = Arc::new(StringVector::from(vec!["def"; rows_num]));
255        let enforceds = Arc::new(StringVector::from(vec!["YES"; rows_num]));
256
257        let columns: Vec<VectorRef> = vec![
258            constraint_catalogs,
259            Arc::new(self.constraint_schemas.finish()),
260            Arc::new(self.constraint_names.finish()),
261            Arc::new(self.table_schemas.finish()),
262            Arc::new(self.table_names.finish()),
263            Arc::new(self.constraint_types.finish()),
264            enforceds,
265        ];
266
267        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
268    }
269}
270
271impl DfPartitionStream for InformationSchemaTableConstraints {
272    fn schema(&self) -> &ArrowSchemaRef {
273        self.schema.arrow_schema()
274    }
275
276    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
277        let schema = self.schema.arrow_schema().clone();
278        let mut builder = self.builder();
279        Box::pin(DfRecordBatchStreamAdapter::new(
280            schema,
281            futures::stream::once(async move {
282                builder
283                    .make_table_constraints(None)
284                    .await
285                    .map(|x| x.into_df_record_batch())
286                    .map_err(Into::into)
287            }),
288        ))
289    }
290}
catalog/system_schema/information_schema/table_constraints.rs

catalog/system_schema/information_schema/
table_constraints.rs
