templates/load_data: add an optional parameter headers (#1710)

* templates/load_data: add an optional parameter headers ...

... now `load_data` function supports setting extra headers

* docs/templates/overview: cover some edge-cases in the explanation

* templates/load_data: fix caching logic with headers

* docs/templates: change wording for load_data headers explanations
This commit is contained in:
liushuyu 2022-01-11 13:36:21 -07:00 committed by Vincent Prouillet
parent 78132a8d1d
commit 37b31cb61f
2 changed files with 191 additions and 3 deletions

View file

@ -6,7 +6,7 @@ use std::str::FromStr;
use std::sync::{Arc, Mutex};
use csv::Reader;
use reqwest::header::{HeaderValue, CONTENT_TYPE};
use reqwest::header::{HeaderMap, HeaderName, HeaderValue, CONTENT_TYPE};
use reqwest::{blocking::Client, header};
use tera::{from_value, to_value, Error, Function as TeraFn, Map, Result, Value};
use url::Url;
@ -120,12 +120,14 @@ impl DataSource {
method: Method,
post_body: &Option<String>,
post_content_type: &Option<String>,
headers: &Option<Vec<String>>,
) -> u64 {
let mut hasher = DefaultHasher::new();
format.hash(&mut hasher);
method.hash(&mut hasher);
post_body.hash(&mut hasher);
post_content_type.hash(&mut hasher);
headers.hash(&mut hasher);
self.hash(&mut hasher);
hasher.finish()
}
@ -162,6 +164,31 @@ fn get_output_format_from_args(
}
}
fn add_headers_from_args(header_args: Option<Vec<String>>) -> Result<HeaderMap> {
let mut headers = HeaderMap::new();
if let Some(header_args) = header_args {
for arg in header_args {
let mut splitter = arg.splitn(2, '=');
let key = splitter
.next()
.ok_or_else(|| {
format!("Invalid header argument. Expecting header key, got '{}'", arg)
})?
.to_string();
let value = splitter.next().ok_or_else(|| {
format!("Invalid header argument. Expecting header value, got '{}'", arg)
})?;
headers.append(
HeaderName::from_str(&key)
.map_err(|e| format!("Invalid header name '{}': {}", key, e))?,
value.parse().map_err(|e| format!("Invalid header value '{}': {}", value, e))?,
);
}
}
Ok(headers)
}
/// A Tera function to load data from a file or from a URL
/// Currently the supported formats are json, toml, csv, bibtex and plain text
#[derive(Debug)]
@ -223,6 +250,11 @@ impl TeraFn for LoadData {
},
_ => Method::Get,
};
let headers = optional_arg!(
Vec<String>,
args.get("headers"),
"`load_data`: `headers` needs to be an argument with a list of strings of format <name>=<value>."
);
// If the file doesn't exist, source is None
let data_source = match (
@ -255,8 +287,13 @@ impl TeraFn for LoadData {
};
let file_format = get_output_format_from_args(format_arg, &data_source)?;
let cache_key =
data_source.get_cache_key(&file_format, method, &post_body_arg, &post_content_type);
let cache_key = data_source.get_cache_key(
&file_format,
method,
&post_body_arg,
&post_content_type,
&headers,
);
let mut cache = self.result_cache.lock().expect("result cache lock");
if let Some(cached_result) = cache.get(&cache_key) {
@ -271,10 +308,12 @@ impl TeraFn for LoadData {
let req = match method {
Method::Get => response_client
.get(url.as_str())
.headers(add_headers_from_args(headers)?)
.header(header::ACCEPT, file_format.as_accept_header()),
Method::Post => {
let mut resp = response_client
.post(url.as_str())
.headers(add_headers_from_args(headers)?)
.header(header::ACCEPT, file_format.as_accept_header());
if let Some(content_type) = post_content_type {
match HeaderValue::from_str(&content_type) {
@ -660,12 +699,14 @@ mod tests {
Method::Get,
&None,
&None,
&Some(vec![]),
);
let cache_key_2 = DataSource::Path(get_test_file("test.toml")).get_cache_key(
&OutputFormat::Toml,
Method::Get,
&None,
&None,
&Some(vec![]),
);
assert_eq!(cache_key, cache_key_2);
}
@ -677,12 +718,14 @@ mod tests {
Method::Get,
&None,
&None,
&Some(vec![]),
);
let json_cache_key = DataSource::Path(get_test_file("test.json")).get_cache_key(
&OutputFormat::Toml,
Method::Get,
&None,
&None,
&Some(vec![]),
);
assert_ne!(toml_cache_key, json_cache_key);
}
@ -694,16 +737,37 @@ mod tests {
Method::Get,
&None,
&None,
&Some(vec![]),
);
let json_cache_key = DataSource::Path(get_test_file("test.toml")).get_cache_key(
&OutputFormat::Json,
Method::Get,
&None,
&None,
&Some(vec![]),
);
assert_ne!(toml_cache_key, json_cache_key);
}
#[test]
fn different_cache_key_per_headers() {
let header1_cache_key = DataSource::Path(get_test_file("test.toml")).get_cache_key(
&OutputFormat::Json,
Method::Get,
&None,
&None,
&Some(vec!["a=b".to_string()]),
);
let header2_cache_key = DataSource::Path(get_test_file("test.toml")).get_cache_key(
&OutputFormat::Json,
Method::Get,
&None,
&None,
&Some(vec![]),
);
assert_ne!(header1_cache_key, header2_cache_key);
}
#[test]
fn can_load_remote_data() {
let _m = mock("GET", "/zpydpkjj67")
@ -1002,4 +1066,91 @@ mod tests {
_mjson.assert();
}
#[test]
fn is_custom_headers_working() {
let _mjson = mock("POST", "/kr1zdgbm4y4")
.with_header("content-type", "application/json")
.match_header("accept", "text/plain")
.match_header("x-custom-header", "some-values")
.with_body("{i_am:'json'}")
.expect(1)
.create();
let url = format!("{}{}", mockito::server_url(), "/kr1zdgbm4y4");
let static_fn = LoadData::new(PathBuf::from("../utils"), None, PathBuf::new());
let mut args = HashMap::new();
args.insert("url".to_string(), to_value(&url).unwrap());
args.insert("format".to_string(), to_value("plain").unwrap());
args.insert("method".to_string(), to_value("post").unwrap());
args.insert("content_type".to_string(), to_value("text/plain").unwrap());
args.insert("body".to_string(), to_value("this is a match").unwrap());
args.insert("headers".to_string(), to_value(["x-custom-header=some-values"]).unwrap());
let result = static_fn.call(&args);
assert!(result.is_ok());
_mjson.assert();
}
#[test]
fn is_custom_headers_working_with_multiple_values() {
let _mjson = mock("POST", "/kr1zdgbm4y5")
.with_status(201)
.with_header("content-type", "application/json")
.match_header("authorization", "Bearer 123")
// Mockito currently does not have a way to validate multiple headers with the same name
// see https://github.com/lipanski/mockito/issues/117
.match_header("accept", mockito::Matcher::Any)
.match_header("x-custom-header", "some-values")
.match_header("x-other-header", "some-other-values")
.with_body("<html>I am a server that needs authentication and returns HTML with Accept set to JSON</html>")
.expect(1)
.create();
let url = format!("{}{}", mockito::server_url(), "/kr1zdgbm4y5");
let static_fn = LoadData::new(PathBuf::from("../utils"), None, PathBuf::new());
let mut args = HashMap::new();
args.insert("url".to_string(), to_value(&url).unwrap());
args.insert("format".to_string(), to_value("plain").unwrap());
args.insert("method".to_string(), to_value("post").unwrap());
args.insert("content_type".to_string(), to_value("text/plain").unwrap());
args.insert("body".to_string(), to_value("this is a match").unwrap());
args.insert(
"headers".to_string(),
to_value([
"x-custom-header=some-values",
"x-other-header=some-other-values",
"accept=application/json",
"authorization=Bearer 123",
])
.unwrap(),
);
let result = static_fn.call(&args);
assert!(result.is_ok());
_mjson.assert();
}
#[test]
fn fails_when_specifying_invalid_headers() {
let _mjson = mock("GET", "/kr1zdgbm4y6").with_status(204).expect(0).create();
let static_fn = LoadData::new(PathBuf::from("../utils"), None, PathBuf::new());
let url = format!("{}{}", mockito::server_url(), "/kr1zdgbm4y6");
let mut args = HashMap::new();
args.insert("url".to_string(), to_value(&url).unwrap());
args.insert("format".to_string(), to_value("plain").unwrap());
args.insert("headers".to_string(), to_value(["bad-entry::bad-header"]).unwrap());
let result = static_fn.call(&args);
assert!(result.is_err());
let static_fn = LoadData::new(PathBuf::from("../utils"), None, PathBuf::new());
let mut args = HashMap::new();
args.insert("url".to_string(), to_value(&url).unwrap());
args.insert("format".to_string(), to_value("plain").unwrap());
args.insert("headers".to_string(), to_value(["\n=\r"]).unwrap());
let result = static_fn.call(&args);
assert!(result.is_err());
_mjson.assert();
}
}

View file

@ -409,6 +409,43 @@ This example will make a POST request to the kroki service to generate a SVG.
{{postdata|safe}}
```
If you need additional handling for the HTTP headers, you can use the `headers` parameter.
You might need this parameter when the resource requires authentication or require passing additional
parameters via special headers.
Please note that the headers will be appended to the default headers set by Zola itself instead of replacing them.
This example will make a POST request to the GitHub markdown rendering service.
```jinja2
{% set postdata = load_data(url="https://api.github.com/markdown", format="plain", method="POST", content_type="application/json", headers=["accept=application/vnd.github.v3+json"], body='{"text":"headers support added in #1710, commit before it: b3918f124d13ec1bedad4860c15a060dd3751368","context":"getzola/zola","mode":"gfm"}')%}
{{postdata|safe}}
```
The following example shows how to send a GraphQL query to GitHub (requires authentication).
If you want to try this example on your own machine, you need to provide a GitHub PAT (Personal Access Token),
you can acquire the access token at this link: https://github.com/settings/tokens and then set `GITHUB_TOKEN`
environment variable to the access token you have obtained.
```jinja2
{% set token = get_env(name="GITHUB_TOKEN") %}
{% set postdata = load_data(url="https://api.github.com/graphql", format="json", method="POST" ,content_type="application/json", headers=["accept=application/vnd.github.v4.idl", "authentication=Bearer " ~ token], body='{"query":"query { viewer { login }}"}')%}
{{postdata|safe}}
```
In case you need to specify multiple headers with the same name, you can specify them like this:
```
headers=["accept=application/json,text/html"]
```
Which is equivalent to two `Accept` headers with `application/json` and `text/html`.
If it doesn't work, you can instead specify the headers multiple times to achieve a similar effect:
```
headers=["accept=application/json", "accept=text/html"]
```
#### Data caching
Data file loading and remote requests are cached in memory during the build, so multiple requests aren't made