arrow_flight/sql/metadata/
db_schemas.rs

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

arrow_flight/sql/metadata/
db_schemas.rs
