Start adding execution

Author: matthewmturner

Cargo.toml

@@ -75,7 +75,7 @@ flightsql = [
 ]
 functions-json = ["datafusion-app/functions-json"]
 functions-parquet = ["datafusion-app/functions-parquet"]
-http = ["axum"]
+http = ["axum", "dep:metrics", "dep:metrics-exporter-prometheus", "dep:tower-http"]
 hudi = ["datafusion-app/hudi"]
 huggingface = ["datafusion-app/huggingface"]
 iceberg = ["datafusion-app/iceberg"]

crates/datafusion-app/src/config.rs

@@ -21,7 +21,6 @@ use std::path::PathBuf;
 
 #[cfg(feature = "udfs-wasm")]
 use datafusion_udfs_wasm::WasmInputDataType;
-use log::info;
 use serde::Deserialize;
 use std::collections::HashMap;
 
@@ -109,7 +108,6 @@ impl Default for ExecutionConfig {
 }
 
 fn default_ddl_path() -> Option<PathBuf> {
-    info!("Creating default ExecutionConfig");
     if let Some(user_dirs) = directories::UserDirs::new() {
         let ddl_path = user_dirs
             .home_dir()
@@ -241,7 +239,7 @@ pub struct WasmUdfConfig {
 }
 
 #[cfg(feature = "flightsql")]
-#[derive(Default)]
+#[derive(Clone, Default)]
 pub struct FlightSQLConfig {
     pub connection_url: String,
     pub benchmark_iterations: usize,
@@ -259,14 +257,12 @@ impl FlightSQLConfig {
     }
 }
 
-#[cfg(feature = "flightsql")]
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct AuthConfig {
     pub basic_auth: Option<BasicAuth>,
     pub bearer_token: Option<String>,
 }
 
-#[cfg(feature = "flightsql")]
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct BasicAuth {
     pub username: String,

crates/datafusion-app/src/flightsql.rs

@@ -34,7 +34,7 @@ use crate::{config::FlightSQLConfig, flightsql_benchmarks::FlightSQLBenchmarkSta
 
 pub type FlightSQLClient = Mutex<Option<FlightSqlServiceClient<Channel>>>;
 
-#[derive(Default)]
+#[derive(Clone, Default)]
 pub struct FlightSQLContext {
     config: FlightSQLConfig,
     flightsql_client: FlightSQLClient,

src/args.rs

@@ -88,9 +88,9 @@ pub struct DftArgs {
     #[clap(short = 'n', help = "Set the number of benchmark iterations to run")]
     pub benchmark_iterations: Option<usize>,
 
-    #[cfg(any(feature = "flightsql", feature = "flightsql"))]
-    #[clap(long, help = "Set the host and port to be used for FlightSQL")]
-    pub flightsql_host: Option<String>,
+    #[cfg(any(feature = "flightsql", feature = "http"))]
+    #[clap(long, help = "Set the host and port to be used for server")]
+    pub host: Option<String>,
 
     #[clap(
         long,

src/cli/mod.rs

@@ -601,9 +601,7 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
                 auth,
             );
             let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
-            flightsql_ctx
-                .create_client(cli.flightsql_host.clone())
-                .await?;
+            flightsql_ctx.create_client(cli.host.clone()).await?;
             app_execution.with_flightsql_ctx(flightsql_ctx);
         }
     }

src/config.rs

@@ -25,7 +25,7 @@ use lazy_static::lazy_static;
 use log::{debug, error};
 use serde::Deserialize;
 
-#[cfg(feature = "flightsql")]
+#[cfg(any(feature = "flightsql", feature = "http"))]
 use datafusion_app::config::AuthConfig;
 
 lazy_static! {
@@ -125,8 +125,29 @@ impl Default for FlightSQLClientConfig {
 }
 
 #[cfg(feature = "http")]
-#[derive(Clone, Debug, Default, Deserialize)]
-pub struct HttpServerConfig {}
+#[derive(Clone, Debug, Deserialize)]
+pub struct HttpServerConfig {
+    #[serde(default = "default_execution_config")]
+    pub execution: ExecutionConfig,
+    #[serde(default = "default_connection_url")]
+    pub connection_url: String,
+    #[serde(default = "default_server_metrics_port")]
+    pub server_metrics_port: String,
+    #[serde(default = "default_auth_config")]
+    pub auth: AuthConfig,
+}
+
+#[cfg(feature = "http")]
+impl Default for HttpServerConfig {
+    fn default() -> Self {
+        Self {
+            execution: default_execution_config(),
+            connection_url: default_connection_url(),
+            server_metrics_port: default_server_metrics_port(),
+            auth: default_auth_config(),
+        }
+    }
+}
 
 #[derive(Clone, Debug, Default, Deserialize)]
 pub struct AppConfig {
@@ -196,12 +217,12 @@ fn default_paste() -> bool {
     false
 }
 
-#[cfg(feature = "flightsql")]
+#[cfg(any(feature = "flightsql", feature = "http"))]
 pub fn default_connection_url() -> String {
     "http://localhost:50051".to_string()
 }
 
-#[cfg(feature = "flightsql")]
+#[cfg(any(feature = "flightsql", feature = "http"))]
 fn default_server_metrics_port() -> String {
     "0.0.0.0:9000".to_string()
 }
@@ -215,7 +236,7 @@ fn default_editor_config() -> EditorConfig {
     EditorConfig::default()
 }
 
-#[cfg(feature = "flightsql")]
+#[cfg(any(feature = "flightsql", feature = "http"))]
 fn default_auth_config() -> AuthConfig {
     AuthConfig::default()
 }

src/execution.rs

@@ -22,12 +22,6 @@ use datafusion::prelude::*;
 use datafusion_app::flightsql::{FlightSQLClient, FlightSQLContext};
 use datafusion_app::local::ExecutionContext;
 
-pub enum AppType {
-    Cli,
-    Tui,
-    FlightSQLServer,
-}
-
 /// Provides all core execution functionality for execution queries from either a local
 /// `SessionContext` or a remote `FlightSQL` service
 pub struct AppExecution {

src/lib.rs

@@ -2,10 +2,7 @@ pub mod args;
 pub mod cli;
 pub mod config;
 pub mod execution;
-#[cfg(feature = "flightsql")]
-pub mod flightsql_server;
-#[cfg(feature = "http")]
-pub mod http_server;
+pub mod server;
 pub mod telemetry;
 pub mod test_utils;
 pub mod tui;

src/main.rs

@@ -17,8 +17,8 @@
 
 use clap::Parser;
 use color_eyre::Result;
-#[cfg(feature = "flightsql")]
-use datafusion_dft::{args::Command, flightsql_server, http_server};
+#[cfg(any(feature = "flightsql", feature = "http"))]
+use datafusion_dft::{args::Command, server};
 use datafusion_dft::{args::DftArgs, cli, config::create_config, tui};
 
 fn main() -> Result<()> {
@@ -42,6 +42,11 @@ fn should_init_env_logger(cli: &DftArgs) -> bool {
     if let Some(Command::ServeFlightSql { .. }) = cli.command {
         return true;
     }
+    #[cfg(feature = "http")]
+    if let Some(Command::ServeHttp { .. }) = cli.command {
+        return true;
+    }
+
     if !cli.files.is_empty() || !cli.commands.is_empty() {
         return true;
     }
@@ -56,11 +61,13 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
 
     #[cfg(feature = "flightsql")]
     if let Some(Command::ServeFlightSql { .. }) = cli.command {
-        flightsql_server::try_run(cli.clone(), cfg.clone()).await?;
+        server::flightsql::try_run(cli.clone(), cfg.clone()).await?;
+        return Ok(());
     }
     #[cfg(feature = "http")]
     if let Some(Command::ServeHttp { .. }) = cli.command {
-        http_server::try_run(cli.clone(), cfg.clone()).await?;
+        server::http::try_run(cli.clone(), cfg.clone()).await?;
+        return Ok(());
     }
 
     if !cli.files.is_empty() || !cli.commands.is_empty() {

src/flightsql_server/mod.rs src/server/flightsql/mod.rssrc/flightsql_server/mod.rs renamed to src/server/flightsql/mod.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-pub mod services;
+pub mod service;
 
 use crate::args::DftArgs;
 use crate::config::AppConfig;
@@ -25,9 +25,7 @@ use datafusion_app::config::merge_configs;
 use datafusion_app::extensions::DftSessionStateBuilder;
 use datafusion_app::local::ExecutionContext;
 use log::info;
-use metrics::{describe_counter, describe_histogram};
-use metrics_exporter_prometheus::{Matcher, PrometheusBuilder};
-use services::flightsql::FlightSqlServiceImpl;
+use service::FlightSqlServiceImpl;
 use std::net::SocketAddr;
 use std::time::Duration;
 use tokio::net::TcpListener;
@@ -37,25 +35,12 @@ use tonic::transport::Server;
 #[cfg(feature = "flightsql")]
 use tower_http::validate_request::ValidateRequestHeaderLayer;
 
-const DEFAULT_TIMEOUT_SECONDS: u64 = 60;
-
-fn initialize_metrics() {
-    describe_counter!("requests", "Incoming requests by FlightSQL endpoint");
+use super::try_start_metrics_server;
 
-    describe_histogram!(
-        "get_flight_info_latency_ms",
-        metrics::Unit::Milliseconds,
-        "Get flight info latency ms"
-    );
-
-    describe_histogram!(
-        "do_get_fallback_latency_ms",
-        metrics::Unit::Milliseconds,
-        "Do get fallback latency ms"
-    )
-}
+const DEFAULT_TIMEOUT_SECONDS: u64 = 60;
+const DEFAULT_SERVER_ADDRESS: &str = "127.0.0.1:50051";
 
-fn create_server_handle(
+pub fn create_server_handle(
     config: &AppConfig,
     flightsql: FlightSqlServiceImpl,
     listener: TcpListener,
@@ -139,37 +124,20 @@ impl FlightSqlApp {
         addr: &str,
         metrics_addr: &str,
     ) -> Result<Self> {
-        let flightsql = services::flightsql::FlightSqlServiceImpl::new(app_execution);
-        // let OS choose a free port
+        info!("Listening to FlightSQL on {addr}");
+        let flightsql = service::FlightSqlServiceImpl::new(app_execution);
         let listener = TcpListener::bind(addr).await.unwrap();
-        let addr = listener.local_addr().unwrap();
 
         // prepare the shutdown channel
         let (tx, rx) = tokio::sync::oneshot::channel();
         let handle = create_server_handle(config, flightsql, listener, rx)?;
 
-        {
-            let builder = PrometheusBuilder::new();
-            let addr: SocketAddr = metrics_addr.parse()?;
-            info!("Listening to metrics on {addr}");
-            builder
-                .with_http_listener(addr)
-                .set_buckets_for_metric(
-                    Matcher::Suffix("latency_ms".to_string()),
-                    &[
-                        1.0, 3.0, 5.0, 10.0, 25.0, 50.0, 75.0, 100.0, 250.0, 500.0, 1000.0, 2500.0,
-                        5000.0, 10000.0, 20000.0,
-                    ],
-                )?
-                .install()
-                .expect("failed to install metrics recorder/exporter");
-
-            initialize_metrics();
-        }
+        let metrics_addr: SocketAddr = metrics_addr.parse()?;
+        try_start_metrics_server(metrics_addr)?;
 
         let app = Self {
             shutdown: Some(tx),
-            addr,
+            addr: metrics_addr,
             handle: Some(handle),
         };
         Ok(app)
@@ -188,7 +156,7 @@ impl FlightSqlApp {
         }
     }
 
-    pub async fn run_app(self) {
+    pub async fn run(self) {
         if let Some(handle) = self.handle {
             handle
                 .await
@@ -208,10 +176,8 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
     let session_state_builder = DftSessionStateBuilder::try_new(Some(merged_exec_config.clone()))?
         .with_extensions()
         .await?;
-    // FlightSQL Server mode: start a FlightSQL server
-    const DEFAULT_SERVER_ADDRESS: &str = "127.0.0.1:50051";
-    info!("Starting FlightSQL server on {}", DEFAULT_SERVER_ADDRESS);
     let session_state = session_state_builder.build()?;
+    // FlightSQL Server mode: start a FlightSQL server
     let execution_ctx = ExecutionContext::try_new(&merged_exec_config, session_state)?;
     if cli.run_ddl {
         execution_ctx.execute_ddl().await;
@@ -220,11 +186,10 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
     let app = FlightSqlApp::try_new(
         app_execution,
         &config,
-        &cli.flightsql_host
-            .unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
+        &cli.host.unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
         &config.flightsql_server.server_metrics_port,
     )
     .await?;
-    app.run_app().await;
+    app.run().await;
     Ok(())
 }