cache-control example

uses cache-control headers from subgraphs to determine an overall cache-control policy

partially addresses apollographql#326

Author: lennyburdette

Cargo.lock

@@ -756,6 +756,18 @@ dependencies = [
  "either",
 ]
 
+[[package]]
+name = "cache-control"
+version = "0.1.0"
+dependencies = [
+ "anyhow",
+ "apollo-router",
+ "http",
+ "serde_json",
+ "tokio",
+ "tower",
+]
+
 [[package]]
 name = "cargo-scaffold"
 version = "0.8.7"

Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "apollo-router-scaffold",
     "examples/add-timestamp-header/rhai",
     "examples/async-auth/rust",
+    "examples/cache-control/rhai",
     "examples/context/rust",
     "examples/cookies-to-headers/rhai",
     "examples/data-response-mutate/rhai",
@@ -35,7 +36,7 @@ strip = "debuginfo"
 incremental = false
 
 # If building a dhat feature, you must use this profile
-# e.g. heap allocation tracing: cargo build --profile release-dhat --features dhat-heap 
+# e.g. heap allocation tracing: cargo build --profile release-dhat --features dhat-heap
 # e.g. heap and ad-hoc allocation tracing: cargo build --profile release-dhat --features dhat-heap,dhat-ad-hoc
 [profile.release-dhat]
 inherits = "release"

apollo-router/src/plugins/rhai/mod.rs

@@ -641,6 +641,7 @@ macro_rules! gen_map_request {
                                 .and_data(body.data)
                                 .and_label(body.label)
                                 .and_path(body.path)
+                                .headers(HeaderMap::new())
                                 .build()
                         } else {
                             $base::Response::error_builder()

apollo-router/src/plugins/traffic_shaping/cache.rs

@@ -179,6 +179,7 @@ where
             .data(data)
             .extensions(Object::new())
             .context(request.context)
+            .headers(http::HeaderMap::new())
             .build())
     }
 }

apollo-router/src/services/subgraph.rs

@@ -138,6 +138,7 @@ impl Response {
         extensions: Object,
         status_code: Option<StatusCode>,
         context: Context,
+        headers: http::HeaderMap<http::HeaderValue>,
     ) -> Response {
         // Build a response
         let res = graphql::Response::builder()
@@ -149,11 +150,13 @@ impl Response {
             .build();
 
         // Build an http Response
-        let response = http::Response::builder()
+        let mut response = http::Response::builder()
             .status(status_code.unwrap_or(StatusCode::OK))
             .body(res)
             .expect("Response is serializable; qed");
 
+        *response.headers_mut() = headers;
+
         Self { response, context }
     }
 
@@ -172,6 +175,7 @@ impl Response {
         extensions: JsonMap<ByteString, Value>,
         status_code: Option<StatusCode>,
         context: Option<Context>,
+        headers: Option<http::HeaderMap<http::HeaderValue>>,
     ) -> Response {
         Response::new(
             label,
@@ -181,6 +185,7 @@ impl Response {
             extensions,
             status_code,
             context.unwrap_or_default(),
+            headers.unwrap_or_default(),
         )
     }
 
@@ -201,6 +206,7 @@ impl Response {
             Default::default(),
             status_code,
             context,
+            Default::default(),
         ))
     }
 }

apollo-router/src/services/supergraph.rs

