Add path to observability reqs

Author: matthewmturner

Cargo.toml

@@ -81,6 +81,7 @@ default = ["functions-parquet"]
 deltalake = ["datafusion-app/deltalake"]
 flightsql = [
   "datafusion-app/flightsql",
+  "datafusion-app/observability",
   "dep:arrow-flight",
   "dep:metrics",
   "dep:metrics-exporter-prometheus",

crates/datafusion-app/src/observability/mod.rs

@@ -87,6 +87,7 @@ impl ObservabilityContext {
             let values = Values {
                 schema,
                 values: vec![vec![
+                    lit(req.path),
                     lit(req.sql),
                     cast(
                         lit(req.start_ms),
@@ -123,6 +124,7 @@ impl ObservabilityContext {
 
 /// Details that will be recorded in the configured observability request table
 pub struct ObservabilityRequestDetails {
+    pub path: String,
     pub sql: String,
     pub start_ms: i64,
     pub duration_ms: i64,
@@ -132,6 +134,7 @@ pub struct ObservabilityRequestDetails {
 
 fn req_fields() -> Vec<Field> {
     vec![
+        Field::new("path", DataType::Utf8, false),
         Field::new("sql", DataType::Utf8, false),
         Field::new(
             "timestamp",
@@ -183,6 +186,7 @@ mod test {
 
         let ctx = execution.session_ctx();
         let req = ObservabilityRequestDetails {
+            path: "/sql".to_string(),
             sql: "SELECT 1".to_string(),
             start_ms: 100,
             duration_ms: 200,

src/cli/mod.rs

@@ -30,7 +30,7 @@ use datafusion_app::extensions::DftSessionStateBuilder;
 use datafusion_app::local::ExecutionContext;
 use datafusion_app::local_benchmarks::LocalBenchmarkStats;
 use futures::{Stream, StreamExt};
-use log::{error, info};
+use log::info;
 use parquet::{arrow::ArrowWriter, file::properties::WriterProperties};
 use std::error::Error;
 use std::fs::File;
@@ -43,6 +43,7 @@ use {
         flightsql::FlightSQLContext,
         flightsql_benchmarks::FlightSQLBenchmarkStats,
     },
+    log::error,
     tonic::IntoRequest,
 };
 
@@ -606,11 +607,8 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
                 auth,
             );
             let flightsql_ctx = FlightSQLContext::new(flightsql_cfg);
-            if let Err(e) = flightsql_ctx.create_client(cli.host.clone()).await {
-                error!("{}", e.to_string())
-            } else {
-                app_execution.with_flightsql_ctx(flightsql_ctx);
-            }
+            flightsql_ctx.create_client(cli.host.clone()).await?;
+            app_execution.with_flightsql_ctx(flightsql_ctx);
         }
     }
     let app = CliApp::new(app_execution, cli.clone());

src/server/http/mod.rs

@@ -123,9 +123,7 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
         execution_ctx.execute_ddl().await;
     }
 
-    #[cfg(not(feature = "flightsql"))]
-    let app_execution = AppExecution::new(execution_ctx);
-    #[cfg(feature = "flightsql")]
+    #[allow(unused_mut)]
     let mut app_execution = AppExecution::new(execution_ctx);
     #[cfg(feature = "flightsql")]
     {

src/server/http/router.rs

@@ -19,7 +19,7 @@ use std::{io::Cursor, time::Duration};
 
 use axum::{
     body::Body,
-    extract::{Json, Path, Query, State},
+    extract::{Json, OriginalUri, Path, Query, State},
     response::{IntoResponse, Response},
     routing::{get, post},
     Router,
@@ -29,7 +29,7 @@ use datafusion_app::{observability::ObservabilityRequestDetails, ExecOptions, Ex
 use http::{HeaderValue, StatusCode};
 use jiff::Timestamp;
 use log::error;
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 use tokio_stream::StreamExt;
 use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};
 use tracing::info;
@@ -38,6 +38,7 @@ use crate::{config::HttpServerConfig, execution::AppExecution};
 
 #[derive(Debug)]
 struct ExecRequest {
+    path: String,
     sql: String,
 }
 
@@ -83,17 +84,22 @@ struct PostSqlBody {
     flightsql: bool,
 }
 
-async fn post_sql_handler(state: State<ExecutionState>, Json(body): Json<PostSqlBody>) -> Response {
+async fn post_sql_handler(
+    state: State<ExecutionState>,
+    OriginalUri(uri): OriginalUri,
+    Json(body): Json<PostSqlBody>,
+) -> Response {
     if body.flightsql && !cfg!(feature = "flightsql") {
         return (
             StatusCode::BAD_REQUEST,
             "FlightSQL is not enabled on this server",
         )
             .into_response();
     }
-    let rt = state.execution.session_ctx().runtime_env();
-    println!("Runtime {rt:?}");
-    let req = ExecRequest { sql: body.sql };
+    let req = ExecRequest {
+        path: uri.path().to_string(),
+        sql: body.sql.to_string(),
+    };
     let opts = ExecOptions::new(Some(state.config.result_limit), body.flightsql);
     create_response(&state, req, opts).await
 }
@@ -106,6 +112,7 @@ struct GetCatalogQueryParams {
 
 async fn get_catalog_handler(
     state: State<ExecutionState>,
+    OriginalUri(uri): OriginalUri,
     Query(query): Query<GetCatalogQueryParams>,
 ) -> Response {
     let opts = ExecOptions::new(None, query.flightsql);
@@ -117,11 +124,14 @@ async fn get_catalog_handler(
             .into_response();
     }
     let sql = "SHOW TABLES".to_string();
-    let req = ExecRequest { sql };
+    let req = ExecRequest {
+        path: uri.path().to_string(),
+        sql,
+    };
     create_response(&state, req, opts).await
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Serialize)]
 struct GetTablePathParams {
     catalog: String,
     schema: String,
@@ -136,19 +146,23 @@ struct GetTableQueryParams {
 
 async fn get_table_handler(
     state: State<ExecutionState>,
-    Path(params): Path<GetTablePathParams>,
+    Path(path): Path<GetTablePathParams>,
     Query(query): Query<GetTableQueryParams>,
+    OriginalUri(uri): OriginalUri,
 ) -> Response {
     let GetTablePathParams {
         catalog,
         schema,
         table,
-    } = params;
+    } = path;
     let sql = format!(
         "SELECT * FROM \"{catalog}\".\"{schema}\".\"{table}\" LIMIT {}",
         state.config.result_limit
     );
-    let req = ExecRequest { sql };
+    let req = ExecRequest {
+        path: uri.path().to_string(),
+        sql,
+    };
     let opts = ExecOptions::new(Some(state.config.result_limit), query.flightsql);
     create_response(&state, req, opts).await
 }
@@ -247,6 +261,7 @@ async fn create_response(
     let (res, details) = response_for_sql(state, req.sql.clone(), opts).await;
     let elapsed = Timestamp::now() - start;
     let req = ObservabilityRequestDetails {
+        path: req.path,
         sql: req.sql,
         start_ms: start.as_millisecond(),
         duration_ms: elapsed.get_milliseconds(),