Move host arg (#319)

Author: matthewmturner

.github/workflows/test.yml

@@ -313,3 +313,31 @@ jobs:
       - name: Run tests
         run: |
           cargo t --features=http,flightsql server::http::router
+      - name: Run CLI cases
+        run: |
+          cargo t --features=http,flightsql server_cases::http
+  test-flightsql-server:
+    name: Extension / FlightSQL Server
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [amd64]
+    steps:
+      - uses: actions/checkout@v2
+        with:
+          submodules: true
+      - name: Cache Cargo
+        uses: actions/cache@v4
+        with:
+          path: /home/runner/.cargo
+          key: cargo-dft-cache-
+      - name: Cache Rust dependencies
+        uses: actions/cache@v4
+        with:
+          path: target
+          key: ${{ runner.os }}-cargo-target-${{ hashFiles('Cargo.lock') }}
+      - name: Setup Rust Toolchain
+        uses: ./.github/actions/setup-rust
+      - name: Run CLI cases
+        run: |
+          cargo t --features=http,flightsql server_cases::http

Cargo.toml

@@ -42,6 +42,7 @@ serde = { version = "1.0.197", features = ["derive"] }
 strum = "0.26.2"
 tokio = { version = "1.36.0", features = [
   "macros",
+  "process",
   "rt-multi-thread",
   "signal",
 ] }

docs/config.md

@@ -112,7 +112,7 @@ benchmark_iterations = 10
 
 [flightsql_server]
 connection_url = "http://localhost:50051"
-server_metrics_port = "0.0.0.0:9000"
+server_metrics_addr = "0.0.0.0:9000"
 ```
 
 ## Editor Config

docs/flightsql_server.md

@@ -49,8 +49,8 @@ Prometheus metrics are automatically published to help you monitor server perfor
 
 ```toml
 [flightsql_server]
-# Configure metrics port
-server_metrics_port = "0.0.0.0:9000"
+# Configure metrics addr
+server_metrics_addr = "0.0.0.0:9000"
 ```
 
 Available metrics include:

docs/http_server.md

@@ -36,7 +36,7 @@ Prometheus metrics are automatically published.
 
 ```toml
 [flightsql_server]
-server_metrics_port = "0.0.0.0:9000"
+server_metrics_addr = "0.0.0.0:9000"
 ```
 
 ## Benchmarking

src/args.rs

@@ -19,7 +19,10 @@
 
 use crate::config::get_data_dir;
 use clap::{Parser, Subcommand};
-use std::path::{Path, PathBuf};
+use std::{
+    net::SocketAddr,
+    path::{Path, PathBuf},
+};
 
 const LONG_ABOUT: &str = "
 dft - DataFusion TUI
@@ -92,12 +95,7 @@ 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 = "http"))]
-    #[clap(
-        long,
-        global = true,
-        help = "Set the host and port to be used for server"
-    )]
+    #[clap(long, help = "Host address to connect to")]
     pub host: Option<String>,
 
     #[clap(
@@ -107,13 +105,18 @@ pub struct DftArgs {
     )]
     pub output: Option<PathBuf>,
 
+    #[cfg(any(feature = "flightsql", feature = "http"))]
     #[command(subcommand)]
     pub command: Option<Command>,
 }
 
 impl DftArgs {
     pub fn config_path(&self) -> PathBuf {
-        if let Some(Command::ServeFlightSql { config: Some(cfg) }) = &self.command {
+        #[cfg(any(feature = "flightsql", feature = "http"))]
+        if let Some(Command::ServeFlightSql {
+            config: Some(cfg), ..
+        }) = &self.command
+        {
             return Path::new(cfg).to_path_buf();
         }
         if let Some(config) = self.config.as_ref() {
@@ -132,12 +135,20 @@ pub enum Command {
     ServeHttp {
         #[clap(short, long)]
         config: Option<String>,
+        #[clap(long, help = "Set the port to be used for server")]
+        addr: Option<SocketAddr>,
+        #[clap(long, help = "Set the port to be used for serving metrics")]
+        metrics_addr: Option<SocketAddr>,
     },
     /// Start a FlightSQL server
     #[command(name = "serve-flightsql")]
     ServeFlightSql {
         #[clap(short, long)]
         config: Option<String>,
+        #[clap(long, help = "Set the port to be used for server")]
+        addr: Option<SocketAddr>,
+        #[clap(long, help = "Set the port to be used for serving metrics")]
+        metrics_addr: Option<SocketAddr>,
     },
 }
 

src/config.rs

@@ -19,6 +19,8 @@
 
 use std::path::PathBuf;
 
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+
 use datafusion_app::config::ExecutionConfig;
 use directories::{ProjectDirs, UserDirs};
 use lazy_static::lazy_static;
@@ -84,8 +86,8 @@ pub struct FlightSQLServerConfig {
     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_server_metrics_addr")]
+    pub server_metrics_addr: SocketAddr,
     #[serde(default = "default_auth_config")]
     pub auth: AuthConfig,
 }
@@ -96,7 +98,7 @@ impl Default for FlightSQLServerConfig {
         Self {
             execution: default_execution_config(),
             connection_url: default_connection_url(),
-            server_metrics_port: default_server_metrics_port(),
+            server_metrics_addr: default_server_metrics_addr(),
             auth: default_auth_config(),
         }
     }
@@ -131,8 +133,8 @@ pub struct HttpServerConfig {
     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_server_metrics_addr")]
+    pub server_metrics_addr: SocketAddr,
     #[serde(default = "default_auth_config")]
     pub auth: AuthConfig,
     #[serde(default = "default_timeout_seconds")]
@@ -147,7 +149,7 @@ impl Default for HttpServerConfig {
         Self {
             execution: default_execution_config(),
             connection_url: default_connection_url(),
-            server_metrics_port: default_server_metrics_port(),
+            server_metrics_addr: default_server_metrics_addr(),
             auth: default_auth_config(),
             timeout_seconds: default_timeout_seconds(),
             result_limit: default_result_limit(),
@@ -229,8 +231,8 @@ pub fn default_connection_url() -> String {
 }
 
 #[cfg(any(feature = "flightsql", feature = "http"))]
-fn default_server_metrics_port() -> String {
-    "0.0.0.0:9000".to_string()
+fn default_server_metrics_addr() -> SocketAddr {
+    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9000)
 }
 
 #[derive(Clone, Debug, Default, Deserialize)]

src/main.rs

@@ -64,19 +64,19 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
     }
     #[cfg(feature = "http")]
     {
-        tracing_subscriber::registry()
-            .with(
-                tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
-                    format!(
-                        "{}=debug,tower_http=debug,axum=trace",
-                        env!("CARGO_CRATE_NAME")
-                    )
-                    .into()
-                }),
-            )
-            .with(tracing_subscriber::fmt::layer().without_time())
-            .init();
         if let Some(Command::ServeHttp { .. }) = cli.command {
+            tracing_subscriber::registry()
+                .with(
+                    tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| {
+                        format!(
+                            "{}=debug,tower_http=debug,axum=trace",
+                            env!("CARGO_CRATE_NAME")
+                        )
+                        .into()
+                    }),
+                )
+                .with(tracing_subscriber::fmt::layer().without_time())
+                .init();
             server::http::try_run(cli.clone(), cfg.clone()).await?;
             return Ok(());
         }

src/server/flightsql/mod.rs

@@ -17,7 +17,7 @@
 
 pub mod service;
 
-use crate::args::DftArgs;
+use crate::args::{Command, DftArgs};
 use crate::config::AppConfig;
 use crate::execution::AppExecution;
 use color_eyre::{eyre::eyre, Result};
@@ -26,7 +26,7 @@ use datafusion_app::extensions::DftSessionStateBuilder;
 use datafusion_app::local::ExecutionContext;
 use log::info;
 use service::FlightSqlServiceImpl;
-use std::net::SocketAddr;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::time::Duration;
 use tokio::net::TcpListener;
 use tokio::sync::oneshot;
@@ -38,7 +38,6 @@ use tower_http::validate_request::ValidateRequestHeaderLayer;
 use super::try_start_metrics_server;
 
 const DEFAULT_TIMEOUT_SECONDS: u64 = 60;
-const DEFAULT_SERVER_ADDRESS: &str = "127.0.0.1:50051";
 
 pub fn create_server_handle(
     config: &AppConfig,
@@ -121,8 +120,8 @@ impl FlightSqlApp {
     pub async fn try_new(
         app_execution: AppExecution,
         config: &AppConfig,
-        addr: &str,
-        metrics_addr: &str,
+        addr: SocketAddr,
+        metrics_addr: SocketAddr,
     ) -> Result<Self> {
         info!("Listening to FlightSQL on {addr}");
         let flightsql = service::FlightSqlServiceImpl::new(app_execution);
@@ -132,7 +131,6 @@ impl FlightSqlApp {
         let (tx, rx) = tokio::sync::oneshot::channel();
         let handle = create_server_handle(config, flightsql, listener, rx)?;
 
-        let metrics_addr: SocketAddr = metrics_addr.parse()?;
         try_start_metrics_server(metrics_addr)?;
 
         let app = Self {
@@ -188,13 +186,41 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
         execution_ctx.execute_ddl().await;
     }
     let app_execution = AppExecution::new(execution_ctx);
-    let app = FlightSqlApp::try_new(
-        app_execution,
-        &config,
-        &cli.host.unwrap_or(DEFAULT_SERVER_ADDRESS.to_string()),
-        &config.flightsql_server.server_metrics_port,
-    )
-    .await?;
+
+    let (addr, metrics_addr) = if let Some(cmd) = cli.command.clone() {
+        match cmd {
+            Command::ServeFlightSql {
+                addr: Some(addr),
+                metrics_addr: Some(metrics_addr),
+                ..
+            } => (addr, metrics_addr),
+            Command::ServeFlightSql {
+                addr: Some(addr),
+                metrics_addr: None,
+                ..
+            } => (addr, config.flightsql_server.server_metrics_addr),
+            Command::ServeFlightSql {
+                addr: None,
+                metrics_addr: Some(metrics_addr),
+                ..
+            } => (
+                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 50051),
+                metrics_addr,
+            ),
+
+            _ => (
+                SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 50051),
+                config.flightsql_server.server_metrics_addr,
+            ),
+        }
+    } else {
+        (
+            SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 50051),
+            config.flightsql_server.server_metrics_addr,
+        )
+    };
+
+    let app = FlightSqlApp::try_new(app_execution, &config, addr, metrics_addr).await?;
     app.run().await;
     Ok(())
 }