Get schemas

Author: matthewmturner

crates/datafusion-app/src/extensions/s3.rs

@@ -19,7 +19,7 @@
 
 use crate::config::ExecutionConfig;
 use crate::extensions::{DftSessionStateBuilder, Extension};
-use log::info;
+use log::{debug, info};
 use std::sync::Arc;
 
 use url::Url;
@@ -52,20 +52,18 @@ impl Extension for AwsS3Extension {
         for s3_config in s3_configs {
             match s3_config.to_object_store() {
                 Ok(object_store) => {
-                    info!("Created object store: {}", object_store);
+                    debug!("created object store: {}", object_store);
                     if let Some(object_store_url) = s3_config.object_store_url() {
-                        info!("Endpoint exists");
                         if let Ok(parsed_endpoint) = Url::parse(object_store_url) {
-                            info!("Parsed endpoint");
                             builder
                                 .runtime_env()
                                 .register_object_store(&parsed_endpoint, Arc::new(object_store));
-                            info!("Registered s3 object store");
+                            info!("registered s3 object store at {object_store_url}");
                         }
                     }
                 }
                 Err(e) => {
-                    log::error!("Error creating object store: {:?}", e);
+                    log::error!("error creating object store: {:?}", e);
                 }
             }
         }

crates/datafusion-app/src/flightsql.rs

@@ -18,7 +18,9 @@
 use std::sync::Arc;
 
 use arrow_flight::{
-    decode::FlightRecordBatchStream, sql::client::FlightSqlServiceClient, FlightInfo,
+    decode::FlightRecordBatchStream,
+    sql::{client::FlightSqlServiceClient, CommandGetDbSchemas},
+    FlightInfo,
 };
 #[cfg(feature = "flightsql")]
 use base64::engine::{general_purpose::STANDARD, Engine as _};
@@ -195,19 +197,19 @@ impl FlightSQLContext {
                         }
                     }
                     Err(e) => Err(DataFusionError::External(
-                        format!("Call to do_get failed: {}", e.to_string()).into(),
+                        format!("Call to do_get failed: {}", e).into(),
                     )),
                 }
             } else {
-                return Err(DataFusionError::External("Missing ticket".into()));
+                Err(DataFusionError::External("Missing ticket".into()))
             }
         } else {
-            return Err(DataFusionError::External("Missing client".into()));
+            Err(DataFusionError::External("Missing client".into()))
         }
     }
 
     pub async fn get_catalogs_flight_info(&self) -> DFResult<FlightInfo> {
-        let client = self.client.clone();
+        let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;
         if let Some(client) = guard.as_mut() {
             client
@@ -221,8 +223,31 @@ impl FlightSQLContext {
         }
     }
 
+    pub async fn get_db_schemas_flight_info(
+        &self,
+        catalog: Option<String>,
+        schema_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,
+            };
+            client
+                .get_db_schemas(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 = self.client.clone();
+        let client = Arc::clone(&self.client);
         let mut guard = client.lock().await;
         if let Some(client) = guard.as_mut() {
             let mut streams = Vec::new();

crates/datafusion-app/src/local.rs

@@ -72,24 +72,6 @@ pub struct ExecutionContext {
     observability: ObservabilityContext,
 }
 
-impl Default for ExecutionContext {
-    fn default() -> Self {
-        let cfg = SessionConfig::new().with_information_schema(true);
-        let session_ctx = SessionContext::new_with_config(cfg);
-        #[cfg(feature = "observability")]
-        let observability =
-            ObservabilityContext::try_new(ObservabilityConfig::default(), "test").unwrap();
-        Self {
-            config: ExecutionConfig::default(),
-            session_ctx,
-            ddl_path: None,
-            executor: None,
-            #[cfg(feature = "observability")]
-            observability,
-        }
-    }
-}
-
 impl std::fmt::Debug for ExecutionContext {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("ExecutionContext").finish()
@@ -181,6 +163,27 @@ impl ExecutionContext {
         Ok(ctx)
     }
 
+    /// Useful for testing execution functionality
+    pub fn test() -> Self {
+        let cfg = SessionConfig::new().with_information_schema(true);
+        let session_ctx = SessionContext::new_with_config(cfg);
+        let exec_cfg = ExecutionConfig::default();
+        // Okay to `unwrap` in a test
+        let app_catalog = create_app_catalog(&exec_cfg, "test", ".0.1.0").unwrap();
+        session_ctx.register_catalog("test", app_catalog);
+        #[cfg(feature = "observability")]
+        let observability =
+            ObservabilityContext::try_new(ObservabilityConfig::default(), "test").unwrap();
+        Self {
+            config: ExecutionConfig::default(),
+            session_ctx,
+            ddl_path: None,
+            executor: None,
+            #[cfg(feature = "observability")]
+            observability,
+        }
+    }
+
     pub fn config(&self) -> &ExecutionConfig {
         &self.config
     }

src/args.rs

@@ -129,10 +129,21 @@ impl DftArgs {
 
 #[derive(Clone, Debug, Subcommand)]
 pub enum FlightSqlCommand {
-    /// Executes `GetFlightInfo` and `DoGet` on the provided SQL query
-    StatementQuery { sql: String },
-    /// Executes `GetCatalogsFlightInfo` and `DoGet`
+    /// Executes `CommandStatementQuery` and `DoGet` to return results
+    StatementQuery {
+        /// The query to execute
+        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)]
+        catalog: Option<String>,
+        #[clap(long, short)]
+        schema_pattern: Option<String>,
+    },
 }
 
 #[derive(Clone, Debug, Subcommand)]

src/cli/mod.rs

@@ -102,7 +102,24 @@ impl CliApp {
                     .await?;
                 let flight_batch_stream = stream::select_all(streams);
                 self.print_any_stream(flight_batch_stream).await;
-
+                Ok(())
+            }
+            FlightSqlCommand::GetDbSchemas {
+                catalog,
+                schema_pattern,
+            } => {
+                let flight_info = self
+                    .app_execution
+                    .flightsql_ctx()
+                    .get_db_schemas_flight_info(catalog, schema_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(())
             }
         }

src/server/flightsql/service.rs

@@ -21,7 +21,8 @@ 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, CommandStatementQuery, SqlInfo, TicketStatementQuery,
+    Any, CommandGetCatalogs, CommandGetDbSchemas, CommandStatementQuery, SqlInfo,
+    TicketStatementQuery,
 };
 use arrow_flight::{FlightDescriptor, FlightEndpoint, FlightInfo, Ticket};
 use color_eyre::Result;
@@ -90,13 +91,13 @@ impl FlightSqlServiceImpl {
                         .boxed();
                     Ok(Response::new(flight_data_stream))
                 } else {
-                    Err(Status::internal("Plan not found for id"))
+                    Err(Status::internal("plan not found for id"))
                 }
             }
             Err(e) => {
                 error!("error decoding handle to uuid for {request_id}: {:?}", e);
                 Err(Status::internal(
-                    "Error decoding handle to uuid for {request_id}",
+                    "error decoding handle to uuid for {request_id}",
                 ))
             }
         }
