hasura · hallettj · Apr 12, 2024 · Apr 12, 2024 · Apr 12, 2024 · Apr 12, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/dc-api-types/src/query_response.rs b/crates/dc-api-types/src/query_response.rs
@@ -42,13 +42,29 @@ pub struct ForEachRow {
     pub query: RowSet,
 }
 
+/// A row set must contain either rows, or aggregates, or possibly both
 #[skip_serializing_none]
-#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
-pub struct RowSet {
-    /// The results of the aggregates returned by the query
-    pub aggregates: Option<HashMap<String, serde_json::Value>>,
-    /// The rows returned by the query, corresponding to the query's fields
-    pub rows: Option<Vec<HashMap<String, ResponseFieldValue>>>,
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum RowSet {
+    Aggregate {
+        /// The results of the aggregates returned by the query
+        aggregates: HashMap<String, serde_json::Value>,
+        /// The rows returned by the query, corresponding to the query's fields
+        rows: Option<Vec<HashMap<String, ResponseFieldValue>>>,
+    },
+    Rows {
+        /// Rows returned by a query that did not request aggregates.
+        rows: Vec<HashMap<String, ResponseFieldValue>>,
+    },
+}
+
+impl Default for RowSet {
+    fn default() -> Self {
+        RowSet::Rows {
+            rows: Default::default(),
+        }
+    }
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

diff --git a/crates/dc-api/src/interface_types/json_response.rs b/crates/dc-api/src/interface_types/json_response.rs
diff --git a/crates/dc-api/src/interface_types/mod.rs b/crates/dc-api/src/interface_types/mod.rs
@@ -1,4 +1,3 @@
 mod agent_error;
-mod json_response;
 
-pub use self::{agent_error::AgentError, json_response::JsonResponse};
+pub use self::agent_error::AgentError;
diff --git a/crates/dc-api/src/lib.rs b/crates/dc-api/src/lib.rs
@@ -1,3 +1,3 @@
 mod interface_types;
 
-pub use self::interface_types::{AgentError, JsonResponse};
+pub use self::interface_types::AgentError;
diff --git a/crates/mongodb-agent-common/src/query/execute_query_request.rs b/crates/mongodb-agent-common/src/query/execute_query_request.rs
@@ -1,17 +1,15 @@
 use anyhow::anyhow;
-use bytes::Bytes;
-use dc_api::JsonResponse;
 use dc_api_types::{QueryRequest, QueryResponse};
 use futures_util::TryStreamExt;
-use mongodb::bson::{doc, Document};
+use mongodb::bson::{self, doc, Document};
 
 use super::pipeline::{pipeline_for_query_request, ResponseShape};
 use crate::{interface_types::MongoAgentError, mongodb::CollectionTrait};
 
 pub async fn execute_query_request(
     collection: &impl CollectionTrait<Document>,
     query_request: QueryRequest,
-) -> Result<JsonResponse<QueryResponse>, MongoAgentError> {
+) -> Result<QueryResponse, MongoAgentError> {
     let (pipeline, response_shape) = pipeline_for_query_request(&query_request)?;
     tracing::debug!(pipeline = %serde_json::to_string(&pipeline).unwrap(), "aggregate pipeline");
 
@@ -33,9 +31,8 @@ pub async fn execute_query_request(
             ))
         })?,
     };
+    tracing::debug!(response_document = %serde_json::to_string(&response_document).unwrap(), "response from MongoDB");
 
-    let bytes: Bytes = serde_json::to_vec(&response_document)
-        .map_err(MongoAgentError::Serialization)?
-        .into();
-    Ok(JsonResponse::Serialized(bytes))
+    let response = bson::from_document(response_document)?;
+    Ok(response)
 }
