Normalize query HTTP response shapes

Author: shisoft

crates/morpheus-http-client/src/lib.rs

@@ -397,7 +397,7 @@ impl MorpheusHttpClient {
             .send()
             .await
             .map_err(ClientError::http)?;
-        parse_json(res).await
+        parse_query_result_json(res).await
     }
 
     pub async fn query_explain(
@@ -554,6 +554,100 @@ async fn parse_json<T: serde::de::DeserializeOwned>(
     }))
 }
 
+async fn parse_query_result_json(
+    res: reqwest::Response,
+) -> Result<ApiResponse<QueryResult>, ClientError> {
+    let status = res.status();
+    if status.is_success() {
+        let body = res
+            .json::<serde_json::Value>()
+            .await
+            .map_err(ClientError::http)?;
+        return normalize_query_result_response(body);
+    }
+
+    let body = match res.json::<ApiErrorBody>().await {
+        Ok(body) => body,
+        Err(_) => ApiErrorBody {
+            code: "http_error".to_string(),
+            message: format!("HTTP {status}"),
+        },
+    };
+    Err(ClientError::api(ApiError {
+        status: status.as_u16(),
+        body,
+    }))
+}
+
+fn normalize_query_result_response(
+    body: serde_json::Value,
+) -> Result<ApiResponse<QueryResult>, ClientError> {
+    if let Ok(parsed) = serde_json::from_value::<ApiResponse<QueryResult>>(body.clone()) {
+        return Ok(parsed);
+    }
+
+    #[derive(serde::Deserialize)]
+    struct LegacyQueryResponse {
+        data: Vec<CellDto>,
+    }
+
+    if let Ok(parsed) = serde_json::from_value::<LegacyQueryResponse>(body) {
+        return Ok(ApiResponse {
+            data: QueryResult {
+                cells: parsed.data,
+                hit_table: Default::default(),
+            },
+        });
+    }
+
+    Err(ClientError::invalid_argument(
+        "query response body did not match expected formats",
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+
+    #[test]
+    fn normalize_query_result_response_accepts_structured_shape() {
+        let parsed = normalize_query_result_response(json!({
+            "data": {
+                "cells": [{
+                    "id": "JEFLHiLiF34zS6AVQMKYQy",
+                    "header": {},
+                    "data": {"content": "x"}
+                }],
+                "hit_table": {
+                    "JEFLHiLiF34zS6AVQMKYQy": {
+                        "1:bm25": 0.5
+                    }
+                }
+            }
+        }))
+        .expect("structured response should parse");
+
+        assert_eq!(parsed.data.cells.len(), 1);
+        assert_eq!(parsed.data.hit_table.len(), 1);
+    }
+
+    #[test]
+    fn normalize_query_result_response_accepts_legacy_array_shape() {
+        let parsed = normalize_query_result_response(json!({
+            "data": [{
+                "id": "JEFLHiLiF34zS6AVQMKYQy",
+                "header": {},
+                "data": {"content": "x"}
+            }]
+        }))
+        .expect("legacy response should parse");
+
+        assert_eq!(parsed.data.cells.len(), 1);
+        assert!(parsed.data.hit_table.is_empty());
+    }
+}
+
 pub fn validate_base58_id(s: &str) -> Result<(), ClientError> {
     let bytes = bs58::decode(s)
         .into_vec()