Much better

Author: matthewmturner

Cargo.lock

@@ -39,13 +39,15 @@ prost = "0.13.1"
 ratatui = "0.28.0"
 serde = { version = "1.0.197", features = ["derive"] }
 strum = "0.26.2"
-tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread", "signal"] }
 tokio-stream = { version = "0.1.15", features = ["net"] }
 tokio-util = "0.7.10"
 toml = "0.8.12"
 tonic = { version = "0.12.3", optional = true }
 tower = { version = "0.5.0" }
-tower-http = { version = "0.6.2", features = ["auth"], optional = true }
+tower-http = { version = "0.6.2", features = ["auth", "trace", "timeout"], optional = true }
+tracing = { version = "0.1.41", features = ["log"] }
+tracing-subscriber = { version = "0.3.19", features = ["env-filter"] }
 tui-logger = { version = "0.12", features = ["tracing-support"] }
 tui-textarea = { version = "0.6.1", features = ["search"] }
 url = { version = "2.5.2", optional = true }

Cargo.toml

@@ -20,6 +20,7 @@ use color_eyre::Result;
 #[cfg(any(feature = "flightsql", feature = "http"))]
 use datafusion_dft::{args::Command, server};
 use datafusion_dft::{args::DftArgs, cli, config::create_config, tui};
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 fn main() -> Result<()> {
     let cli = DftArgs::parse();
@@ -42,10 +43,6 @@ 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;
@@ -65,6 +62,18 @@ async fn app_entry_point(cli: DftArgs) -> Result<()> {
         return Ok(());
     }
     #[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 {
         server::http::try_run(cli.clone(), cfg.clone()).await?;
         return Ok(());

src/main.rs

@@ -17,99 +17,49 @@
 
 mod router;
 
-use std::{net::SocketAddr, time::Duration};
+use std::net::SocketAddr;
 
 use crate::{args::DftArgs, config::AppConfig, execution::AppExecution};
-use axum::{extract::State, routing::get, Router};
-use color_eyre::{eyre::eyre, Result};
-use datafusion::arrow::json::{writer::LineDelimited, Writer};
+use axum::Router;
+use color_eyre::Result;
 use datafusion_app::{
     config::merge_configs, extensions::DftSessionStateBuilder, local::ExecutionContext,
 };
-use log::info;
-use tokio::{net::TcpListener, sync::oneshot, task::JoinHandle};
-use tokio_stream::StreamExt;
-use tower_http::validate_request::ValidateRequestHeaderLayer;
+use router::create_router;
+use tokio::{net::TcpListener, signal};
+use tracing::info;
 
 use super::try_start_metrics_server;
 
-const DEFAULT_TIMEOUT_SECONDS: u64 = 60;
 const DEFAULT_SERVER_ADDRESS: &str = "127.0.0.1:8080";
 
-pub fn create_router(
-    config: &AppConfig,
-    // flightsql: FlightSqlServiceImpl,
-    listener: TcpListener,
-    rx: oneshot::Receiver<()>,
-    // shutdown_future: impl Future<Output = ()> + Send,
-) -> Router {
-    let server_timeout = Duration::from_secs(DEFAULT_TIMEOUT_SECONDS);
-    // let mut server_builder = Server::builder().timeout(server_timeout);
-    let shutdown_future = async move {
-        rx.await.ok();
+/// From https://github.com/tokio-rs/axum/blob/main/examples/graceful-shutdown/src/main.rs
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
     };
 
-    Router::new().route("/", get(|| async { "Hello, World!" }))
-
-    // axum::serve(listener, router).await.unwrap();
-
-    // TODO: onlu include TrailersLayer for testing
-    // if cfg!(feature = "flightsql") {
-    //     match (
-    //         &config.flightsql_server.auth.basic_auth,
-    //         &config.flightsql_server.auth.bearer_token,
-    //     ) {
-    //         (Some(_), Some(_)) => Err(eyre!("Only one auth type can be used at a time")),
-    //         (Some(basic), None) => {
-    //             let basic_auth_layer =
-    //                 ValidateRequestHeaderLayer::basic(&basic.username, &basic.password);
-    //             let f = server_builder
-    //                 .layer(basic_auth_layer)
-    //                 .add_service(flightsql.service())
-    //                 .serve_with_incoming_shutdown(
-    //                     tokio_stream::wrappers::TcpListenerStream::new(listener),
-    //                     shutdown_future,
-    //                 );
-    //             Ok(tokio::task::spawn(f))
-    //         }
-    //         (None, Some(token)) => {
-    //             let bearer_auth_layer = ValidateRequestHeaderLayer::bearer(token);
-    //             let f = server_builder
-    //                 .layer(bearer_auth_layer)
-    //                 .add_service(flightsql.service())
-    //                 .serve_with_incoming_shutdown(
-    //                     tokio_stream::wrappers::TcpListenerStream::new(listener),
-    //                     shutdown_future,
-    //                 );
-    //             Ok(tokio::task::spawn(f))
-    //         }
-    //         (None, None) => {
-    //             let f = server_builder
-    //                 .add_service(flightsql.service())
-    //                 .serve_with_incoming_shutdown(
-    //                     tokio_stream::wrappers::TcpListenerStream::new(listener),
-    //                     shutdown_future,
-    //                 );
-    //             Ok(tokio::task::spawn(f))
-    //         }
-    //     }
-    // } else {
-    //     let f = server_builder
-    //         .add_service(flightsql.service())
-    //         .serve_with_incoming_shutdown(
-    //             tokio_stream::wrappers::TcpListenerStream::new(listener),
-    //             shutdown_future,
-    //         );
-    //     Ok(tokio::task::spawn(f))
-    // }
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
 }
 
 /// Creates and manages a running FlightSqlServer with a background task
 pub struct HttpApp {
-    execution: AppExecution,
-    /// channel to send shutdown command
-    shutdown: Option<tokio::sync::oneshot::Sender<()>>,
-
     /// Address the server is listening on
     listener: TcpListener,
 
@@ -124,60 +74,30 @@ impl HttpApp {
         let listener = TcpListener::bind(addr).await.unwrap();
 
         // prepare the shutdown channel
-        let (tx, rx) = tokio::sync::oneshot::channel();
         let state = execution.execution_ctx().clone();
 
-        let router = Router::new()
-            .route(
-                "/",
-                get(|State(state): State<ExecutionContext>| async { "Hello, World!" }),
-            )
-            .route(
-                "/query",
-                get(|State(state): State<ExecutionContext>| async move {
-                    let r = state.execute_sql("SELECT 1").await;
-                    match r {
-                        Ok(mut ba) => {
-                            let mut buf = Vec::new();
-                            let mut writer: Writer<&mut [u8], LineDelimited> =
-                                datafusion::arrow::json::LineDelimitedWriter::new(&mut buf);
-                            while let Some(b) = ba.next().await {
-                                writer.write(&b.unwrap()).unwrap();
-                            }
-                            writer.finish().unwrap();
-                            let r = String::from_utf8(buf).unwrap();
-                            r
-                        }
-                        Err(e) => "Meep".to_string(),
-                    }
-                }),
-            )
-            .with_state(state);
+        let router = create_router(state);
 
         let metrics_addr: SocketAddr = metrics_addr.parse()?;
         try_start_metrics_server(metrics_addr)?;
 
-        let app = Self {
-            execution,
-            shutdown: Some(tx),
-            listener,
-            router,
-        };
+        let app = Self { listener, router };
         Ok(app)
     }
 
     pub async fn run(self) {
-        match axum::serve(self.listener, self.router).await {
-            Ok(_) => {}
+        match axum::serve(self.listener, self.router)
+            .with_graceful_shutdown(shutdown_signal())
+            .await
+        {
+            Ok(_) => {
+                info!("Shutting down app")
+            }
             Err(_) => {
                 panic!("Error serving HTTP app")
             }
         }
     }
-
-    async fn root_handler() -> String {
-        "Hi there".to_string()
-    }
 }
 
 pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {

src/server/http/mod.rs

@@ -14,3 +14,89 @@
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
+
+use std::{io::Cursor, time::Duration};
+
+use axum::{
+    body::Body,
+    extract::State,
+    response::{IntoResponse, Response},
+    routing::get,
+    Router,
+};
+use datafusion::arrow::json::ArrayWriter;
+use datafusion_app::local::ExecutionContext;
+use http::{HeaderValue, StatusCode};
+use log::error;
+use tokio_stream::StreamExt;
+use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};
+
+const DEFAULT_TIMEOUT_SECONDS: u64 = 10;
+
+pub fn create_router(execution: ExecutionContext) -> Router {
+    Router::new()
+        .route(
+            "/",
+            get(|State(_): State<ExecutionContext>| async { "Hello, from DFT!" }),
+        )
+        .route("/sql", get(execute_sql))
+        .layer((
+            TraceLayer::new_for_http(),
+            // Graceful shutdown will wait for outstanding requests to complete. Add a timeout so
+            // requests don't hang forever.
+            TimeoutLayer::new(Duration::from_secs(DEFAULT_TIMEOUT_SECONDS)),
+        ))
+        .with_state(execution)
+}
+
+async fn execute_sql(State(state): State<ExecutionContext>) -> Response {
+    let results = state.execute_sql("SELECT 1, 2").await;
+    match results {
+        Ok(mut batch_stream) => {
+            let mut buf: Cursor<Vec<u8>> = Cursor::new(Vec::new());
+            let mut writer = ArrayWriter::new(&mut buf);
+
+            while let Some(maybe_batch) = batch_stream.next().await {
+                match maybe_batch {
+                    Ok(batch) => {
+                        if let Err(e) = writer.write(&batch) {
+                            error!("Error serializing result batches: {}", e);
+                            return (StatusCode::INTERNAL_SERVER_ERROR, "Serialization error")
+                                .into_response();
+                        }
+                    }
+                    Err(e) => {
+                        error!("Error executing query: {}", e);
+                        return (StatusCode::INTERNAL_SERVER_ERROR, "Query execution error")
+                            .into_response();
+                    }
+                }
+            }
+
+            if let Err(e) = writer.finish() {
+                error!("Error finalizing JSON writer: {}", e);
+                return (StatusCode::INTERNAL_SERVER_ERROR, "Finalization error").into_response();
+            }
+
+            match String::from_utf8(buf.into_inner()) {
+                Ok(json) => {
+                    let mut res = Response::new(Body::new(json));
+                    res.headers_mut()
+                        .insert("content-type", HeaderValue::from_static("application/json"));
+                    res
+                }
+                Err(_) => {
+                    (StatusCode::INTERNAL_SERVER_ERROR, "UTF-8 conversion error").into_response()
+                }
+            }
+        }
+        Err(e) => {
+            error!("Error executing SQL: {}", e);
+            (
+                StatusCode::BAD_REQUEST,
+                format!("SQL execution failed: {}", e),
+            )
+                .into_response()
+        }
+    }
+}