Split SupergraphFetcher

Author: ardatan

packages/libraries/router/src/registry.rs

@@ -1,18 +1,16 @@
 use crate::registry_logger::Logger;
 use anyhow::{anyhow, Result};
+use hive_console_sdk::supergraph_fetcher::SupergraphFetcher;
 use sha2::Digest;
 use sha2::Sha256;
 use std::env;
 use std::io::Write;
 use std::thread;
 
-#[derive(Debug, Clone)]
+#[derive(Debug)]
 pub struct HiveRegistry {
-    endpoint: String,
-    key: String,
     file_name: String,
-    etag: Option<String>,
-    accept_invalid_certs: bool,
+    fetcher: SupergraphFetcher,
     pub logger: Logger,
 }
 
@@ -24,8 +22,6 @@ pub struct HiveRegistryConfig {
     schema_file_path: Option<String>,
 }
 
-static COMMIT: Option<&'static str> = option_env!("GITHUB_SHA");
-
 impl HiveRegistry {
     #[allow(clippy::new_ret_no_self)]
     pub fn new(user_config: Option<HiveRegistryConfig>) -> Result<()> {
@@ -133,11 +129,9 @@ impl HiveRegistry {
         env::set_var("APOLLO_ROUTER_HOT_RELOAD", "true");
 
         let mut registry = HiveRegistry {
-            endpoint,
-            key,
+            fetcher: SupergraphFetcher::try_new(endpoint, key, accept_invalid_certs)
+                .map_err(|e| anyhow!(e))?,
             file_name,
-            etag: None,
-            accept_invalid_certs,
             logger,
         };
 
@@ -161,52 +155,9 @@ impl HiveRegistry {
         Ok(())
     }
 
-    fn fetch_supergraph(&mut self, etag: Option<String>) -> Result<Option<String>, String> {
-        let client = reqwest::blocking::Client::builder()
-            .danger_accept_invalid_certs(self.accept_invalid_certs)
-            .build()
-            .map_err(|err| err.to_string())?;
-        let mut headers = reqwest::header::HeaderMap::new();
-
-        headers.insert(
-            reqwest::header::USER_AGENT,
-            reqwest::header::HeaderValue::from_str(
-                format!("hive-apollo-router/{}", COMMIT.unwrap_or("local")).as_str(),
-            )
-            .unwrap(),
-        );
-        headers.insert("X-Hive-CDN-Key", self.key.parse().unwrap());
-
-        if let Some(checksum) = etag {
-            headers.insert("If-None-Match", checksum.parse().unwrap());
-        }
-
-        let resp = client
-            .get(self.endpoint.as_str())
-            .headers(headers)
-            .send()
-            .map_err(|e| e.to_string())?;
-
-        match resp.headers().get("etag") {
-            Some(checksum) => {
-                let etag = checksum.to_str().map_err(|e| e.to_string())?;
-                self.update_latest_etag(Some(etag.to_string()));
-            }
-            None => {
-                self.update_latest_etag(None);
-            }
-        }
-
-        if resp.status().as_u16() == 304 {
-            return Ok(None);
-        }
-
-        Ok(Some(resp.text().map_err(|e| e.to_string())?))
-    }
-
     fn initial_supergraph(&mut self) -> Result<(), String> {
         let mut file = std::fs::File::create(self.file_name.clone()).map_err(|e| e.to_string())?;
-        let resp = self.fetch_supergraph(None)?;
+        let resp = self.fetcher.fetch_supergraph()?;
 
         match resp {
             Some(supergraph) => {
@@ -221,12 +172,8 @@ impl HiveRegistry {
         Ok(())
     }
 
-    fn update_latest_etag(&mut self, etag: Option<String>) {
-        self.etag = etag;
-    }
-
     fn poll(&mut self) {
-        match self.fetch_supergraph(self.etag.clone()) {
+        match self.fetcher.fetch_supergraph() {
             Ok(new_supergraph) => {
                 if let Some(new_supergraph) = new_supergraph {
                     let current_file = std::fs::read_to_string(self.file_name.clone())

packages/libraries/sdk-rs/src/lib.rs

@@ -1,3 +1,4 @@
 pub mod agent;
 pub mod graphql;
 pub mod persisted_documents;
+pub mod supergraph_fetcher;

packages/libraries/sdk-rs/src/supergraph_fetcher.rs

@@ -0,0 +1,86 @@
+#[derive(Debug)]
+pub struct SupergraphFetcher {
+    client: reqwest::blocking::Client,
+    endpoint: String,
+    key: String,
+    etag: Option<String>,
+}
+
+static COMMIT: Option<&'static str> = option_env!("GITHUB_SHA");
+
+impl SupergraphFetcher {
+    pub fn try_new(
+        endpoint: String,
+        key: String,
+        accept_invalid_certs: bool,
+    ) -> Result<Self, String> {
+        if endpoint.is_empty() {
+            return Err("Supergraph endpoint is required".into());
+        }
+        if key.is_empty() {
+            return Err("Supergraph key is required".into());
+        }
+        let mut endpoint = endpoint;
+        if !endpoint.ends_with("/supergraph") {
+            if endpoint.ends_with("/") {
+                endpoint.push_str("supergraph")
+            } else {
+                endpoint.push_str("/supergraph")
+            }
+        }
+        let client = reqwest::blocking::Client::builder()
+            .danger_accept_invalid_certs(accept_invalid_certs)
+            .build()
+            .map_err(|err| err.to_string())?;
+        Ok(Self {
+            client,
+            endpoint,
+            key,
+            etag: None,
+        })
+    }
+
+    pub fn fetch_supergraph(&mut self) -> Result<Option<String>, String> {
+        let mut headers = reqwest::header::HeaderMap::new();
+
+        headers.insert(
+            reqwest::header::USER_AGENT,
+            reqwest::header::HeaderValue::from_str(
+                format!("hive-apollo-router/{}", COMMIT.unwrap_or("local")).as_str(),
+            )
+            .unwrap(),
+        );
+        headers.insert("X-Hive-CDN-Key", self.key.parse().unwrap());
+
+        if let Some(checksum) = &self.etag {
+            headers.insert("If-None-Match", checksum.parse().unwrap());
+        }
+
+        let resp = self
+            .client
+            .get(self.endpoint.as_str())
+            .headers(headers)
+            .send()
+            .map_err(|e| e.to_string())?;
+
+        match resp.headers().get("etag") {
+            Some(checksum) => {
+                let etag = checksum.to_str().map_err(|e| e.to_string())?;
+                self.update_latest_etag(Some(etag.to_string()));
+            }
+            None => {
+                self.update_latest_etag(None);
+            }
+        }
+
+        if resp.status().as_u16() == 304 {
+            return Ok(None);
+        }
+
+        Ok(Some(resp.text().map_err(|e| e.to_string())?))
+    }
+
+    fn update_latest_etag(&mut self, etag: Option<String>) {
+        self.etag = etag;
+    }
+}