feat(FlightSQL): support do_get_tables. (#10858)

Author: youngsofun

src/query/service/src/servers/flight_sql/flight_sql_service/catalog.rs

@@ -0,0 +1,132 @@
+// Copyright 2023 Datafuse Labs.
+//
+// Licensed 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.
+
+use std::sync::Arc;
+
+use arrow_array::builder::StringBuilder;
+use arrow_array::ArrayRef;
+use arrow_array::RecordBatch;
+use arrow_flight::utils::batches_to_flight_data;
+use arrow_schema::DataType;
+use arrow_schema::Field;
+use arrow_schema::Schema;
+use common_catalog::catalog::Catalog;
+use common_catalog::catalog::CatalogManager;
+use common_catalog::table_context::TableContext;
+use common_exception::ErrorCode;
+use futures_util::stream;
+use tonic::Status;
+
+use crate::servers::flight_sql::flight_sql_service::DoGetStream;
+
+pub(super) struct CatalogInfoProvider {}
+
+impl CatalogInfoProvider {
+    fn batch_to_get_stream(batch: RecordBatch) -> Result<DoGetStream, Status> {
+        let schema = (*batch.schema()).clone();
+        let batches = vec![batch];
+        let flight_data = batches_to_flight_data(schema, batches)
+            .map_err(|e| Status::internal(format!("{e:?}")))?
+            .into_iter()
+            .map(Ok);
+        let stream = stream::iter(flight_data);
+        Ok(Box::pin(stream))
+    }
+
+    async fn get_tables_internal(
+        ctx: Arc<dyn TableContext>,
+        catalog_name: Option<String>,
+        database_name: Option<String>,
+    ) -> common_exception::Result<(Vec<String>, Vec<String>, Vec<String>, Vec<String>)> {
+        let tenant = ctx.get_tenant();
+        let catalog_mgr = CatalogManager::instance();
+        let catalogs: Vec<(String, Arc<dyn Catalog>)> = if let Some(catalog_name) = catalog_name {
+            vec![(
+                catalog_name.clone(),
+                catalog_mgr.get_catalog(&catalog_name)?,
+            )]
+        } else {
+            catalog_mgr
+                .catalogs
+                .iter()
+                .map(|r| (r.key().to_string(), r.value().clone()))
+                .collect()
+        };
+
+        let mut catalog_names = vec![];
+        let mut database_names = vec![];
+        let mut table_names = vec![];
+        let mut table_types = vec![];
+        let table_type = "table".to_string();
+        for (catalog_name, catalog) in catalogs.into_iter() {
+            let dbs = if let Some(database_name) = &database_name {
+                vec![catalog.get_database(tenant.as_str(), database_name).await?]
+            } else {
+                catalog.list_databases(tenant.as_str()).await?
+            };
+            for db in dbs {
+                let db_name = db.name().to_string().into_boxed_str();
+                let db_name: &str = Box::leak(db_name);
+                let tables = match catalog.list_tables(tenant.as_str(), db_name).await {
+                    Ok(tables) => tables,
+                    Err(err) if err.code() == ErrorCode::EMPTY_SHARE_ENDPOINT_CONFIG => {
+                        tracing::warn!("list tables failed on db {}: {}", db.name(), err);
+                        continue;
+                    }
+                    Err(err) => return Err(err),
+                };
+                for table in tables {
+                    catalog_names.push(catalog_name.clone());
+                    database_names.push(db_name.to_string());
+                    table_names.push(table.name().to_string());
+                    table_types.push(table_type.clone());
+                }
+            }
+        }
+        Ok((catalog_names, database_names, table_names, table_types))
+    }
+
+    pub(crate) async fn get_tables(
+        ctx: Arc<dyn TableContext>,
+        catalog_name: Option<String>,
+        database_name: Option<String>,
+    ) -> Result<DoGetStream, Status> {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("catalog_name", DataType::Utf8, false),
+            Field::new("db_schema_name", DataType::Utf8, false),
+            Field::new("table_name", DataType::Utf8, false),
+            Field::new("table_type", DataType::Utf8, false),
+        ]));
+        let (catalog_name, db_schema_name, table_name, table_type) =
+            Self::get_tables_internal(ctx.clone(), catalog_name, database_name)
+                .await
+                .map_err(|e| Status::internal(format!("{e:?}")))?;
+        let batch = RecordBatch::try_new(schema, vec![
+            Self::string_array(catalog_name),
+            Self::string_array(db_schema_name),
+            Self::string_array(table_name),
+            Self::string_array(table_type),
+        ])
+        .map_err(|e| Status::internal(format!("RecordBatch::try_new fail {:?}", e)))?;
+        Self::batch_to_get_stream(batch)
+    }
+
+    fn string_array(values: Vec<String>) -> ArrayRef {
+        let mut builder = StringBuilder::new();
+        for v in &values {
+            builder.append_value(v);
+        }
+        Arc::new(builder.finish())
+    }
+}

src/query/service/src/servers/flight_sql/flight_sql_service/mod.rs

@@ -14,6 +14,7 @@
 
 // The servers module used for external communication with user, such as MySQL wired protocol, etc.
 
+mod catalog;
 mod query;
 mod service;
 mod session;
@@ -23,6 +24,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use arrow_flight::FlightData;
+use catalog::CatalogInfoProvider;
 use common_sql::plans::Plan;
 use common_sql::PlanExtras;
 use dashmap::DashMap;

src/query/service/src/servers/flight_sql/flight_sql_service/service.rs

@@ -377,10 +377,18 @@ impl FlightSqlService for FlightSqlServiceImpl {
     async fn do_get_tables(
         &self,
         query: CommandGetTables,
-        _request: Request<Ticket>,
+        request: Request<Ticket>,
     ) -> Result<Response<<Self as FlightService>::DoGetStream>, Status> {
         tracing::info!("do_get_tables({query:?})");
-        Err(Status::unimplemented("do_get_tables not implemented"))
+        let session = self.get_session(&request)?;
+        let context = session
+            .create_query_context()
+            .await
+            .map_err(|e| status!("Could not create_query_context", e))?;
+        Ok(Response::new(
+            super::CatalogInfoProvider::get_tables(context.clone(), query.catalog.clone(), None)
+                .await?,
+        ))
     }
 
     #[async_backtrace::framed]