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

arrow_flight/sql/metadata/
catalogs.rs
