From 40c439024b1a53b9bc31b1f6859f537ded0699b2 Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Fri, 15 Aug 2025 21:45:57 +0100 Subject: [PATCH 1/5] selects only id when checking existence of data Signed-off-by: Aminu Oluwaseun Joshua --- crates/key-value-azure/src/store.rs | 32 +++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index 518246e03d..95baea73a8 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -190,7 +190,7 @@ impl Store for AzureCosmosStore { } async fn exists(&self, key: &str) -> Result { - Ok(self.get_entity::(key).await?.is_some()) + Ok(self.get_entity_by_id::(key).await?.is_some()) } async fn get_keys(&self) -> Result, Error> { @@ -424,6 +424,28 @@ impl AzureCosmosStore { .map(|(p, _)| p.clone())) } + async fn get_entity_by_id(&self, key: &str) -> Result, Error> + where + F: CosmosEntity + Send + Sync + serde::de::DeserializeOwned + Clone, + { + let query = self + .client + .query_documents(Query::new(self.get_id_query(key))) + .query_cross_partition(true) + .max_item_count(1); + + // There can be no duplicated keys, so we create the stream and only take the first result. + let mut stream = query.into_stream::(); + let Some(res) = stream.next().await else { + return Ok(None); + }; + Ok(res + .map_err(log_error)? + .results + .first() + .map(|(p, _)| p.clone())) + } + async fn get_keys(&self) -> Result, Error> { let query = self .client @@ -446,8 +468,14 @@ impl AzureCosmosStore { query } + fn get_id_query(&self, key: &str) -> String { + let mut query = format!("SELECT c.id FROM c WHERE c.id='{key}'"); + self.append_store_id(&mut query, true); + query + } + fn get_keys_query(&self) -> String { - let mut query = "SELECT * FROM c".to_owned(); + let mut query = "SELECT c.id, c.store_id FROM c".to_owned(); self.append_store_id(&mut query, false); query } From c9c16ea8fa5c4c8d7388484765400a90050ab42e Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Mon, 18 Aug 2025 18:08:38 +0100 Subject: [PATCH 2/5] cleaner code Signed-off-by: Aminu Oluwaseun Joshua --- crates/key-value-azure/src/store.rs | 35 ++++++++++------------------- 1 file changed, 12 insertions(+), 23 deletions(-) diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index 95baea73a8..e90d5cfa64 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -190,7 +190,18 @@ impl Store for AzureCosmosStore { } async fn exists(&self, key: &str) -> Result { - Ok(self.get_entity_by_id::(key).await?.is_some()) + let mut stream = self + .client + .query_documents(Query::new(self.get_id_query(key))) + .query_cross_partition(true) + .max_item_count(1) + .into_stream::(); + + match stream.next().await { + Some(Ok(res)) => Ok(!res.results.is_empty()), + Some(Err(e)) => Err(log_error(e)), + None => Ok(false), + } } async fn get_keys(&self) -> Result, Error> { @@ -424,28 +435,6 @@ impl AzureCosmosStore { .map(|(p, _)| p.clone())) } - async fn get_entity_by_id(&self, key: &str) -> Result, Error> - where - F: CosmosEntity + Send + Sync + serde::de::DeserializeOwned + Clone, - { - let query = self - .client - .query_documents(Query::new(self.get_id_query(key))) - .query_cross_partition(true) - .max_item_count(1); - - // There can be no duplicated keys, so we create the stream and only take the first result. - let mut stream = query.into_stream::(); - let Some(res) = stream.next().await else { - return Ok(None); - }; - Ok(res - .map_err(log_error)? - .results - .first() - .map(|(p, _)| p.clone())) - } - async fn get_keys(&self) -> Result, Error> { let query = self .client From 21bd073ae69d7c9fb70fca502ac440949ccdf681 Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Mon, 18 Aug 2025 18:18:12 +0100 Subject: [PATCH 3/5] store_id added to query Signed-off-by: Aminu Oluwaseun Joshua --- crates/key-value-azure/src/store.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index e90d5cfa64..49b38c2464 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -458,7 +458,7 @@ impl AzureCosmosStore { } fn get_id_query(&self, key: &str) -> String { - let mut query = format!("SELECT c.id FROM c WHERE c.id='{key}'"); + let mut query = format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"); self.append_store_id(&mut query, true); query } From f9079a4e91822a7a49cdc38f8f90fe5b1dc8f8df Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Tue, 19 Aug 2025 23:56:37 +0100 Subject: [PATCH 4/5] handle optional store_id Signed-off-by: Aminu Oluwaseun Joshua --- crates/key-value-azure/src/store.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index 49b38c2464..6e0ed4dcd5 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -458,13 +458,19 @@ impl AzureCosmosStore { } fn get_id_query(&self, key: &str) -> String { - let mut query = format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"); + let mut query = match self.store_id { + Some(_) => format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"), + None => format!("SELECT c.id FROM c WHERE c.id='{key}'"), + }; self.append_store_id(&mut query, true); query } fn get_keys_query(&self) -> String { - let mut query = "SELECT c.id, c.store_id FROM c".to_owned(); + let mut query = match self.store_id { + Some(_) => "SELECT c.id, c.store_id FROM c".to_owned(), + None => "SELECT c.id FROM c".to_owned(), + }; self.append_store_id(&mut query, false); query } From 3d8f57440ffdc188474fb28f898e3cff5ec66aff Mon Sep 17 00:00:00 2001 From: Aminu Oluwaseun Joshua Date: Wed, 20 Aug 2025 19:24:47 +0100 Subject: [PATCH 5/5] revert Signed-off-by: Aminu Oluwaseun Joshua --- crates/key-value-azure/src/store.rs | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index 6e0ed4dcd5..49b38c2464 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -458,19 +458,13 @@ impl AzureCosmosStore { } fn get_id_query(&self, key: &str) -> String { - let mut query = match self.store_id { - Some(_) => format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"), - None => format!("SELECT c.id FROM c WHERE c.id='{key}'"), - }; + let mut query = format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"); self.append_store_id(&mut query, true); query } fn get_keys_query(&self) -> String { - let mut query = match self.store_id { - Some(_) => "SELECT c.id, c.store_id FROM c".to_owned(), - None => "SELECT c.id FROM c".to_owned(), - }; + let mut query = "SELECT c.id, c.store_id FROM c".to_owned(); self.append_store_id(&mut query, false); query }