fix #2616 by using RawValue, and fix query rewriting

Author: analogrelay

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -38,7 +38,7 @@ workspace = true
 [features]
 default = ["hmac_rust"]
 key_auth = [] # Enables support for key-based authentication (Primary Keys and Resource Tokens)
-preview_query_engine = [] # Enables support for the PREVIEW external query engine
+preview_query_engine = ["serde_json/raw_value"] # Enables support for the PREVIEW external query engine
 hmac_rust = ["azure_core/hmac_rust"]
 hmac_openssl = ["azure_core/hmac_openssl"]
 

sdk/cosmos/azure_data_cosmos/src/query/engine.rs

@@ -20,7 +20,7 @@ pub struct PipelineResult {
     pub is_completed: bool,
 
     /// The items yielded by the pipeline.
-    pub items: Vec<Vec<u8>>,
+    pub items: Vec<Box<serde_json::value::RawValue>>,
 
     /// Additional requests that must be made before the pipeline can continue.
     pub requests: Vec<QueryRequest>,

sdk/cosmos/azure_data_cosmos/src/query/executor.rs

@@ -15,7 +15,7 @@ pub struct QueryExecutor<T: DeserializeOwned> {
     items_link: ResourceLink,
     context: Context<'static>,
     query_engine: QueryEngineRef,
-    base_request: Request,
+    base_request: Option<Request>,
     query: Query,
     pipeline: Option<OwnedQueryPipeline>,
 
@@ -45,7 +45,7 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
             items_link,
             context,
             query_engine,
-            base_request,
+            base_request: None,
             query,
             pipeline: None,
             phantom: std::marker::PhantomData,
@@ -69,8 +69,13 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
     /// An item to yield, or None if execution is complete.
     #[tracing::instrument(skip_all)]
     async fn step(&mut self) -> azure_core::Result<Option<FeedPage<T>>> {
-        let pipeline = match self.pipeline.as_mut() {
-            Some(pipeline) => pipeline,
+        let (pipeline, base_request) = match self.pipeline.as_mut() {
+            Some(pipeline) => (
+                pipeline,
+                self.base_request
+                    .as_ref()
+                    .expect("base_request should be set when pipeline is set"),
+            ),
             None => {
                 // Initialize the pipeline.
                 let query_plan = get_query_plan(
@@ -97,8 +102,16 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
                 let pipeline =
                     self.query_engine
                         .create_pipeline(&self.query.text, &query_plan, &pkranges)?;
+                self.query.text = pipeline.query().into();
+                self.base_request = Some(crate::pipeline::create_base_query_request(
+                    self.http_pipeline.url(&self.items_link),
+                    &self.query,
+                )?);
                 self.pipeline = Some(pipeline);
-                self.pipeline.as_mut().expect("we just set it")
+                (
+                    self.pipeline.as_mut().expect("we just set it"),
+                    self.base_request.as_ref().expect("we just set it"),
+                )
             }
         };
 
@@ -113,7 +126,7 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
                 let items = results
                     .items
                     .into_iter()
-                    .map(|item| serde_json::from_slice::<T>(&item))
+                    .map(|item| serde_json::from_str::<T>(item.get()))
                     .collect::<Result<Vec<_>, _>>()?;
 
                 // TODO: Provide a continuation token.
@@ -122,7 +135,7 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
 
             // No items, so make any requests we need to make and provide them to the pipeline.
             for request in results.requests {
-                let mut query_request = self.base_request.clone();
+                let mut query_request = base_request.clone();
                 query_request.insert_header(
                     constants::PARTITION_KEY_RANGE_ID,
                     request.partition_key_range_id.clone(),

sdk/cosmos/azure_data_cosmos/src/query/mod.rs

@@ -68,7 +68,7 @@ pub use engine::*;
 pub struct Query {
     /// The query text itself.
     #[serde(rename = "query")]
-    text: String,
+    pub(crate) text: String,
 
     /// A list of parameters used in the query and their associated value.
     #[serde(skip_serializing_if = "Vec::is_empty")] // Don't serialize an empty array.

sdk/cosmos/azure_data_cosmos/tests/framework/query_engine.rs

@@ -6,6 +6,7 @@ use std::{collections::VecDeque, sync::Mutex};
 use serde::{Deserialize, Serialize};
 
 use azure_data_cosmos::query::{PipelineResult, QueryEngine, QueryPipeline};
+use serde_json::value::RawValue;
 
 #[derive(Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -202,7 +203,19 @@ impl QueryPipeline for MockQueryPipeline {
             if let Some((index, _)) = state {
                 // Add this item to the result.
                 if let Some(item) = self.partitions[index].pop_item()? {
-                    items.push(item);
+                    let s = String::from_utf8(item).map_err(|_| {
+                        azure_core::Error::message(
+                            azure_core::error::ErrorKind::DataConversion,
+                            "item is not valid UTF-8",
+                        )
+                    })?;
+                    let raw_value = RawValue::from_string(s).map_err(|_| {
+                        azure_core::Error::message(
+                            azure_core::error::ErrorKind::DataConversion,
+                            "failed to create RawValue",
+                        )
+                    })?;
+                    items.push(raw_value);
                 }
             } else {
                 // All partitions are exhausted, or have no items to produce