diff --git a/crates/mongodb-agent-common/src/query/foreach.rs b/crates/mongodb-agent-common/src/query/foreach.rs
@@ -176,14 +176,14 @@ mod tests {
                 "$facet": {
                     "__FACET___0": [
                         { "$match": { "$and": [{ "artistId": {"$eq":1 }}]}},
-                        { "$replaceWith": { 
+                        { "$replaceWith": {
                             "albumId": { "$ifNull": ["$albumId", null] },
                             "title": { "$ifNull": ["$title", null] }
                         } },
                     ],
                     "__FACET___1": [
                         { "$match": { "$and": [{ "artistId": {"$eq":2}}]}},
-                        { "$replaceWith": { 
+                        { "$replaceWith": {
                             "albumId": { "$ifNull": ["$albumId", null] },
                             "title": { "$ifNull": ["$title", null] }
                         } },
@@ -248,9 +248,7 @@ mod tests {
                 })?)]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
 
         Ok(())
@@ -364,7 +362,6 @@ mod tests {
             ]
         }))?;
 
-
         let mut collection = MockCollectionTrait::new();
         collection
             .expect_aggregate()
@@ -398,9 +395,7 @@ mod tests {
                 })?)]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
 
         Ok(())

diff --git a/crates/mongodb-agent-common/src/query/make_selector.rs b/crates/mongodb-agent-common/src/query/make_selector.rs
@@ -10,7 +10,7 @@ use mongodb_support::BsonScalarType;
 
 use crate::{
     comparison_function::ComparisonFunction, interface_types::MongoAgentError,
-    query::serialization::json_to_bson_scalar, query::column_ref::column_ref,
+    query::column_ref::column_ref, query::serialization::json_to_bson_scalar,
 };
 
 use BinaryArrayComparisonOperator as ArrOp;
@@ -112,7 +112,7 @@ fn make_selector_helper(
                         "comparisons between columns",
                     )),
                     ArrayComparisonValue::Variable(name) => {
-                        Ok(variable_to_mongo_expression(variables, name, value_type)?.into())
+                        variable_to_mongo_expression(variables, name, value_type)
                     }
                 })
                 .collect::<Result<_, MongoAgentError>>()?;
@@ -149,11 +149,10 @@ fn variable_to_mongo_expression(
     variables: Option<&BTreeMap<String, serde_json::Value>>,
     variable: &str,
     value_type: &str,
-) -> Result<bson::Document, MongoAgentError> {
+) -> Result<bson::Bson, MongoAgentError> {
     let value = variables
         .and_then(|vars| vars.get(variable))
         .ok_or_else(|| MongoAgentError::VariableNotDefined(variable.to_owned()))?;
-    Ok(doc! {
-        "$literal": bson_from_scalar_value(value, value_type)?
-    })
+
+    bson_from_scalar_value(value, value_type)
 }
diff --git a/crates/mongodb-agent-common/src/query/mod.rs b/crates/mongodb-agent-common/src/query/mod.rs
@@ -9,7 +9,6 @@ mod pipeline;
 mod relations;
 pub mod serialization;
 
-use dc_api::JsonResponse;
 use dc_api_types::{QueryRequest, QueryResponse, Target};
 use mongodb::bson::Document;
 
@@ -28,7 +27,7 @@ pub fn collection_name(query_request_target: &Target) -> String {
 pub async fn handle_query_request(
     config: &MongoConfig,
     query_request: QueryRequest,
-) -> Result<JsonResponse<QueryResponse>, MongoAgentError> {
+) -> Result<QueryResponse, MongoAgentError> {
     tracing::debug!(?config, query_request = %serde_json::to_string(&query_request).unwrap(), "executing query");
 
     let database = config.client.database(&config.database);
@@ -91,9 +90,7 @@ mod tests {
                 ]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
         Ok(())
     }
@@ -170,9 +167,7 @@ mod tests {
                 })?)]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
         Ok(())
     }
@@ -255,9 +250,7 @@ mod tests {
                 })?)]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
         Ok(())
     }
@@ -317,9 +310,7 @@ mod tests {
                 })?)]))
             });
 
-        let result = execute_query_request(&collection, query_request)
-            .await?
-            .into_value()?;
+        let result = execute_query_request(&collection, query_request).await?;
         assert_eq!(expected_response, result);
         Ok(())
     }