catalog/system_schema/information_schema/
schemata.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_SCHEMATA_TABLE_ID;
19use common_error::ext::BoxedError;
20use common_meta::key::schema_name::SchemaNameKey;
21use common_recordbatch::adapter::RecordBatchStreamAdapter;
22use common_recordbatch::{RecordBatch, SendableRecordBatchStream};
23use datafusion::execution::TaskContext;
24use datafusion::physical_plan::SendableRecordBatchStream as DfSendableRecordBatchStream;
25use datafusion::physical_plan::stream::RecordBatchStreamAdapter as DfRecordBatchStreamAdapter;
26use datafusion::physical_plan::streaming::PartitionStream as DfPartitionStream;
27use datatypes::prelude::{ConcreteDataType, ScalarVectorBuilder, VectorRef};
28use datatypes::schema::{ColumnSchema, Schema, SchemaRef};
29use datatypes::value::Value;
30use datatypes::vectors::StringVectorBuilder;
31use snafu::{OptionExt, ResultExt};
32use store_api::storage::{ScanRequest, TableId};
33
34use crate::CatalogManager;
35use crate::error::{
36    CreateRecordBatchSnafu, InternalSnafu, Result, TableMetadataManagerSnafu,
37    UpgradeWeakCatalogManagerRefSnafu,
38};
39use crate::system_schema::SystemTable;
40use crate::system_schema::information_schema::{
41    InformationSchemaTableScope, InformationTable, Predicates, SCHEMATA,
42};
43use crate::system_schema::utils;
44
45pub const CATALOG_NAME: &str = "catalog_name";
46pub const SCHEMA_NAME: &str = "schema_name";
47const DEFAULT_CHARACTER_SET_NAME: &str = "default_character_set_name";
48const DEFAULT_COLLATION_NAME: &str = "default_collation_name";
49/// The database options
50pub const SCHEMA_OPTS: &str = "options";
51const INIT_CAPACITY: usize = 42;
52
53/// The `information_schema.schemata` table implementation.
54#[derive(Debug)]
55pub(super) struct InformationSchemaSchemata {
56    schema: SchemaRef,
57    catalog_name: String,
58    catalog_manager: Weak<dyn CatalogManager>,
59}
60
61impl InformationSchemaSchemata {
62    pub(super) fn new(catalog_name: String, catalog_manager: Weak<dyn CatalogManager>) -> Self {
63        Self {
64            schema: Self::schema(),
65            catalog_name,
66            catalog_manager,
67        }
68    }
69
70    pub(crate) fn schema() -> SchemaRef {
71        Arc::new(Schema::new(vec![
72            ColumnSchema::new(CATALOG_NAME, ConcreteDataType::string_datatype(), false),
73            ColumnSchema::new(SCHEMA_NAME, ConcreteDataType::string_datatype(), false),
74            ColumnSchema::new(
75                DEFAULT_CHARACTER_SET_NAME,
76                ConcreteDataType::string_datatype(),
77                false,
78            ),
79            ColumnSchema::new(
80                DEFAULT_COLLATION_NAME,
81                ConcreteDataType::string_datatype(),
82                false,
83            ),
84            ColumnSchema::new("sql_path", ConcreteDataType::string_datatype(), true),
85            ColumnSchema::new(SCHEMA_OPTS, ConcreteDataType::string_datatype(), true),
86        ]))
87    }
88
89    fn builder(&self) -> InformationSchemaSchemataBuilder {
90        InformationSchemaSchemataBuilder::new(
91            self.schema.clone(),
92            self.catalog_name.clone(),
93            self.catalog_manager.clone(),
94        )
95    }
96}
97
98impl InformationTable for InformationSchemaSchemata {
99    fn scope(&self) -> InformationSchemaTableScope {
100        InformationSchemaTableScope::Catalog
101    }
102}
103
104impl SystemTable for InformationSchemaSchemata {
105    fn table_id(&self) -> TableId {
106        INFORMATION_SCHEMA_SCHEMATA_TABLE_ID
107    }
108
109    fn table_name(&self) -> &'static str {
110        SCHEMATA
111    }
112
113    fn schema(&self) -> SchemaRef {
114        self.schema.clone()
115    }
116
117    fn to_stream(&self, request: ScanRequest) -> Result<SendableRecordBatchStream> {
118        let schema = self.schema.arrow_schema().clone();
119        let mut builder = self.builder();
120        let stream = Box::pin(DfRecordBatchStreamAdapter::new(
121            schema,
122            futures::stream::once(async move {
123                builder
124                    .make_schemata(Some(request))
125                    .await
126                    .map(|x| x.into_df_record_batch())
127                    .map_err(Into::into)
128            }),
129        ));
130        Ok(Box::pin(
131            RecordBatchStreamAdapter::try_new(stream)
132                .map_err(BoxedError::new)
133                .context(InternalSnafu)?,
134        ))
135    }
136}
137
138/// Builds the `information_schema.schemata` table row by row
139///
140/// Columns are based on <https://docs.pingcap.com/tidb/stable/information-schema-schemata>
141struct InformationSchemaSchemataBuilder {
142    schema: SchemaRef,
143    catalog_name: String,
144    catalog_manager: Weak<dyn CatalogManager>,
145
146    catalog_names: StringVectorBuilder,
147    schema_names: StringVectorBuilder,
148    charset_names: StringVectorBuilder,
149    collation_names: StringVectorBuilder,
150    sql_paths: StringVectorBuilder,
151    schema_options: StringVectorBuilder,
152}
153
154impl InformationSchemaSchemataBuilder {
155    fn new(
156        schema: SchemaRef,
157        catalog_name: String,
158        catalog_manager: Weak<dyn CatalogManager>,
159    ) -> Self {
160        Self {
161            schema,
162            catalog_name,
163            catalog_manager,
164            catalog_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
165            schema_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
166            charset_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
167            collation_names: StringVectorBuilder::with_capacity(INIT_CAPACITY),
168            sql_paths: StringVectorBuilder::with_capacity(INIT_CAPACITY),
169            schema_options: StringVectorBuilder::with_capacity(INIT_CAPACITY),
170        }
171    }
172
173    /// Construct the `information_schema.schemata` virtual table
174    async fn make_schemata(&mut self, request: Option<ScanRequest>) -> Result<RecordBatch> {
175        let catalog_name = self.catalog_name.clone();
176        let catalog_manager = self
177            .catalog_manager
178            .upgrade()
179            .context(UpgradeWeakCatalogManagerRefSnafu)?;
180        let table_metadata_manager = utils::table_meta_manager(&self.catalog_manager)?;
181        let predicates = Predicates::from_scan_request(&request);
182
183        for schema_name in catalog_manager.schema_names(&catalog_name, None).await? {
184            let opts = if let Some(table_metadata_manager) = &table_metadata_manager {
185                table_metadata_manager
186                    .schema_manager()
187                    .get(SchemaNameKey::new(&catalog_name, &schema_name))
188                    .await
189                    .context(TableMetadataManagerSnafu)?
190                    // information_schema is not available from this
191                    // table_metadata_manager and we return None
192                    .map(|schema_opts| format!("{}", schema_opts.into_inner()))
193            } else {
194                None
195            };
196
197            self.add_schema(
198                &predicates,
199                &catalog_name,
200                &schema_name,
201                opts.as_deref().unwrap_or(""),
202            );
203        }
204
205        self.finish()
206    }
207
208    fn add_schema(
209        &mut self,
210        predicates: &Predicates,
211        catalog_name: &str,
212        schema_name: &str,
213        schema_options: &str,
214    ) {
215        let row = [
216            (CATALOG_NAME, &Value::from(catalog_name)),
217            (SCHEMA_NAME, &Value::from(schema_name)),
218            (DEFAULT_CHARACTER_SET_NAME, &Value::from("utf8")),
219            (DEFAULT_COLLATION_NAME, &Value::from("utf8_bin")),
220            (SCHEMA_OPTS, &Value::from(schema_options)),
221        ];
222
223        if !predicates.eval(&row) {
224            return;
225        }
226
227        self.catalog_names.push(Some(catalog_name));
228        self.schema_names.push(Some(schema_name));
229        self.charset_names.push(Some("utf8"));
230        self.collation_names.push(Some("utf8_bin"));
231        self.sql_paths.push(None);
232        self.schema_options.push(Some(schema_options));
233    }
234
235    fn finish(&mut self) -> Result<RecordBatch> {
236        let columns: Vec<VectorRef> = vec![
237            Arc::new(self.catalog_names.finish()),
238            Arc::new(self.schema_names.finish()),
239            Arc::new(self.charset_names.finish()),
240            Arc::new(self.collation_names.finish()),
241            Arc::new(self.sql_paths.finish()),
242            Arc::new(self.schema_options.finish()),
243        ];
244        RecordBatch::new(self.schema.clone(), columns).context(CreateRecordBatchSnafu)
245    }
246}
247
248impl DfPartitionStream for InformationSchemaSchemata {
249    fn schema(&self) -> &ArrowSchemaRef {
250        self.schema.arrow_schema()
251    }
252
253    fn execute(&self, _: Arc<TaskContext>) -> DfSendableRecordBatchStream {
254        let schema = self.schema.arrow_schema().clone();
255        let mut builder = self.builder();
256        Box::pin(DfRecordBatchStreamAdapter::new(
257            schema,
258            futures::stream::once(async move {
259                builder
260                    .make_schemata(None)
261                    .await
262                    .map(|x| x.into_df_record_batch())
263                    .map_err(Into::into)
264            }),
265        ))
266    }
267}
catalog/system_schema/information_schema/schemata.rs

catalog/system_schema/information_schema/
schemata.rs
