A little more work

Author: matthewmturner

Cargo.toml

@@ -83,6 +83,7 @@ flightsql = [
   "datafusion-app/flightsql",
   "datafusion-app/observability",
   "dep:arrow-flight",
+  "dep:jiff",
   "dep:metrics",
   "dep:metrics-exporter-prometheus",
   "dep:tonic",

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

@@ -25,6 +25,7 @@ use datafusion::{
     logical_expr::{logical_plan::dml::InsertOp, LogicalPlan, Values},
     physical_plan::execute_stream,
     prelude::{cast, lit, SessionContext},
+    scalar::ScalarValue,
     sql::TableReference,
 };
 use log::error;
@@ -94,7 +95,7 @@ impl ObservabilityContext {
                         DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".into())),
                     ),
                     lit(req.duration_ms),
-                    lit(req.rows),
+                    lit(ScalarValue::UInt64(req.rows)),
                     lit(req.status),
                 ]],
             };
@@ -128,7 +129,7 @@ pub struct ObservabilityRequestDetails {
     pub sql: String,
     pub start_ms: i64,
     pub duration_ms: i64,
-    pub rows: u64,
+    pub rows: Option<u64>,
     pub status: u16,
 }
 
@@ -142,7 +143,7 @@ fn req_fields() -> Vec<Field> {
             false,
         ),
         Field::new("duration_ms", DataType::Int64, false),
-        Field::new("rows", DataType::UInt64, false),
+        Field::new("rows", DataType::UInt64, true),
         Field::new("status", DataType::UInt16, false),
     ]
 }

src/server/flightsql/service.rs

@@ -25,7 +25,9 @@ use arrow_flight::{FlightDescriptor, FlightEndpoint, FlightInfo, Ticket};
 use datafusion::logical_expr::LogicalPlan;
 use datafusion::sql::parser::DFParser;
 use datafusion_app::local::ExecutionContext;
+use datafusion_app::observability::ObservabilityRequestDetails;
 use futures::{StreamExt, TryStreamExt};
+use jiff::Timestamp;
 use log::{debug, error, info};
 use metrics::{counter, histogram};
 use prost::Message;
@@ -60,12 +62,11 @@ impl FlightSqlServiceImpl {
 
     async fn get_flight_info_statement_handler(
         &self,
-        query: CommandStatementQuery,
+        query: String,
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
         info!("get_flight_info_statement query: {:?}", query);
         debug!("get_flight_info_statement request: {:?}", request);
-        let CommandStatementQuery { query, .. } = query;
         let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
         match DFParser::parse_sql_with_dialect(&query, &dialect) {
             Ok(statements) => {
@@ -192,11 +193,26 @@ impl FlightSqlService for FlightSqlServiceImpl {
         request: Request<FlightDescriptor>,
     ) -> Result<Response<FlightInfo>, Status> {
         counter!("requests", "endpoint" => "get_flight_info").increment(1);
-        let start = Instant::now();
-        let res = self.get_flight_info_statement_handler(query, request).await;
-        let duration = start.elapsed();
-        histogram!("get_flight_info_latency_ms").record(duration.as_millis() as f64);
-        res
+        let start = Timestamp::now();
+        let CommandStatementQuery { query, .. } = query;
+        let res = self
+            .get_flight_info_statement_handler(query.clone(), request)
+            .await?;
+        let duration = Timestamp::now() - start;
+        let headers = res.metadata();
+        println!("HEADAERS: {headers:?}");
+
+        let ctx = self.execution.session_ctx();
+        // let req = ObservabilityRequestDetails {
+        //     path: "GetFlightInfo".to_string(),
+        //     sql: query,
+        //     rows: None,
+        //     start_ms: start.as_millisecond(),
+        //     duration_ms: duration.get_milliseconds(),
+        // };
+        // self.execution.observability().try_record_request(ctx, req);
+        histogram!("get_flight_info_latency_ms").record(duration.get_milliseconds() as f64);
+        Ok(res)
     }
 
     async fn do_get_statement(

src/server/http/router.rs

@@ -265,7 +265,7 @@ async fn create_response(
         sql: req.sql,
         start_ms: start.as_millisecond(),
         duration_ms: elapsed.get_milliseconds(),
-        rows: details.rows,
+        rows: Some(details.rows),
         status: res.status().as_u16(),
     };
     let obs = state.execution.execution_ctx().observability();