arrow_flight/sql/metadata/
db_schemas.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`GetDbSchemasBuilder`] for building responses to [`CommandGetDbSchemas`] queries.
19//!
20//! [`CommandGetDbSchemas`]: crate::sql::CommandGetDbSchemas
21
22use std::sync::{Arc, LazyLock};
23
24use arrow_arith::boolean::and;
25use arrow_array::{ArrayRef, RecordBatch, StringArray, builder::StringBuilder};
26use arrow_ord::cmp::eq;
27use arrow_schema::{DataType, Field, Schema, SchemaRef};
28use arrow_select::{filter::filter_record_batch, take::take};
29use arrow_string::like::like;
30
31use super::lexsort_to_indices;
32use crate::error::*;
33use crate::sql::CommandGetDbSchemas;
34
35/// A builder for a [`CommandGetDbSchemas`] response.
36///
37/// Builds rows like this:
38///
39/// * catalog_name: utf8,
40/// * db_schema_name: utf8 not null
41pub struct GetDbSchemasBuilder {
42    // Specifies the Catalog to search for the tables.
43    // - An empty string retrieves those without a catalog.
44    // - If omitted the catalog name is not used to narrow the search.
45    catalog_filter: Option<String>,
46    // Optional filters to apply
47    db_schema_filter_pattern: Option<String>,
48    // array builder for catalog names
49    catalog_name: StringBuilder,
50    // array builder for schema names
51    db_schema_name: StringBuilder,
52}
53
54impl CommandGetDbSchemas {
55    /// Create a builder suitable for constructing a response
56    pub fn into_builder(self) -> GetDbSchemasBuilder {
57        self.into()
58    }
59}
60
61impl From<CommandGetDbSchemas> for GetDbSchemasBuilder {
62    fn from(value: CommandGetDbSchemas) -> Self {
63        Self::new(value.catalog, value.db_schema_filter_pattern)
64    }
65}
66
67impl GetDbSchemasBuilder {
68    /// Create a new instance of [`GetDbSchemasBuilder`]
69    ///
70    /// # Parameters
71    ///
72    /// - `catalog`:  Specifies the Catalog to search for the tables.
73    ///   - An empty string retrieves those without a catalog.
74    ///   - If omitted the catalog name is not used to narrow the search.
75    /// - `db_schema_filter_pattern`: Specifies a filter pattern for schemas to search for.
76    ///   When no pattern is provided, the pattern will not be used to narrow the search.
77    ///   In the pattern string, two special characters can be used to denote matching rules:
78    ///     - "%" means to match any substring with 0 or more characters.
79    ///     - "_" means to match any one character.
80    ///
81    /// [`CommandGetDbSchemas`]: crate::sql::CommandGetDbSchemas
82    pub fn new(
83        catalog: Option<impl Into<String>>,
84        db_schema_filter_pattern: Option<impl Into<String>>,
85    ) -> Self {
86        Self {
87            catalog_filter: catalog.map(|v| v.into()),
88            db_schema_filter_pattern: db_schema_filter_pattern.map(|v| v.into()),
89            catalog_name: StringBuilder::new(),
90            db_schema_name: StringBuilder::new(),
91        }
92    }
93
94    /// Append a row
95    ///
96    /// In case the catalog should be considered as empty, pass in an empty string '""'.
97    pub fn append(&mut self, catalog_name: impl AsRef<str>, schema_name: impl AsRef<str>) {
98        self.catalog_name.append_value(catalog_name);
99        self.db_schema_name.append_value(schema_name);
100    }
101
102    /// builds a `RecordBatch` with the correct schema for a `CommandGetDbSchemas` response
103    pub fn build(self) -> Result<RecordBatch> {
104        let schema = self.schema();
105        let Self {
106            catalog_filter,
107            db_schema_filter_pattern,
108            mut catalog_name,
109            mut db_schema_name,
110        } = self;
111
112        // Make the arrays
113        let catalog_name = catalog_name.finish();
114        let db_schema_name = db_schema_name.finish();
115
116        let mut filters = vec![];
117
118        if let Some(db_schema_filter_pattern) = db_schema_filter_pattern {
119            // use like kernel to get wildcard matching
120            let scalar = StringArray::new_scalar(db_schema_filter_pattern);
121            filters.push(like(&db_schema_name, &scalar)?)
122        }
123
124        if let Some(catalog_filter_name) = catalog_filter {
125            let scalar = StringArray::new_scalar(catalog_filter_name);
126            filters.push(eq(&catalog_name, &scalar)?);
127        }
128
129        // `AND` any filters together
130        let mut total_filter = None;
131        while let Some(filter) = filters.pop() {
132            let new_filter = match total_filter {
133                Some(total_filter) => and(&total_filter, &filter)?,
134                None => filter,
135            };
136            total_filter = Some(new_filter);
137        }
138
139        let batch = RecordBatch::try_new(
140            schema,
141            vec![
142                Arc::new(catalog_name) as ArrayRef,
143                Arc::new(db_schema_name) as ArrayRef,
144            ],
145        )?;
146
147        // Apply the filters if needed
148        let filtered_batch = if let Some(filter) = total_filter {
149            filter_record_batch(&batch, &filter)?
150        } else {
151            batch
152        };
153
154        // Order filtered results by catalog_name, then db_schema_name
155        let indices = lexsort_to_indices(filtered_batch.columns());
156        let columns = filtered_batch
157            .columns()
158            .iter()
159            .map(|c| take(c, &indices, None))
160            .collect::<std::result::Result<Vec<_>, _>>()?;
161
162        Ok(RecordBatch::try_new(filtered_batch.schema(), columns)?)
163    }
164
165    /// Return the schema of the RecordBatch that will be returned
166    /// from [`CommandGetDbSchemas`]
167    pub fn schema(&self) -> SchemaRef {
168        get_db_schemas_schema()
169    }
170}
171
172fn get_db_schemas_schema() -> SchemaRef {
173    Arc::clone(&GET_DB_SCHEMAS_SCHEMA)
174}
175
176/// The schema for GetDbSchemas
177static GET_DB_SCHEMAS_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
178    Arc::new(Schema::new(vec![
179        Field::new("catalog_name", DataType::Utf8, true),
180        Field::new("db_schema_name", DataType::Utf8, false),
181    ]))
182});
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use arrow_array::{StringArray, UInt32Array};
188
189    fn get_ref_batch() -> RecordBatch {
190        RecordBatch::try_new(
191            get_db_schemas_schema(),
192            vec![
193                Arc::new(StringArray::from(vec![
194                    "a_catalog",
195                    "a_catalog",
196                    "b_catalog",
197                    "b_catalog",
198                ])) as ArrayRef,
199                Arc::new(StringArray::from(vec![
200                    "a_schema", "b_schema", "a_schema", "b_schema",
201                ])) as ArrayRef,
202            ],
203        )
204        .unwrap()
205    }
206
207    #[test]
208    fn test_schemas_are_filtered() {
209        let ref_batch = get_ref_batch();
210
211        let mut builder = GetDbSchemasBuilder::new(None::<String>, None::<String>);
212        builder.append("a_catalog", "a_schema");
213        builder.append("a_catalog", "b_schema");
214        builder.append("b_catalog", "a_schema");
215        builder.append("b_catalog", "b_schema");
216        let schema_batch = builder.build().unwrap();
217
218        assert_eq!(schema_batch, ref_batch);
219
220        let mut builder = GetDbSchemasBuilder::new(None::<String>, Some("a%"));
221        builder.append("a_catalog", "a_schema");
222        builder.append("a_catalog", "b_schema");
223        builder.append("b_catalog", "a_schema");
224        builder.append("b_catalog", "b_schema");
225        let schema_batch = builder.build().unwrap();
226
227        let indices = UInt32Array::from(vec![0, 2]);
228        let ref_filtered = RecordBatch::try_new(
229            get_db_schemas_schema(),
230            ref_batch
231                .columns()
232                .iter()
233                .map(|c| take(c, &indices, None))
234                .collect::<std::result::Result<Vec<_>, _>>()
235                .unwrap(),
236        )
237        .unwrap();
238
239        assert_eq!(schema_batch, ref_filtered);
240    }
241
242    #[test]
243    fn test_schemas_are_sorted() {
244        let ref_batch = get_ref_batch();
245
246        let mut builder = GetDbSchemasBuilder::new(None::<String>, None::<String>);
247        builder.append("a_catalog", "b_schema");
248        builder.append("b_catalog", "a_schema");
249        builder.append("a_catalog", "a_schema");
250        builder.append("b_catalog", "b_schema");
251        let schema_batch = builder.build().unwrap();
252
253        assert_eq!(schema_batch, ref_batch)
254    }
255
256    #[test]
257    fn test_builder_from_query() {
258        let ref_batch = get_ref_batch();
259        let query = CommandGetDbSchemas {
260            catalog: Some("a_catalog".into()),
261            db_schema_filter_pattern: Some("b%".into()),
262        };
263
264        let mut builder = query.into_builder();
265        builder.append("a_catalog", "a_schema");
266        builder.append("a_catalog", "b_schema");
267        builder.append("b_catalog", "a_schema");
268        builder.append("b_catalog", "b_schema");
269        let schema_batch = builder.build().unwrap();
270
271        let indices = UInt32Array::from(vec![1]);
272        let ref_filtered = RecordBatch::try_new(
273            get_db_schemas_schema(),
274            ref_batch
275                .columns()
276                .iter()
277                .map(|c| take(c, &indices, None))
278                .collect::<std::result::Result<Vec<_>, _>>()
279                .unwrap(),
280        )
281        .unwrap();
282
283        assert_eq!(schema_batch, ref_filtered);
284    }
285}
arrow_flight/sql/metadata/db_schemas.rs

arrow_flight/sql/metadata/
db_schemas.rs
