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

arrow_flight/sql/metadata/
catalogs.rs
