feat: add catalog/schema subcommands to flight_sql_client. (#6332)

* feat: add catalog/schema subcommands to flight_sql_client.

With this change basic commands are added to query the catalogs and
schemas of a Flight SQL server.

* fix: adds tests for flight_sql_client cli

Additionally adds a builder pattern for the CommandGetTableTypes similar
to CommandGetDbSchemas, while its implementation is trivial it helps to
have a pattern to follow when implementing the command.

* fix: add default to GetTableTypesBuilder

arrow-flight/src/bin/flight_sql_client.rs

@@ -20,7 +20,10 @@ use std::{sync::Arc, time::Duration};
 use anyhow::{bail, Context, Result};
 use arrow_array::{ArrayRef, Datum, RecordBatch, StringArray};
 use arrow_cast::{cast_with_options, pretty::pretty_format_batches, CastOptions};
-use arrow_flight::{sql::client::FlightSqlServiceClient, FlightInfo};
+use arrow_flight::{
+    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetTables},
+    FlightInfo,
+};
 use arrow_schema::Schema;
 use clap::{Parser, Subcommand};
 use futures::TryStreamExt;
@@ -111,6 +114,51 @@ struct Args {
 /// Different available commands.
 #[derive(Debug, Subcommand)]
 enum Command {
+    /// Get catalogs.
+    Catalogs,
+    /// Get db schemas for a catalog.
+    DbSchemas {
+        /// Name of a catalog.
+        ///
+        /// Required.
+        catalog: String,
+        /// Specifies a filter pattern for schemas to search for.
+        /// When no schema_filter is provided, the pattern will not be used to narrow the search.
+        /// In the pattern string, two special characters can be used to denote matching rules:
+        ///     - "%" means to match any substring with 0 or more characters.
+        ///     - "_" means to match any one character.
+        #[clap(short, long)]
+        db_schema_filter: Option<String>,
+    },
+    /// Get tables for a catalog.
+    Tables {
+        /// Name of a catalog.
+        ///
+        /// Required.
+        catalog: String,
+        /// Specifies a filter pattern for schemas to search for.
+        /// When no schema_filter is provided, the pattern will not be used to narrow the search.
+        /// In the pattern string, two special characters can be used to denote matching rules:
+        ///     - "%" means to match any substring with 0 or more characters.
+        ///     - "_" means to match any one character.
+        #[clap(short, long)]
+        db_schema_filter: Option<String>,
+        /// Specifies a filter pattern for tables to search for.
+        /// When no table_filter is provided, all tables matching other filters are searched.
+        /// In the pattern string, two special characters can be used to denote matching rules:
+        ///     - "%" means to match any substring with 0 or more characters.
+        ///     - "_" means to match any one character.
+        #[clap(short, long)]
+        table_filter: Option<String>,
+        /// Specifies a filter of table types which must match.
+        /// The table types depend on vendor/implementation. It is usually used to separate tables from views or system tables.
+        /// TABLE, VIEW, and SYSTEM TABLE are commonly supported.
+        #[clap(long)]
+        table_types: Vec<String>,
+    },
+    /// Get table types.
+    TableTypes,
+
     /// Execute given statement.
     StatementQuery {
         /// SQL query.
@@ -150,6 +198,36 @@ async fn main() -> Result<()> {
         .context("setup client")?;
 
     let flight_info = match args.cmd {
+        Command::Catalogs => client.get_catalogs().await.context("get catalogs")?,
+        Command::DbSchemas {
+            catalog,
+            db_schema_filter,
+        } => client
+            .get_db_schemas(CommandGetDbSchemas {
+                catalog: Some(catalog),
+                db_schema_filter_pattern: db_schema_filter,
+            })
+            .await
+            .context("get db schemas")?,
+        Command::Tables {
+            catalog,
+            db_schema_filter,
+            table_filter,
+            table_types,
+        } => client
+            .get_tables(CommandGetTables {
+                catalog: Some(catalog),
+                db_schema_filter_pattern: db_schema_filter,
+                table_name_filter_pattern: table_filter,
+                table_types,
+                // Schema is returned as ipc encoded bytes.
+                // We do not support returning the schema as there is no trivial mechanism
+                // to display the information to the user.
+                include_schema: false,
+            })
+            .await
+            .context("get tables")?,
+        Command::TableTypes => client.get_table_types().await.context("get table types")?,
         Command::StatementQuery { query } => client
             .execute(query, None)
             .await

arrow-flight/src/sql/metadata/mod.rs

@@ -33,6 +33,7 @@
 mod catalogs;
 mod db_schemas;
 mod sql_info;
+mod table_types;
 mod tables;
 mod xdbc_info;
 

arrow-flight/src/sql/metadata/table_types.rs

@@ -0,0 +1,158 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`GetTableTypesBuilder`] for building responses to [`CommandGetTableTypes`] queries.
+//!
+//! [`CommandGetTableTypes`]: crate::sql::CommandGetTableTypes
+
+use std::sync::Arc;
+
+use arrow_array::{builder::StringBuilder, ArrayRef, RecordBatch};
+use arrow_schema::{DataType, Field, Schema, SchemaRef};
+use arrow_select::take::take;
+use once_cell::sync::Lazy;
+
+use crate::error::*;
+use crate::sql::CommandGetTableTypes;
+
+use super::lexsort_to_indices;
+
+/// A builder for a [`CommandGetTableTypes`] response.
+///
+/// Builds rows like this:
+///
+/// * table_type: utf8,
+#[derive(Default)]
+pub struct GetTableTypesBuilder {
+    // array builder for table types
+    table_type: StringBuilder,
+}
+
+impl CommandGetTableTypes {
+    /// Create a builder suitable for constructing a response
+    pub fn into_builder(self) -> GetTableTypesBuilder {
+        self.into()
+    }
+}
+
+impl From<CommandGetTableTypes> for GetTableTypesBuilder {
+    fn from(_value: CommandGetTableTypes) -> Self {
+        Self::new()
+    }
+}
+
+impl GetTableTypesBuilder {
+    /// Create a new instance of [`GetTableTypesBuilder`]
+    pub fn new() -> Self {
+        Self {
+            table_type: StringBuilder::new(),
+        }
+    }
+
+    /// Append a row
+    pub fn append(&mut self, table_type: impl AsRef<str>) {
+        self.table_type.append_value(table_type);
+    }
+
+    /// builds a `RecordBatch` with the correct schema for a `CommandGetTableTypes` response
+    pub fn build(self) -> Result<RecordBatch> {
+        let schema = self.schema();
+        let Self { mut table_type } = self;
+
+        // Make the arrays
+        let table_type = table_type.finish();
+
+        let batch = RecordBatch::try_new(schema, vec![Arc::new(table_type) as ArrayRef])?;
+
+        // Order filtered results by table_type
+        let indices = lexsort_to_indices(batch.columns());
+        let columns = batch
+            .columns()
+            .iter()
+            .map(|c| take(c, &indices, None))
+            .collect::<std::result::Result<Vec<_>, _>>()?;
+
+        Ok(RecordBatch::try_new(batch.schema(), columns)?)
+    }
+
+    /// Return the schema of the RecordBatch that will be returned
+    /// from [`CommandGetTableTypes`]
+    pub fn schema(&self) -> SchemaRef {
+        get_table_types_schema()
+    }
+}
+
+fn get_table_types_schema() -> SchemaRef {
+    Arc::clone(&GET_TABLE_TYPES_SCHEMA)
+}
+
+/// The schema for [`CommandGetTableTypes`].
+static GET_TABLE_TYPES_SCHEMA: Lazy<SchemaRef> = Lazy::new(|| {
+    Arc::new(Schema::new(vec![Field::new(
+        "table_type",
+        DataType::Utf8,
+        false,
+    )]))
+});
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use arrow_array::StringArray;
+
+    fn get_ref_batch() -> RecordBatch {
+        RecordBatch::try_new(
+            get_table_types_schema(),
+            vec![Arc::new(StringArray::from(vec![
+                "a_table_type",
+                "b_table_type",
+                "c_table_type",
+                "d_table_type",
+            ])) as ArrayRef],
+        )
+        .unwrap()
+    }
+
+    #[test]
+    fn test_table_types_are_sorted() {
+        let ref_batch = get_ref_batch();
+
+        let mut builder = GetTableTypesBuilder::new();
+        builder.append("b_table_type");
+        builder.append("a_table_type");
+        builder.append("d_table_type");
+        builder.append("c_table_type");
+        let schema_batch = builder.build().unwrap();
+
+        assert_eq!(schema_batch, ref_batch)
+    }
+
+    #[test]
+    fn test_builder_from_query() {
+        let ref_batch = get_ref_batch();
+        let query = CommandGetTableTypes {};
+
+        let mut builder = query.into_builder();
+        builder.append("a_table_type");
+        builder.append("b_table_type");
+        builder.append("c_table_type");
+        builder.append("d_table_type");
+        let schema_batch = builder.build().unwrap();
+
+        assert_eq!(schema_batch, ref_batch)
+    }
+}