Merge null to fields that are undefined (#291)

Author: trueleo

server/src/handlers/event.rs

@@ -16,7 +16,7 @@
  *
  */
 
-use actix_web::{web, HttpRequest, HttpResponse};
+use actix_web::{web, HttpRequest, HttpResponse, Responder};
 use serde_json::Value;
 use std::time::Instant;
 
@@ -33,17 +33,28 @@ use self::error::{PostError, QueryError};
 const PREFIX_TAGS: &str = "x-p-tag-";
 const PREFIX_META: &str = "x-p-meta-";
 const STREAM_NAME_HEADER_KEY: &str = "x-p-stream";
+const FILL_NULL_OPTION_KEY: &str = "fill_null";
 const SEPARATOR: char = '^';
 
-pub async fn query(_req: HttpRequest, json: web::Json<Value>) -> Result<HttpResponse, QueryError> {
+pub async fn query(
+    _req: HttpRequest,
+    json: web::Json<Value>,
+) -> Result<impl Responder, QueryError> {
     let time = Instant::now();
     let json = json.into_inner();
+
+    let fill_null = json
+        .as_object()
+        .and_then(|map| map.get(FILL_NULL_OPTION_KEY))
+        .and_then(|value| value.as_bool())
+        .unwrap_or_default();
+
     let query = Query::parse(json)?;
 
     let storage = CONFIG.storage().get_object_store();
     let query_result = query.execute(storage).await;
     let query_result = query_result
-        .map(Into::<QueryResponse>::into)
+        .map(|(records, fields)| QueryResponse::new(records, fields, fill_null))
         .map(|response| response.to_http())
         .map_err(|e| e.into());
 

server/src/query.rs

@@ -88,7 +88,7 @@ impl Query {
     pub async fn execute(
         &self,
         storage: Arc<dyn ObjectStorage + Send>,
-    ) -> Result<Vec<RecordBatch>, ExecuteError> {
+    ) -> Result<(Vec<RecordBatch>, Vec<String>), ExecuteError> {
         let dir = StorageDir::new(&self.stream_name);
         // take a look at local dir and figure out what local cache we could use for this query
         let staging_arrows = dir
@@ -128,8 +128,19 @@ impl Query {
         .map_err(ObjectStorageError::DataFusionError)?;
         // execute the query and collect results
         let df = ctx.sql(self.query.as_str()).await?;
+        // dataframe qualifies name by adding table name before columns. \
+        // For now this is just actual names
+        let fields = df
+            .schema()
+            .fields()
+            .iter()
+            .map(|f| f.name())
+            .cloned()
+            .collect_vec();
+
         let results = df.collect().await?;
-        Ok(results)
+
+        Ok((results, fields))
     }
 }
 

server/src/response.rs

@@ -17,34 +17,44 @@
  */
 
 use actix_web::http::StatusCode;
-use actix_web::{HttpResponse, HttpResponseBuilder};
-use datafusion::arrow::json;
+use actix_web::{web, Responder};
+use datafusion::arrow::json::writer::record_batches_to_json_rows;
 use datafusion::arrow::record_batch::RecordBatch;
+use itertools::Itertools;
+use serde_json::Value;
 
 pub struct QueryResponse {
     pub code: StatusCode,
-    pub body: Vec<RecordBatch>,
+    pub records: Vec<RecordBatch>,
+    pub fields: Vec<String>,
+    pub fill_null: bool,
 }
 
 impl QueryResponse {
-    pub fn to_http(&self) -> HttpResponse {
-        log::info!("{}", "Returning query results");
-        let buf = Vec::new();
-        let mut writer = json::ArrayWriter::new(buf);
-        writer.write_batches(&self.body).unwrap();
-        writer.finish().unwrap();
-
-        HttpResponseBuilder::new(self.code)
-            .content_type("json")
-            .body(writer.into_inner())
-    }
-}
-
-impl From<Vec<RecordBatch>> for QueryResponse {
-    fn from(body: Vec<RecordBatch>) -> Self {
+    pub fn new(records: Vec<RecordBatch>, fields: Vec<String>, fill_null: bool) -> Self {
         Self {
             code: StatusCode::OK,
-            body,
+            records,
+            fields,
+            fill_null,
         }
     }
+
+    pub fn to_http(&self) -> impl Responder {
+        log::info!("{}", "Returning query results");
+        let mut json_records = record_batches_to_json_rows(&self.records).unwrap();
+
+        if self.fill_null {
+            for map in &mut json_records {
+                for field in &self.fields {
+                    if !map.contains_key(field) {
+                        map.insert(field.clone(), Value::Null);
+                    }
+                }
+            }
+        }
+
+        let values = json_records.into_iter().map(Value::Object).collect_vec();
+        web::Json(values)
+    }
 }