@@ -131,20 +132,22 @@ impl FlightSqlServiceImpl {
             .try_record_request(ctx, req)
             .await
         {
-            error!("Error recording request: {}", e.to_string())
+            error!("error recording request: {}", e.to_string())
         }
 
         histogram!(latency_metric).record(duration.get_milliseconds() as f64);
     }
 
-    async fn get_flight_info_statement_handler(
+    async fn create_flight_info(
         &self,
         query: String,
         request_id: Uuid,
-        request: Request<FlightDescriptor>,
+        _request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
-        info!("get_flight_info_statement query: {:?}", query);
-        debug!("get_flight_info_statement request: {:?}", request);
+        debug!(
+            "creating flight info for request id {request_id} with query: {:?}",
+            query
+        );
         let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
         match DFParser::parse_sql_with_dialect(&query, &dialect) {
             Ok(statements) => {
@@ -174,25 +177,25 @@ impl FlightSqlServiceImpl {
                             debug!("flight info: {:?}", info);
 
                             let mut guard = self.requests.lock().map_err(|_| {
-                                Status::internal("Failed to acquire lock on requests")
+                                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(Status::internal("error encoding ticket"))
                         }
                     }
                     Err(e) => {
                         error!("error planning SQL query: {:?}", e);
-                        Err(Status::internal("Error planning SQL query"))
+                        Err(Status::internal("error planning SQL query"))
                     }
                 }
             }
             Err(e) => {
                 error!("error parsing SQL query: {:?}", e);
-                Err(Status::internal("Error parsing SQL query"))
+                Err(Status::internal("error parsing SQL query"))
             }
         }
     }
@@ -229,9 +232,40 @@ impl FlightSqlService for FlightSqlServiceImpl {
         let start = Timestamp::now();
         let request_id = uuid::Uuid::new_v4();
         let query = "SELECT DISTINCT table_catalog FROM information_schema.tables".to_string();
-        let res = self
-            .get_flight_info_statement_handler(query, request_id, request)
-            .await;
+        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",
+        )
+        .await;
+        res
+    }
+
+    async fn get_flight_info_schemas(
+        &self,
+        command: CommandGetDbSchemas,
+        request: Request<FlightDescriptor>,
+    ) -> Result<Response<FlightInfo>, Status> {
+        counter!("requests", "endpoint" => "get_flight_info").increment(1);
+        let start = Timestamp::now();
+        let request_id = uuid::Uuid::new_v4();
+        let CommandGetDbSchemas {
+            catalog,
+            db_schema_filter_pattern,
+        } = command;
+        let query = match (catalog, db_schema_filter_pattern) {
+            (Some(catalog), Some(filter)) => format!("SELECT DISTINCT table_catalog, table_schema FROM information_schema.tables WHERE table_catalog = '{catalog}' AND table_schema ILIKE '%{filter}%' ORDER BY table_catalog, table_schema"),
+            (None, Some(filter)) => format!("SELECT DISTINCT table_catalog, table_schema FROM information_schema.tables WHERE table_schema ILIKE '%{filter}%' ORDER BY table_catalog, table_schema"),
+            (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(
@@ -255,7 +289,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
         let CommandStatementQuery { query, .. } = query;
         let request_id = uuid::Uuid::new_v4();
         let res = self
-            .get_flight_info_statement_handler(query.clone(), request_id, request)
+            .create_flight_info(query.clone(), request_id, request)
             .await;
 
         // TODO: Move recording to after response is sent to not impact response latency