Start updating http testing (#297)

Author: matthewmturner

.github/workflows/test.yml

@@ -305,6 +305,31 @@ jobs:
   test-http-server:
     name: Extension / HTTP 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: /home/runner/target
+          key: target-dft-cache-
+      - name: Setup Rust Toolchain
+        uses: ./.github/actions/setup-rust
+      - name: Run auth tests
+        run: |
+          cargo t --features=http server::http::router::test
+  test-http-server-flightsql:
+    name: Extension / HTTP Server (FlightSQL)
+    runs-on: ubuntu-latest
     strategy:
       matrix:
         arch: [amd64]
@@ -326,7 +351,7 @@ jobs:
         uses: ./.github/actions/setup-rust
       - name: Start FlightSQL Server
         run: |
-          cargo r --features=http -- serve-http &
+          cargo r --features=flightsql -- serve-flight-sql &
       - name: Run auth tests
         run: |
-          cargo t --features=http server::test
+          cargo t --features=http server::http::router::flightsql_test

Cargo.lock

@@ -66,9 +66,11 @@ assert_cmd = "2.0.16"
 datafusion-udfs-wasm = { version = "0.1.0", path = "crates/datafusion-udfs-wasm", features = [
   "serde",
 ] }
+http-body-util = "0.1.3"
 insta = { version = "1.40.0", features = ["yaml"] }
 predicates = "3.1.2"
 reqwest = { version = "0.12.14", features = ["json"] }
+serde_json = "1.0.140"
 tempfile = "3.2.0"
 url = "2.5.2"
 

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};
+#[cfg(feature = "http")]
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 fn main() -> Result<()> {

src/main.rs

@@ -32,11 +32,11 @@ use datafusion_app::{
 };
 use router::create_router;
 use tokio::{net::TcpListener, signal};
-use tracing::info;
+use tracing::{debug, info};
 
 use super::try_start_metrics_server;
 
-const DEFAULT_SERVER_ADDRESS: &str = "127.0.0.1:8080";
+const DEFAULT_SERVER_ADDRESS: &str = "localhost:8080";
 
 /// From https://github.com/tokio-rs/axum/blob/main/examples/graceful-shutdown/src/main.rs
 async fn shutdown_signal() {
@@ -142,7 +142,7 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
             .await?;
         app_execution.with_flightsql_ctx(flightsql_context);
     }
-    info!("Created AppExecution: {app_execution:?}");
+    debug!("Created AppExecution: {app_execution:?}");
     let app = HttpApp::try_new(
         app_execution,
         config.clone(),

src/server/http/mod.rs

@@ -100,6 +100,13 @@ async fn get_catalog_handler(
     Query(query): Query<GetCatalogQueryParams>,
 ) -> Response {
     let opts = ExecOptions::new(None, query.flightsql);
+    if opts.flightsql && !cfg!(feature = "flightsql") {
+        return (
+            StatusCode::BAD_REQUEST,
+            "FlightSQL is not enabled on this server",
+        )
+            .into_response();
+    }
     let sql = "SHOW TABLES".to_string();
     execute_sql_with_opts(state, sql, opts).await
 }
@@ -150,7 +157,7 @@ async fn execute_sql_with_opts(
         )
             .into_response(),
 
-        Err(e) => (StatusCode::BAD_REQUEST, format!("Execution failed: {}", e)).into_response(),
+        Err(e) => (StatusCode::BAD_REQUEST, format!("{}", e)).into_response(),
     }
 }
 
@@ -169,6 +176,7 @@ async fn batch_stream_to_response(batch_stream: SendableRecordBatchStream) -> Re
             }
             Err(e) => {
                 error!("Error executing query: {}", e);
+                // TODO: Use more appropriate errors, like 404 for table that doesnt exist
                 return (StatusCode::INTERNAL_SERVER_ERROR, "Query execution error")
                     .into_response();
             }
@@ -190,3 +198,164 @@ async fn batch_stream_to_response(batch_stream: SendableRecordBatchStream) -> Re
         Err(_) => (StatusCode::INTERNAL_SERVER_ERROR, "UTF-8 conversion error").into_response(),
     }
 }
