Cosmos: fix #2616 by using RawValue, and fix query rewriting (#2617)

* fix #2616 by using RawValue, and fix query rewriting

* remove unused variable

* replace  'expect' with 'unwrap'

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>,
 
@@ -37,15 +37,13 @@ impl<T: DeserializeOwned + 'static> QueryExecutor<T> {
     ) -> azure_core::Result<Self> {
         let items_link = container_link.feed(ResourceType::Items);
         let context = options.method_options.context.into_owned();
-        let base_request =
-            pipeline::create_base_query_request(http_pipeline.url(&items_link), &query)?;
         Ok(Self {
             http_pipeline,
             container_link,
             items_link,
             context,
             query_engine,
-            base_request,
+            base_request: None,
             query,
             pipeline: None,
             phantom: std::marker::PhantomData,
@@ -69,8 +67,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 +100,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().unwrap(),
+                    self.base_request.as_ref().unwrap(),
+                )
             }
         };
 
@@ -113,7 +124,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 +133,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")]
@@ -104,10 +105,10 @@ impl PartitionState {
         self.next_continuation = continuation;
     }
 
-    pub fn pop_item(&mut self) -> azure_core::Result<Option<Vec<u8>>> {
+    pub fn pop_item(&mut self) -> azure_core::Result<Option<Box<RawValue>>> {
         match self.queue.pop_front() {
             Some(item) => {
-                let item = serde_json::to_vec(&item)?;
+                let item = serde_json::value::to_raw_value(&item)?;
                 Ok(Some(item))
             }
             None => Ok(None),