feat(hive-console-sdk): async SupergraphFetcher, retry_count in SupergraphFetcher & user_agent in PersistedDocumentsManager (#7246)

Author: ardatan

.changeset/lemon-hairs-sort.md

@@ -0,0 +1,30 @@
+---
+'hive-console-sdk-rs': minor
+---
+Breaking;
+- `SupergraphFetcher` now has two different modes: async and sync. You can choose between `SupergraphFetcherAsyncClient` and `SupergraphFetcherSyncClient` based on your needs. See the examples at the bottom.
+- `SupergraphFetcher` now has a new `retry_count` parameter to specify how many times to retry fetching the supergraph in case of failures.
+- `PersistedDocumentsManager` new needs `user_agent` parameter to be sent to Hive Console when fetching persisted queries.
+- `UsageAgent::new` is now `UsageAgent::try_new` and it returns a `Result` with `Arc`, so you can freely clone it across threads. This change was made to handle potential errors during the creation of the HTTP client. Make sure to handle the `Result` when creating a `UsageAgent`.
+
+```rust
+// Sync Mode
+let fetcher = SupergraphFetcher::try_new_sync(/* params */)
+.map_err(|e| anyhow!("Failed to create SupergraphFetcher: {}", e))?;
+
+// Use the fetcher to fetch the supergraph (Sync)
+let supergraph = fetcher
+    .fetch_supergraph()
+    .map_err(|e| anyhow!("Failed to fetch supergraph: {}", e))?;
+
+// Async Mode
+
+let fetcher = SupergraphFetcher::try_new_async(/* params */)
+.map_err(|e| anyhow!("Failed to create SupergraphFetcher: {}", e))?;
+
+// Use the fetcher to fetch the supergraph (Async)
+let supergraph = fetcher
+    .fetch_supergraph()
+    .await
+    .map_err(|e| anyhow!("Failed to fetch supergraph: {}", e))?;
+```

configs/cargo/Cargo.lock

@@ -24,6 +24,8 @@ use std::time::Duration;
 use tower::{BoxError, ServiceBuilder, ServiceExt};
 use tracing::{debug, info, warn};
 
+use crate::consts::PLUGIN_VERSION;
+
 pub static PERSISTED_DOCUMENT_HASH_KEY: &str = "hive::persisted_document_hash";
 
 #[derive(Clone, Debug, Deserialize, JsonSchema, Default)]
@@ -107,6 +109,7 @@ impl PersistedDocumentsPlugin {
                 Duration::from_secs(config.request_timeout.unwrap_or(15)),
                 config.retry_count.unwrap_or(3),
                 config.cache_size.unwrap_or(1000),
+                format!("hive-apollo-router/{}", PLUGIN_VERSION),
             ))),
             allow_arbitrary_documents,
         })

packages/libraries/router/src/persisted_documents.rs

@@ -2,6 +2,7 @@ use crate::consts::PLUGIN_VERSION;
 use crate::registry_logger::Logger;
 use anyhow::{anyhow, Result};
 use hive_console_sdk::supergraph_fetcher::SupergraphFetcher;
+use hive_console_sdk::supergraph_fetcher::SupergraphFetcherSyncState;
 use sha2::Digest;
 use sha2::Sha256;
 use std::env;
@@ -12,7 +13,7 @@ use std::time::Duration;
 #[derive(Debug)]
 pub struct HiveRegistry {
     file_name: String,
-    fetcher: SupergraphFetcher,
+    fetcher: SupergraphFetcher<SupergraphFetcherSyncState>,
     pub logger: Logger,
 }
 