+
+#[cfg(test)]
+mod test {
+    use axum::body::Body;
+    use datafusion_app::{
+        config::ExecutionConfig, extensions::DftSessionStateBuilder, local::ExecutionContext,
+    };
+    use http::{Request, StatusCode};
+    use http_body_util::BodyExt;
+
+    use crate::{
+        config::HttpServerConfig, execution::AppExecution, server::http::router::create_router,
+    };
+    use tower::ServiceExt;
+
+    fn setup() -> (AppExecution, HttpServerConfig) {
+        let config = ExecutionConfig::default();
+        let state = DftSessionStateBuilder::try_new(None)
+            .unwrap()
+            .build()
+            .unwrap();
+        let local = ExecutionContext::try_new(&config, state).unwrap();
+        let execution = AppExecution::new(local);
+
+        let http_config = HttpServerConfig::default();
+        (execution, http_config)
+    }
+
+    #[tokio::test]
+    async fn test_get_catalog() {
+        let (execution, http_config) = setup();
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/catalog")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_get_table() {
+        let (execution, http_config) = setup();
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/table/datafusion/information_schema/df_settings")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_get_nonexistent_table() {
+        let (execution, http_config) = setup();
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/table/datafusion/information_schema/df_setting")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
+    }
+
+    #[tokio::test]
+    async fn test_correct_when_flightsql_not_enabled() {
+        let (execution, http_config) = setup();
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/catalog?flightsql=true")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
+        let body = res.into_body().collect().await.unwrap().to_bytes();
+        assert_eq!(body, "FlightSQL is not enabled on this server".as_bytes())
+    }
+}
+
+#[cfg(all(test, feature = "flightsql"))]
+mod flightsql_test {
+    use axum::body::Body;
+    use datafusion_app::{
+        config::{ExecutionConfig, FlightSQLConfig},
+        extensions::DftSessionStateBuilder,
+        flightsql::FlightSQLContext,
+        local::ExecutionContext,
+    };
+    use http::{Request, StatusCode};
+
+    use crate::{
+        config::HttpServerConfig, execution::AppExecution, server::http::router::create_router,
+    };
+    use tower::ServiceExt;
+
+    async fn setup() -> (AppExecution, HttpServerConfig) {
+        let config = ExecutionConfig::default();
+        let state = DftSessionStateBuilder::try_new(None)
+            .unwrap()
+            .build()
+            .unwrap();
+        let local = ExecutionContext::try_new(&config, state).unwrap();
+        let mut execution = AppExecution::new(local);
+        let flightsql_cfg = FlightSQLConfig {
+            connection_url: "localhost:50051".to_string(),
+            ..Default::default()
+        };
+        let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
+        flightsql_ctx
+            .create_client(Some("http://localhost:50051".to_string()))
+            .await
+            .unwrap();
+        execution.with_flightsql_ctx(flightsql_ctx);
+
+        let http_config = HttpServerConfig::default();
+        (execution, http_config)
+    }
+
+    #[tokio::test]
+    async fn test_get_catalog() {
+        let (execution, http_config) = setup().await;
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/catalog?flightsql=true")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_get_table() {
+        let (execution, http_config) = setup().await;
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/table/datafusion/information_schema/df_settings?flightsql=true")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::OK);
+    }
+
+    #[tokio::test]
+    async fn test_get_nonexistent_table() {
+        let (execution, http_config) = setup().await;
+        let router = create_router(execution, http_config);
+
+        let req = Request::builder()
+            .uri("/table/datafusion/information_schema/df_setting?flightsql=true")
+            .body(Body::empty())
+            .unwrap();
+        let res = router.oneshot(req).await.unwrap();
+        assert_eq!(res.status(), StatusCode::BAD_REQUEST);
+    }
+}

src/server/http/router.rs

@@ -198,6 +198,8 @@ impl TestConfigBuilder {
         self
     }
 
+    // TODO: Update this to work with HTTP server
+    #[allow(dead_code)]
     #[cfg(feature = "flightsql")]
     pub fn with_server_auth(
         &mut self,

tests/config.rs

@@ -18,6 +18,4 @@
 mod cli_cases;
 mod config;
 mod extension_cases;
-#[cfg(feature = "http")]
-mod server;
 mod tui_cases;