Get tables (#327)

Author: matthewmturner

crates/datafusion-app/src/catalog/mod.rs

@@ -38,7 +38,7 @@ pub fn create_app_catalog(
 ) -> Result<Arc<dyn CatalogProvider>> {
     let catalog = MemoryCatalogProvider::new();
     let meta_schema = Arc::new(MemorySchemaProvider::new());
-    catalog.register_schema("meta", meta_schema.clone())?;
+    catalog.register_schema("meta", Arc::<MemorySchemaProvider>::clone(&meta_schema))?;
     let versions_table = try_create_meta_versions_table(app_name, app_version)?;
     meta_schema.register_table("versions".to_string(), versions_table)?;
     Ok(Arc::new(catalog))
@@ -56,7 +56,7 @@ fn try_create_meta_versions_table(app_name: &str, app_version: &str) -> Result<A
     let datafusion_version_arr = StringArray::from(vec![DATAFUSION_VERSION]);
     let datafusion_app_version_arr = StringArray::from(vec![env!("CARGO_PKG_VERSION")]);
     let batches = RecordBatch::try_new(
-        schema.clone(),
+        Arc::<Schema>::clone(&schema),
         vec![
             Arc::new(app_version_arr),
             Arc::new(datafusion_version_arr),

crates/datafusion-app/src/flightsql.rs

@@ -19,7 +19,7 @@ use std::sync::Arc;
 
 use arrow_flight::{
     decode::FlightRecordBatchStream,
-    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas},
+    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas, CommandGetTables},
     FlightInfo,
 };
 #[cfg(feature = "flightsql")]
@@ -226,14 +226,14 @@ impl FlightSQLContext {
     pub async fn get_db_schemas_flight_info(
         &self,
         catalog: Option<String>,
-        schema_pattern: Option<String>,
+        schema_filter_pattern: Option<String>,
     ) -> DFResult<FlightInfo> {
         let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;
         if let Some(client) = guard.as_mut() {
             let cmd = CommandGetDbSchemas {
                 catalog,
-                db_schema_filter_pattern: schema_pattern,
+                db_schema_filter_pattern: schema_filter_pattern,
             };
             client
                 .get_db_schemas(cmd)
@@ -246,6 +246,35 @@ impl FlightSQLContext {
         }
     }
 
+    pub async fn get_tables_flight_info(
+        &self,
+        catalog: Option<String>,
+        schema_filter_pattern: Option<String>,
+        table_name_filter_pattern: Option<String>,
+        table_types: Vec<String>,
+        include_schema: bool,
+    ) -> DFResult<FlightInfo> {
+        let client = Arc::clone(&self.client);
+        let mut guard = client.lock().await;
+        if let Some(client) = guard.as_mut() {
+            let cmd = CommandGetTables {
+                catalog,
+                db_schema_filter_pattern: schema_filter_pattern,
+                table_name_filter_pattern,
+                table_types,
+                include_schema,
+            };
+            client
+                .get_tables(cmd)
+                .await
+                .map_err(|e| DataFusionError::ArrowError(e, None))
+        } else {
+            Err(DataFusionError::External(
+                "No FlightSQL client configured.  Add one in `~/.config/dft/config.toml`".into(),
+            ))
+        }
+    }
+
     pub async fn do_get(&self, flight_info: FlightInfo) -> DFResult<Vec<FlightRecordBatchStream>> {
         let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;

crates/datafusion-app/src/observability/mod.rs

@@ -61,7 +61,7 @@ impl ObservabilityContext {
     }
 
     pub fn schema(&self) -> Arc<dyn SchemaProvider> {
-        self.schema.clone()
+        Arc::<dyn SchemaProvider>::clone(&self.schema)
     }
 
     /// Attempts to insert request details in `requests` table.  No verification is performed on

src/args.rs

@@ -127,22 +127,40 @@ impl DftArgs {
     }
 }
 
+/// Parameters for each command match to exactly how they are defined in specification (https://arrow.apache.org/docs/format/FlightSql.html#protocol-buffer-definitions)
 #[derive(Clone, Debug, Subcommand)]
 pub enum FlightSqlCommand {
     /// Executes `CommandStatementQuery` and `DoGet` to return results
     StatementQuery {
         /// The query to execute
+        #[clap(long)]
         sql: String,
     },
     /// Executes `CommandGetCatalogs` and `DoGet` to return results
     GetCatalogs,
     /// Executes `CommandGetDbSchemas` and `DoGet` to return results
     GetDbSchemas {
         /// The catalog to retrieve schemas
-        #[clap(long, short)]
+        #[clap(long)]
         catalog: Option<String>,
-        #[clap(long, short)]
-        schema_pattern: Option<String>,
+        /// Schema filter pattern to apply
+        #[clap(long)]
+        db_schema_filter_pattern: Option<String>,
+    },
+    /// Executes `CommandGetDbSchemas` and `DoGet` to return results
+    GetTables {
+        /// The catalog to retrieve schemas
+        #[clap(long)]
+        catalog: Option<String>,
+        /// Schema filter pattern to apply
+        #[clap(long)]
+        db_schema_filter_pattern: Option<String>,
+        /// Table name filter pattern to apply
+        #[clap(long)]
+        table_name_filter_pattern: Option<String>,
+        /// Specific table types to return
+        #[clap(long)]
+        table_types: Option<Vec<String>>,
     },
 }
 

src/cli/mod.rs

@@ -106,12 +106,39 @@ impl CliApp {
             }
             FlightSqlCommand::GetDbSchemas {
                 catalog,
-                schema_pattern,
+                db_schema_filter_pattern,
             } => {
                 let flight_info = self
                     .app_execution
                     .flightsql_ctx()
-                    .get_db_schemas_flight_info(catalog, schema_pattern)
+                    .get_db_schemas_flight_info(catalog, db_schema_filter_pattern)
+                    .await?;
+                let streams = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .do_get(flight_info)
+                    .await?;
+                let flight_batch_stream = stream::select_all(streams);
+                self.print_any_stream(flight_batch_stream).await;
+                Ok(())
+            }
+
+            FlightSqlCommand::GetTables {
+                catalog,
+                db_schema_filter_pattern,
+                table_name_filter_pattern,
+                table_types,
+            } => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_tables_flight_info(
+                        catalog,
+                        db_schema_filter_pattern,
+                        table_name_filter_pattern,
+                        table_types.unwrap_or_default(),
+                        false,
+                    )
                     .await?;
                 let streams = self
                     .app_execution

src/server/flightsql/service.rs

@@ -21,12 +21,13 @@ use arrow_flight::error::FlightError;
 use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
 use arrow_flight::sql::server::FlightSqlService;
 use arrow_flight::sql::{
-    Any, CommandGetCatalogs, CommandGetDbSchemas, CommandStatementQuery, SqlInfo,
+    Any, CommandGetCatalogs, CommandGetDbSchemas, CommandGetTables, CommandStatementQuery, SqlInfo,
     TicketStatementQuery,
 };
 use arrow_flight::{FlightDescriptor, FlightEndpoint, FlightInfo, Ticket};
 use color_eyre::Result;
 use datafusion::logical_expr::LogicalPlan;
+use datafusion::prelude::{col, lit};
 use datafusion::sql::parser::DFParser;
 use datafusion_app::local::ExecutionContext;
 use datafusion_app::observability::ObservabilityRequestDetails;
@@ -138,11 +139,44 @@ impl FlightSqlServiceImpl {
         histogram!(latency_metric).record(duration.get_milliseconds() as f64);
     }
 
+    async fn create_flight_info_for_logical_plan(
+        &self,
+        logical_plan: LogicalPlan,
+        request_id: Uuid,
+        _request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        let schema = logical_plan.schema();
+        let ticket = TicketStatementQuery {
+            statement_handle: request_id.to_string().into(),
+        };
+        debug!("created ticket handle: {:?}", ticket.statement_handle);
+        let mut bytes: Vec<u8> = Vec::new();
+        if ticket.encode(&mut bytes).is_ok() {
+            let info = FlightInfo::new()
+                .try_with_schema(schema.as_arrow())
+                .unwrap()
+                .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(bytes.clone())))
+                .with_descriptor(FlightDescriptor::new_cmd(bytes));
+            debug!("created flight info: {:?}", info);
+
+            let mut guard = self
+                .requests
+                .lock()
+                .map_err(|_| Status::internal("failed to acquire lock on requests"))?;
+            guard.insert(request_id, logical_plan);
+
+            Ok(Response::new(info))
+        } else {
+            error!("error encoding ticket");
+            Err(Status::internal("error encoding ticket"))
+        }
+    }
+
     async fn create_flight_info(
         &self,
         query: String,
         request_id: Uuid,
-        _request: Request<FlightDescriptor>,
+        request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
         debug!(
             "creating flight info for request id {request_id} with query: {:?}",
@@ -154,44 +188,15 @@ impl FlightSqlServiceImpl {
                 let statement = statements[0].clone();
                 let start = std::time::Instant::now();
 
-                let logical_plan = self.execution.statement_to_logical_plan(statement).await;
-
-                match logical_plan {
-                    Ok(logical_plan) => {
-                        debug!("logical planning took: {:?}", start.elapsed());
-                        let schema = logical_plan.schema();
-
-                        let ticket = TicketStatementQuery {
-                            statement_handle: request_id.to_string().into(),
-                        };
-                        debug!("created ticket handle: {:?}", ticket.statement_handle);
-                        let mut bytes: Vec<u8> = Vec::new();
-                        if ticket.encode(&mut bytes).is_ok() {
-                            let info = FlightInfo::new()
-                                .try_with_schema(schema.as_arrow())
-                                .unwrap()
-                                .with_endpoint(
-                                    FlightEndpoint::new().with_ticket(Ticket::new(bytes)),
-                                )
-                                .with_descriptor(FlightDescriptor::new_cmd(query));
-                            debug!("flight info: {:?}", info);
-
-                            let mut guard = self.requests.lock().map_err(|_| {
-                                Status::internal("failed to acquire lock on requests")
-                            })?;
-                            guard.insert(request_id, logical_plan);
-
-                            Ok(Response::new(info))
-                        } else {
-                            error!("error encoding ticket");
-                            Err(Status::internal("error encoding ticket"))
-                        }
-                    }
-                    Err(e) => {
-                        error!("error planning SQL query: {:?}", e);
-                        Err(Status::internal("error planning SQL query"))
-                    }
-                }
+                let logical_plan = self
+                    .execution
+                    .statement_to_logical_plan(statement)
+                    .await
+                    .map_err(|e| Status::internal(e.to_string()))?;
+
+                debug!("logical planning took: {:?}", start.elapsed());
+                self.create_flight_info_for_logical_plan(logical_plan, request_id, request)
+                    .await
             }
             Err(e) => {
                 error!("error parsing SQL query: {:?}", e);
@@ -228,7 +233,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
         _query: CommandGetCatalogs,
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
-        counter!("requests", "endpoint" => "get_flight_info").increment(1);
+        counter!("requests", "endpoint" => "get_flight_info_catalogs").increment(1);
         let start = Timestamp::now();
         let request_id = uuid::Uuid::new_v4();
         let query = "SELECT DISTINCT table_catalog FROM information_schema.tables".to_string();
@@ -251,7 +256,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
         command: CommandGetDbSchemas,
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
-        counter!("requests", "endpoint" => "get_flight_info").increment(1);
+        counter!("requests", "endpoint" => "get_flight_info_schemas").increment(1);
         let start = Timestamp::now();
         let request_id = uuid::Uuid::new_v4();
         let CommandGetDbSchemas {
@@ -264,16 +269,85 @@ impl FlightSqlService for FlightSqlServiceImpl {
             (Some(catalog), None) => format!("SELECT DISTINCT table_catalog, table_schema FROM information_schema.tables WHERE table_catalog = '{catalog}' ORDER BY table_catalog, table_schema"),
             (None, None) => "SELECT DISTINCT table_catalog, table_schema FROM information_schema.tables ORDER BY table_catalog, table_schema".to_string()
         };
-        println!("QUERY: {query}");
         let res = self.create_flight_info(query, request_id, request).await;
 
         // TODO: Move recording to after response is sent to not impact response latency
         self.record_request(
             start,
             Some(request_id.to_string()),
             res.as_ref().err(),
-            "/get_flight_info_catalogs".to_string(),
-            "get_flight_info_catalogs_latency_ms",
+            "/get_flight_info_schemas".to_string(),
+            "get_flight_info_schemas_latency_ms",
+        )
+        .await;
+        res
+    }
+
+    async fn get_flight_info_tables(
+        &self,
+        command: CommandGetTables,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        counter!("requests", "endpoint" => "get_flight_info_tables").increment(1);
+        let start = Timestamp::now();
+        let CommandGetTables {
+            catalog,
+            db_schema_filter_pattern,
+            table_name_filter_pattern,
+            table_types,
+            include_schema: _include_schema,
+        } = command;
+        let base_query = "SELECT * FROM information_schema.tables";
+
+        let mut df = self
+            .execution
+            .session_ctx()
+            .sql(base_query)
+            .await
+            .map_err(|e| Status::internal(e.to_string()))?;
+
+        if let Some(catalog) = catalog {
+            df = df
+                .filter(col("table_catalog").eq(lit(catalog)))
+                .map_err(|e| Status::internal(e.to_string()))?;
+        };
+        if let Some(schema_filter) = db_schema_filter_pattern {
+            df = df
+                .filter(col("table_schema").ilike(lit(format!("%{schema_filter}%"))))
+                .map_err(|e| Status::internal(e.to_string()))?;
+        };
+        if let Some(table_filter) = table_name_filter_pattern {
+            df = df
+                .filter(col("table_name").ilike(lit(format!("%{table_filter}%"))))
+                .map_err(|e| Status::internal(e.to_string()))?;
+        };
+        if !table_types.is_empty() {
+            let table_exprs = table_types.iter().map(lit).collect();
+            df = df
+                .filter(col("table_type").in_list(table_exprs, false))
+                .map_err(|e| Status::internal(e.to_string()))?;
+        };
+        df = df
+            .sort(vec![
+                col("table_catalog").sort(true, false),
+                col("table_schema").sort(true, false),
+                col("table_name").sort(true, false),
+            ])
+            .map_err(|e| Status::internal(e.to_string()))?;
+        let logical_plan = df.into_unoptimized_plan();
+
+        let request_id = uuid::Uuid::new_v4();
+        let res = self
+            .create_flight_info_for_logical_plan(logical_plan, request_id, request)
+            .await;
+
+        // TODO: Move recording to after response is sent to not impact response latency
+        self.record_request(
+            start,
+            Some(request_id.to_string()),
+            res.as_ref().err(),
+            "/get_flight_info_tables".to_string(),
+            "get_flight_info_tables_latency_ms",
         )
         .await;
         res