@@ -122,16 +123,19 @@ impl HiveRegistry {
         env::set_var("APOLLO_ROUTER_SUPERGRAPH_PATH", file_name.clone());
         env::set_var("APOLLO_ROUTER_HOT_RELOAD", "true");
 
-        let mut registry = HiveRegistry {
-            fetcher: SupergraphFetcher::try_new(
-                endpoint,
-                key,
-                format!("hive-apollo-router/{}", PLUGIN_VERSION),
-                Duration::from_secs(5),
-                Duration::from_secs(60),
-                accept_invalid_certs,
-            )
-            .map_err(|e| anyhow!("Failed to create SupergraphFetcher: {}", e))?,
+        let fetcher = SupergraphFetcher::try_new_sync(
+            endpoint,
+            &key,
+            format!("hive-apollo-router/{}", PLUGIN_VERSION),
+            Duration::from_secs(5),
+            Duration::from_secs(60),
+            accept_invalid_certs,
+            3,
+        )
+        .map_err(|e| anyhow!("Failed to create SupergraphFetcher: {}", e))?;
+
+        let registry = HiveRegistry {
+            fetcher,
             file_name,
             logger,
         };
@@ -156,9 +160,12 @@ impl HiveRegistry {
         Ok(())
     }
 
-    fn initial_supergraph(&mut self) -> Result<(), String> {
+    fn initial_supergraph(&self) -> Result<(), String> {
         let mut file = std::fs::File::create(self.file_name.clone()).map_err(|e| e.to_string())?;
-        let resp = self.fetcher.fetch_supergraph()?;
+        let resp = self
+            .fetcher
+            .fetch_supergraph()
+            .map_err(|err| err.to_string())?;
 
         match resp {
             Some(supergraph) => {
@@ -173,7 +180,7 @@ impl HiveRegistry {
         Ok(())
     }
 
-    fn poll(&mut self) {
+    fn poll(&self) {
         match self.fetcher.fetch_supergraph() {
             Ok(new_supergraph) => {
                 if let Some(new_supergraph) = new_supergraph {

packages/libraries/router/src/registry.rs

@@ -8,6 +8,7 @@ use core::ops::Drop;
 use futures::StreamExt;
 use graphql_parser::parse_schema;
 use graphql_parser::schema::Document;
+use hive_console_sdk::agent::UsageAgentExt;
 use hive_console_sdk::agent::{ExecutionReport, UsageAgent};
 use http::HeaderValue;
 use rand::Rng;
@@ -245,8 +246,8 @@ impl Plugin for UsagePlugin {
 
         let agent = if enabled {
             let flush_interval = Duration::from_secs(flush_interval);
-            let agent = Arc::new(UsageAgent::new(
-                token.expect("token is set"),
+            let agent = UsageAgent::try_new(
+                &token.expect("token is set"),
                 endpoint,
                 target_id,
                 buffer_size,
@@ -255,7 +256,8 @@ impl Plugin for UsagePlugin {
                 accept_invalid_certs,
                 flush_interval,
                 format!("hive-apollo-router/{}", PLUGIN_VERSION),
-            ));
+            )
+            .map_err(Box::new)?;
             start_flush_interval(agent.clone());
             Some(agent)
         } else {

packages/libraries/router/src/usage.rs

@@ -1,5 +1,6 @@
 use super::graphql::OperationProcessor;
 use graphql_parser::schema::Document;
+use reqwest::header::{HeaderMap, HeaderValue};
 use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
 use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
 use serde::Serialize;
@@ -103,17 +104,13 @@ impl Buffer {
         Ok(reports)
     }
 }
-
-#[derive(Clone)]
 pub struct UsageAgent {
-    token: String,
     buffer_size: usize,
     endpoint: String,
-    buffer: Arc<Buffer>,
-    processor: Arc<OperationProcessor>,
+    buffer: Buffer,
+    processor: OperationProcessor,
     client: ClientWithMiddleware,
     flush_interval: Duration,
-    user_agent: String,
 }
 
 fn non_empty_string(value: Option<String>) -> Option<String> {
@@ -130,14 +127,18 @@ pub enum AgentError {
     Forbidden,
     #[error("unable to send report: rate limited")]
     RateLimited,
+    #[error("invalid token provided: {0}")]
+    InvalidToken(String),
+    #[error("unable to instantiate the http client for reports sending: {0}")]
+    HTTPClientCreationError(reqwest::Error),
     #[error("unable to send report: {0}")]
     Unknown(String),
 }
 
 impl UsageAgent {
     #[allow(clippy::too_many_arguments)]
-    pub fn new(
-        token: String,
+    pub fn try_new(
+        token: &str,
         endpoint: String,
         target_id: Option<String>,
         buffer_size: usize,
@@ -146,22 +147,35 @@ impl UsageAgent {
         accept_invalid_certs: bool,
         flush_interval: Duration,
         user_agent: String,
-    ) -> Self {
-        let processor = Arc::new(OperationProcessor::new());
-
+    ) -> Result<Arc<Self>, AgentError> {
         let retry_policy = ExponentialBackoff::builder().build_with_max_retries(3);
 
+        let mut default_headers = HeaderMap::new();
+
+        default_headers.insert("X-Usage-API-Version", HeaderValue::from_static("2"));
+
+        let mut authorization_header = HeaderValue::from_str(&format!("Bearer {}", token))
+            .map_err(|_| AgentError::InvalidToken(token.to_string()))?;
+
+        authorization_header.set_sensitive(true);
+
+        default_headers.insert(reqwest::header::AUTHORIZATION, authorization_header);
+
+        default_headers.insert(
+            reqwest::header::CONTENT_TYPE,
+            HeaderValue::from_static("application/json"),
+        );
+
         let reqwest_agent = reqwest::Client::builder()
             .danger_accept_invalid_certs(accept_invalid_certs)
             .connect_timeout(connect_timeout)
             .timeout(request_timeout)
+            .user_agent(user_agent)
             .build()
-            .map_err(|err| err.to_string())
-            .expect("Couldn't instantiate the http client for reports sending!");
+            .map_err(AgentError::HTTPClientCreationError)?;
         let client = ClientBuilder::new(reqwest_agent)
             .with(RetryTransientMiddleware::new_with_policy(retry_policy))
             .build();
-        let buffer = Arc::new(Buffer::new());
 
         let mut endpoint = endpoint;
 
@@ -171,16 +185,14 @@ impl UsageAgent {
             }
         }
 
-        UsageAgent {
-            buffer,
-            processor,
-            endpoint,
-            token,
+        Ok(Arc::new(Self {
             buffer_size,
+            endpoint,
+            buffer: Buffer::new(),
+            processor: OperationProcessor::new(),
             client,
             flush_interval,
-            user_agent,
-        }
+        }))
     }
 
     fn produce_report(&self, reports: Vec<ExecutionReport>) -> Result<Report, AgentError> {
@@ -256,28 +268,13 @@ impl UsageAgent {
         Ok(report)
     }
 
-    pub fn add_report(&self, execution_report: ExecutionReport) -> Result<(), AgentError> {
-        let size = self.buffer.push(execution_report)?;
-
-        self.flush_if_full(size)?;
-
-        Ok(())
-    }
-
     pub async fn send_report(&self, report: Report) -> Result<(), AgentError> {
         let report_body =
             serde_json::to_vec(&report).map_err(|e| AgentError::Unknown(e.to_string()))?;
         // Based on https://the-guild.dev/graphql/hive/docs/specs/usage-reports#data-structure
         let resp = self
             .client
             .post(&self.endpoint)
-            .header("X-Usage-API-Version", "2")
-            .header(
-                reqwest::header::AUTHORIZATION,
-                format!("Bearer {}", self.token),
-            )
-            .header(reqwest::header::USER_AGENT, self.user_agent.to_string())
-            .header(reqwest::header::CONTENT_TYPE, "application/json")
             .header(reqwest::header::CONTENT_LENGTH, report_body.len())
             .body(report_body)
             .send()
@@ -297,17 +294,6 @@ impl UsageAgent {
         }
     }
 
-    pub fn flush_if_full(&self, size: usize) -> Result<(), AgentError> {
-        if size >= self.buffer_size {
-            let cloned_self = self.clone();
-            tokio::task::spawn(async move {
-                cloned_self.flush().await;
-            });
-        }
-
-        Ok(())
-    }
-
     pub async fn flush(&self) {
         let execution_reports = match self.buffer.drain() {
             Ok(res) => res,
@@ -345,3 +331,29 @@ impl UsageAgent {
         }
     }
 }
+
+pub trait UsageAgentExt {
+    fn add_report(&self, execution_report: ExecutionReport) -> Result<(), AgentError>;
+    fn flush_if_full(&self, size: usize) -> Result<(), AgentError>;
+}
+
+impl UsageAgentExt for Arc<UsageAgent> {
+    fn flush_if_full(&self, size: usize) -> Result<(), AgentError> {
+        if size >= self.buffer_size {
+            let cloned_self = self.clone();
+            tokio::task::spawn(async move {
+                cloned_self.flush().await;
+            });
+        }
+
+        Ok(())
+    }
+
+    fn add_report(&self, execution_report: ExecutionReport) -> Result<(), AgentError> {
+        let size = self.buffer.push(execution_report)?;
+
+        self.flush_if_full(size)?;
+
+        Ok(())
+    }
+}