@@ -131,25 +131,29 @@ impl Request {
     /// Create a request with an example query, for tests
     #[builder(visibility = "pub")]
     fn canned_new(
+        query: Option<String>,
         operation_name: Option<String>,
         // Skip the `Object` type alias in order to use buildstructor’s map special-casing
         extensions: JsonMap<ByteString, Value>,
         context: Option<Context>,
         headers: MultiMap<TryIntoHeaderName, TryIntoHeaderValue>,
     ) -> Result<Request, BoxError> {
-        let query = "
-            query TopProducts($first: Int) { 
-                topProducts(first: $first) { 
-                    upc 
-                    name 
-                    reviews { 
-                        id 
-                        product { name } 
-                        author { id name } 
-                    } 
-                } 
+        let query = query.unwrap_or(
+            "
+            query TopProducts($first: Int) {
+                topProducts(first: $first) {
+                    upc
+                    name
+                    reviews {
+                        id
+                        product { name }
+                        author { id name }
+                    }
+                }
             }
-        ";
+        "
+            .to_string(),
+        );
         let mut variables = JsonMap::new();
         variables.insert("first", 2_usize.into());
         Self::fake_new(

apollo-router/src/test_harness.rs

@@ -214,6 +214,7 @@ impl<'a> TestHarness<'a> {
                     let empty_response = subgraph::Response::builder()
                         .extensions(crate::json_ext::Object::new())
                         .context(request.context)
+                        .headers(http::HeaderMap::new())
                         .build();
                     std::future::ready(Ok(empty_response))
                 })

examples/cache-control/rhai/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "cache-control"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+[dependencies]
+anyhow = "1"
+apollo-router = { path = "../../../apollo-router" }
+http = "0.2"
+serde_json = "1"
+tokio = { version = "1", features = ["full"] }
+tower = { version = "0.4", features = ["full"] }

examples/cache-control/rhai/README.md

@@ -0,0 +1,9 @@
+# Rhai script
+
+Demonstrates header and context manipulation via Rhai script.
+
+Usage:
+
+```bash
+cargo run -- -s ../../graphql/supergraph.graphql -c ./router.yaml
+```

examples/cache-control/rhai/router.yaml

@@ -0,0 +1,3 @@
+rhai:
+  scripts: src
+  main: cache_control.rhai

examples/cache-control/rhai/src/cache_control.rhai

@@ -0,0 +1,80 @@
+fn subgraph_service(service, subgraph) {
+  // collect the max-age and scope values from cache-control headers and store
+  // on the context for use in supergraph_service
+  service.map_response(|response| {
+    let cache_control = response.headers.values("cache-control").get(0);
+
+    // if a subgraph response is uncacheable, the whole response is uncacheable
+    if cache_control == () {
+      response.context.cache_control_uncacheable = true;
+      return;
+    }
+
+    let max_age = get_max_age(cache_control);
+
+    // use the smallest max age
+    response.context.upsert("cache_control_max_age", |current| {
+      if current == () {
+        max_age
+      } else if max_age < current {
+        max_age
+      } else {
+        current
+      }
+    });
+
+    let scope = if cache_control.contains("public") {
+      "public"
+    } else {
+      "private"
+    };
+
+    // if the scope is ever private, it cannot become public
+    response.context.upsert("cache_control_scope", |current| {
+      if current == "private" || scope == "private" {
+        "private"
+      } else {
+        scope
+      }
+    });
+  });
+}
+
+fn supergraph_service(service) {
+  // attach the cache-control header if enough data is available
+  service.map_response(|response| {
+    let uncacheable = response.context.cache_control_uncacheable;
+    let max_age = response.context.cache_control_max_age;
+    let scope = response.context.cache_control_scope;
+
+    if uncacheable != true && max_age != () && scope != () {
+      response.headers["cache-control"] = `max-age=${max_age}, ${scope}`;
+    }
+  });
+}
+
+// find the the max-age= part and parse the value into an integer
+fn get_max_age(str) {
+  let max_age = 0;
+
+  for part in str.split(",") {
+    part.remove(" ");
+
+    if part.starts_with("max-age=") {
+      let num = part.split("=").get(1);
+
+      if num == () || num == "" {
+        break;
+      }
+
+      try {
+        max_age = num.parse_int();
+      } catch (err) {
+        log_error(`error parsing max-age from "${str}": ${err}`);
+      }
+      break;
+    }
+  }
+
+  max_age
+}