Json

Author: matthewmturner

Cargo.lock

@@ -41,6 +41,7 @@ pin-project-lite = { version = "0.2.14" }
 prost = "0.14"
 ratatui = { optional = true, version = "0.28.0" }
 serde = { features = ["derive"], version = "1.0.197" }
+serde_json = "1.0"
 strum = "0.26.2"
 tokio = { features = [
   "macros",

Cargo.toml

@@ -37,6 +37,7 @@ opendal = { features = [
 parking_lot = "0.12.3"
 prost = { optional = true, version = "0.14" }
 serde = { features = ["derive"], version = "1.0.197" }
+serde_json = "1.0"
 tokio = { features = ["macros", "rt-multi-thread"], version = "1.36.0" }
 tokio-metrics = { features = [
   "metrics-rs-integration",

crates/datafusion-app/Cargo.toml

@@ -506,10 +506,14 @@ impl FlightSQLContext {
             return Err(eyre::eyre!("Only a single SQL statement can be analyzed"));
         }
 
-        // 1. Create Action with type "analyze_query" and SQL in body
+        // 1. Create JSON request and encode as Action body
+        let request = crate::stats::AnalyzeQueryRequest::with_sql(query);
+        let request_body = serde_json::to_vec(&request)
+            .map_err(|e| eyre::eyre!("Failed to serialize request: {}", e))?;
+
         let action = Action {
             r#type: "analyze_query".to_string(),
-            body: query.as_bytes().to_vec().into(),
+            body: request_body.into(),
         };
 
         // 2. Call do_action on the FlightSQL service

crates/datafusion-app/src/flightsql.rs

@@ -37,8 +37,40 @@ use datafusion::{
 };
 use itertools::Itertools;
 use log::debug;
+use serde::{Deserialize, Serialize};
 use std::{collections::HashMap, sync::Arc, time::Duration};
 
+/// Request structure for the analyze_query action
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AnalyzeQueryRequest {
+    /// SQL query to analyze (currently the only supported format)
+    pub sql: Option<String>,
+
+    // Future extensibility fields (not yet implemented):
+    // /// Substrait query plan (binary or JSON)
+    // pub substrait: Option<Vec<u8>>,
+    // /// Serialized logical plan
+    // pub logical_plan: Option<String>,
+    // /// Serialized physical plan
+    // pub physical_plan: Option<String>,
+}
+
+impl AnalyzeQueryRequest {
+    /// Create a new request with a SQL query
+    pub fn with_sql(sql: impl Into<String>) -> Self {
+        Self {
+            sql: Some(sql.into()),
+        }
+    }
+
+    /// Get the SQL query, returning an error if not present
+    pub fn sql(&self) -> color_eyre::Result<&str> {
+        self.sql
+            .as_deref()
+            .ok_or_else(|| color_eyre::eyre::eyre!("sql field is required"))
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct ExecutionStats {
     query: String,

crates/datafusion-app/src/stats.rs

@@ -37,19 +37,42 @@ The examples in this specification use SQL for illustration, but the protocol wo
 
 ### Request Format
 
-**Request Body**: UTF-8 encoded SQL query string
+**Request Body**: JSON-encoded query request structure
 
-**Important**: The SQL query must contain exactly one SQL statement. Multiple statements (e.g., separated by semicolons) are not supported and will result in an error.
+The request body should be a JSON object with the following structure:
+
+```json
+{
+  "sql": "SELECT * FROM table WHERE id > 100"
+}
+```
+
+**Current Fields**:
+- `sql` (string, required): The SQL query to analyze. Must contain exactly one SQL statement. Multiple statements (e.g., separated by semicolons) are not supported and will result in an error.
+
+**Future Extensibility**:
+The protocol is designed to be extensible. Future versions may support additional query representation fields:
+- `substrait` (bytes): Substrait query plan (binary or JSON)
+- `logical_plan` (string): Serialized logical plan
+- `physical_plan` (string): Serialized physical plan
+
+Servers should ignore unknown fields and clients should only send one query representation field at a time.
 
 **Example**:
 ```rust
+use serde_json::json;
+
+let request_body = json!({
+    "sql": "SELECT * FROM table WHERE id > 100"
+});
+
 Action {
     r#type: "analyze_query".to_string(),
-    body: "SELECT * FROM table WHERE id > 100".as_bytes().to_vec().into()
+    body: serde_json::to_vec(&request_body)?.into()
 }
 ```
 
-**Request Encoding**: The SQL query string should be encoded as UTF-8 bytes in the `Action.body` field.
+**Request Encoding**: The JSON object should be serialized to UTF-8 bytes in the `Action.body` field.
 
 ### Response Format
 
@@ -306,29 +329,46 @@ To implement this protocol in an Arrow Flight service:
 
 **Pseudo-code**:
 ```rust
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Deserialize, Serialize)]
+struct AnalyzeQueryRequest {
+    sql: Option<String>,
+    // Future fields:
+    // substrait: Option<Vec<u8>>,
+    // logical_plan: Option<String>,
+    // physical_plan: Option<String>,
+}
+
 async fn do_action_fallback(&self, request: Request<Action>) -> Result<Response<Stream>> {
     let action = request.into_inner();
 
     if action.r#type == "analyze_query" {
-        // 1. Extract query
-        let query = String::from_utf8(action.body.to_vec())?;
+        // 1. Parse JSON request
+        let request: AnalyzeQueryRequest = serde_json::from_slice(&action.body)?;
+
+        // 2. Extract SQL query (only supported format for now)
+        let query = request.sql.ok_or_else(||
+            Status::invalid_argument("sql field is required")
+        )?;
 
-        // 2. Execute with metrics
+        // 3. Execute with metrics
         let stats = self.analyze_query(&query).await?;
 
-        // 3. Convert to metrics batch (8-field schema with namespaced metrics)
+        // 4. Convert to metrics batch (8-field schema with namespaced metrics)
         let metrics_batch = stats.to_metrics_table()?;
 
-        // 4. Encode as FlightData (no metadata needed)
+        // 5. Encode as FlightData (no metadata needed)
         let flight_data = batches_to_flight_data(&metrics_batch.schema(), vec![metrics_batch])?;
 
-        // 5. Serialize and wrap in Result messages
+        // 6. Serialize and wrap in Result messages
+        // Note: Query is NOT included in response; clients must retain the original request
         let results: Vec<arrow_flight::Result> = flight_data
             .into_iter()
             .map(|fd| arrow_flight::Result { body: fd.encode_to_vec().into() })
             .collect();
 
-        // 6. Return stream
+        // 7. Return stream
         Ok(Response::new(stream::iter(results.into_iter().map(Ok))))
     } else {
         Err(Status::unimplemented("Unknown action"))
@@ -342,10 +382,18 @@ To consume this protocol:
 
 1. **Send Request**
    ```rust
+   use serde_json::json;
+
    let query = "SELECT * FROM table WHERE id > 100";
+
+   // Construct JSON request
+   let request_body = json!({
+       "sql": query
+   });
+
    let action = Action {
        r#type: "analyze_query".to_string(),
-       body: query.as_bytes().to_vec().into(),
+       body: serde_json::to_vec(&request_body)?.into(),
    };
    let stream = client.do_action(action).await?;
    ```

docs/arrow_flight_analyze_protocol.md

@@ -850,34 +850,41 @@ impl FlightSqlService for FlightSqlServiceImpl {
 
         match action.r#type.as_str() {
             "analyze_query" => {
-                // 1. Extract SQL query from action.body
-                let sql = String::from_utf8(action.body.to_vec())
-                    .map_err(|e| Status::invalid_argument(format!("Invalid UTF-8: {}", e)))?;
+                // 1. Parse JSON request body
+                let request: datafusion_app::stats::AnalyzeQueryRequest =
+                    serde_json::from_slice(&action.body).map_err(|e| {
+                        Status::invalid_argument(format!("Invalid JSON request: {}", e))
+                    })?;
+
+                // 2. Extract SQL query (only supported format for now)
+                let sql = request
+                    .sql()
+                    .map_err(|e| Status::invalid_argument(e.to_string()))?;
 
                 info!("Analyzing query via do_action: {}", sql);
 
-                // 2. Execute analyze_query on ExecutionContext
+                // 3. Execute analyze_query on ExecutionContext
                 let mut stats = self
                     .execution
-                    .analyze_query(&sql)
+                    .analyze_query(sql)
                     .await
                     .map_err(|e| Status::internal(format!("Analyze failed: {}", e)))?;
 
                 stats.collect_stats(); // Collect IO and compute metrics from plan
 
-                // 3. Convert ExecutionStats to metrics table format
+                // 4. Convert ExecutionStats to metrics table format
                 let metrics_batch = stats.to_metrics_table().map_err(|e| {
                     Status::internal(format!("Metrics serialization failed: {}", e))
                 })?;
 
-                // 4. Encode metrics batch as FlightData
+                // 5. Encode metrics batch as FlightData
                 let flight_data =
                     batches_to_flight_data(&metrics_batch.schema(), vec![metrics_batch]).map_err(
                         |e| Status::internal(format!("Failed to encode metrics batch: {}", e)),
                     )?;
 
-                // 5. Convert FlightData to arrow_flight::Result messages
-                // Note: SQL query is NOT included in metadata; clients must retain the original query
+                // 6. Convert FlightData to arrow_flight::Result messages
+                // Note: Query is NOT included in response; clients must retain the original request
                 let results: Vec<arrow_flight::Result> = flight_data
                     .into_iter()
                     .map(|fd| {
@@ -887,7 +894,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
                     })
                     .collect();
 
-                // 6. Create stream of Result messages
+                // 7. Create stream of Result messages
                 let stream = futures::stream::iter(results.into_iter().map(Ok)).boxed();
 
                 // Record metrics
@@ -899,7 +906,7 @@ impl FlightSqlService for FlightSqlServiceImpl {
                 let req = ObservabilityRequestDetails {
                     request_id: None,
                     path: "/do_action/analyze_query".to_string(),
-                    sql: Some(sql),
+                    sql: Some(sql.to_string()),
                     start_ms: start.as_millisecond(),
                     duration_ms: duration.get_milliseconds(),
                     rows: None,