arrow_flight/sql/metadata/
catalogs.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
18use std::sync::{Arc, LazyLock};
19
20use arrow_array::{RecordBatch, StringArray};
21use arrow_schema::{DataType, Field, Schema, SchemaRef};
22
23use crate::error::Result;
24use crate::sql::CommandGetCatalogs;
25
26/// A builder for a [`CommandGetCatalogs`] response.
27///
28/// Builds rows like this:
29///
30/// * catalog_name: utf8,
31pub struct GetCatalogsBuilder {
32    catalogs: Vec<String>,
33}
34
35impl CommandGetCatalogs {
36    /// Create a builder suitable for constructing a response
37    pub fn into_builder(self) -> GetCatalogsBuilder {
38        self.into()
39    }
40}
41
42impl From<CommandGetCatalogs> for GetCatalogsBuilder {
43    fn from(_: CommandGetCatalogs) -> Self {
44        Self::new()
45    }
46}
47
48impl Default for GetCatalogsBuilder {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl GetCatalogsBuilder {
55    /// Create a new instance of [`GetCatalogsBuilder`]
56    pub fn new() -> Self {
57        Self {
58            catalogs: Vec::new(),
59        }
60    }
61
62    /// Append a row
63    pub fn append(&mut self, catalog_name: impl Into<String>) {
64        self.catalogs.push(catalog_name.into());
65    }
66
67    /// builds a `RecordBatch` with the correct schema for a
68    /// [`CommandGetCatalogs`] response
69    pub fn build(self) -> Result<RecordBatch> {
70        let Self { mut catalogs } = self;
71        catalogs.sort_unstable();
72
73        let batch = RecordBatch::try_new(
74            Arc::clone(&GET_CATALOG_SCHEMA),
75            vec![Arc::new(StringArray::from_iter_values(catalogs)) as _],
76        )?;
77
78        Ok(batch)
79    }
80
81    /// Returns the schema that will result from [`CommandGetCatalogs`]
82    ///
83    /// [`CommandGetCatalogs`]: crate::sql::CommandGetCatalogs
84    pub fn schema(&self) -> SchemaRef {
85        get_catalogs_schema()
86    }
87}
88
89fn get_catalogs_schema() -> SchemaRef {
90    Arc::clone(&GET_CATALOG_SCHEMA)
91}
92
93/// The schema for GetCatalogs
94static GET_CATALOG_SCHEMA: LazyLock<SchemaRef> = LazyLock::new(|| {
95    Arc::new(Schema::new(vec![Field::new(
96        "catalog_name",
97        DataType::Utf8,
98        false,
99    )]))
100});
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105
106    #[test]
107    fn test_catalogs_are_sorted() {
108        let batch = ["a_catalog", "c_catalog", "b_catalog"]
109            .into_iter()
110            .fold(GetCatalogsBuilder::new(), |mut builder, catalog| {
111                builder.append(catalog);
112                builder
113            })
114            .build()
115            .unwrap();
116        let catalogs = batch
117            .column(0)
118            .as_any()
119            .downcast_ref::<StringArray>()
120            .unwrap()
121            .iter()
122            .flatten()
123            .collect::<Vec<_>>();
124        assert!(catalogs.is_sorted());
125        assert_eq!(catalogs, ["a_catalog", "b_catalog", "c_catalog"]);
126    }
127}
arrow_flight/sql/metadata/catalogs.rs

arrow_flight/sql/metadata/
catalogs.rs
