Merge pull request #3017 from kate-goldenring/enable-connection-pooling-in-cosmos-sdk

Use connection pooling when querying CosmosDB

Author: kate-goldenring

Cargo.lock

@@ -17,6 +17,12 @@ futures = { workspace = true }
 serde = { workspace = true }
 spin-core = { path = "../core" }
 spin-factor-key-value = { path = "../factor-key-value" }
+reqwest = { version = "0.12", default-features = false }
 
 [lints]
 workspace = true
+
+[features]
+# Enables reusing connections to the Azure Cosmos DB service.
+connection-pooling = []
+default = ["connection-pooling"]

crates/key-value-azure/Cargo.toml

@@ -1,7 +1,8 @@
-use anyhow::Result;
-use azure_data_cosmos::prelude::Operation;
+use anyhow::{Context, Result};
 use azure_data_cosmos::{
-    prelude::{AuthorizationToken, CollectionClient, CosmosClient, Query},
+    prelude::{
+        AuthorizationToken, CollectionClient, CosmosClient, CosmosClientBuilder, Operation, Query,
+    },
     CosmosEntity,
 };
 use futures::StreamExt;
@@ -88,14 +89,28 @@ impl KeyValueAzureCosmos {
                 )
             }
         };
-        let cosmos_client = CosmosClient::new(account, token);
+        let cosmos_client = cosmos_client(account, token)?;
         let database_client = cosmos_client.database_client(database);
         let client = database_client.collection_client(container);
 
         Ok(Self { client, app_id })
     }
 }
 
+fn cosmos_client(account: impl Into<String>, token: AuthorizationToken) -> Result<CosmosClient> {
+    if cfg!(feature = "connection-pooling") {
+        let client = reqwest::ClientBuilder::new()
+            .build()
+            .context("failed to build reqwest client")?;
+        let transport_options = azure_core::TransportOptions::new(std::sync::Arc::new(client));
+        Ok(CosmosClientBuilder::new(account, token)
+            .transport(transport_options)
+            .build())
+    } else {
+        Ok(CosmosClient::new(account, token))
+    }
+}
+
 #[async_trait]
 impl StoreManager for KeyValueAzureCosmos {
     async fn get(&self, name: &str) -> Result<Arc<dyn Store>, Error> {