Skip to content

Commit

Permalink
Feature/cred custom headers (#68)
Browse files Browse the repository at this point in the history 
* wip

* support additional headers in upload
  • Loading branch information
@dbernaciak
dbernaciak authored Sep 13, 2024
1 parent 43b20f9 commit afc5e4d
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 4 deletions.
8 changes: 8 additions & 0 deletions py_integ/test_integ.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,5 +156,13 @@ def test_generic_dl(client: Fusion) -> None:
print(f"Passed integ tests for {test_nm}") # noqa: T201


def test_custom_headers(client: Fusion) -> None:
client.credentials.headers = {"my_special_header": "my_special_value", "my_special_header2": "my_special_value2"}
r = client.session.get(client.root_url + "catalogs/common")
assert r.request.headers["my_special_header"] == "my_special_value"
assert r.request.headers["my_special_header2"] == "my_special_value2"
print("Passed custom headers test") # noqa: T201


if __name__ == "__main__":
gen_generic_dl()
42 changes: 38 additions & 4 deletions src/auth.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,7 @@ struct FusionCredsPersistent {
grant_type: String,
//#[serde(deserialize_with = "deserialize_fusion_e2e")]
fusion_e2e: Option<String>,
headers: Option<HashMap<String, String>>,
}

#[pyclass(module = "fusion._fusion")]
Expand Down Expand Up @@ -315,6 +316,9 @@ pub struct FusionCredentials {
#[pyo3(get, set)]
fusion_e2e: Option<String>,

#[pyo3(get, set)]
headers: HashMap<String, String>,

#[pyo3(get)]
proxies: HashMap<String, String>,

Expand Down Expand Up @@ -342,6 +346,7 @@ impl Default for FusionCredentials {
proxies: HashMap::new(),
grant_type: "client_credentials".to_string(),
fusion_e2e: None,
headers: HashMap::new(),
http_client: None,
}
}
Expand Down Expand Up @@ -372,6 +377,7 @@ impl FusionCredentials {
Option<HashMap<String, String>>,
Option<String>,
Option<String>,
Option<HashMap<String, String>>,
)> {
Ok((
self.client_id.clone(),
Expand All @@ -384,11 +390,12 @@ impl FusionCredentials {
Some(self.proxies.clone()),
Some(self.grant_type.clone()),
self.fusion_e2e.clone(),
Some(self.headers.clone()),
))
}

#[classmethod]
#[pyo3(signature = (client_id=None, client_secret=None, resource=None, auth_url=None, proxies=None, fusion_e2e=None))]
#[pyo3(signature = (client_id=None, client_secret=None, resource=None, auth_url=None, proxies=None, fusion_e2e=None, headers=None))]
fn from_client_id(
_cls: &Bound<'_, PyType>,
client_id: Option<String>,
Expand All @@ -397,6 +404,7 @@ impl FusionCredentials {
auth_url: Option<String>,
proxies: Option<HashMap<String, String>>,
fusion_e2e: Option<String>,
headers: Option<HashMap<String, String>>,
) -> PyResult<Self> {
Ok(Self {
client_id,
Expand All @@ -406,6 +414,7 @@ impl FusionCredentials {
proxies: proxies.unwrap_or_default(),
grant_type: "client_credentials".to_string(),
fusion_e2e,
headers: headers.unwrap_or_default(),
fusion_token: HashMap::new(),
bearer_token: None,
username: None,
Expand All @@ -423,7 +432,7 @@ impl FusionCredentials {
}

#[classmethod]
#[pyo3(signature = (client_id=None, username=None, password=None, resource=None, auth_url=None, proxies=None, fusion_e2e=None))]
#[pyo3(signature = (client_id=None, username=None, password=None, resource=None, auth_url=None, proxies=None, fusion_e2e=None, headers=None))]
fn from_user_id(
_cls: &Bound<'_, PyType>,
client_id: Option<String>,
Expand All @@ -433,6 +442,7 @@ impl FusionCredentials {
auth_url: Option<String>,
proxies: Option<HashMap<String, String>>,
fusion_e2e: Option<String>,
headers: Option<HashMap<String, String>>,
) -> PyResult<Self> {
Ok(Self {
client_id,
Expand All @@ -443,6 +453,7 @@ impl FusionCredentials {
proxies: proxies.unwrap_or_default(),
grant_type: "password".to_string(),
fusion_e2e,
headers: headers.unwrap_or_default(),
fusion_token: HashMap::new(),
bearer_token: None,
client_secret: None,
Expand All @@ -451,20 +462,22 @@ impl FusionCredentials {
}

#[classmethod]
#[pyo3(signature = (bearer_token=None, bearer_token_expiry=None, proxies=None, fusion_e2e=None))]
#[pyo3(signature = (bearer_token=None, bearer_token_expiry=None, proxies=None, fusion_e2e=None, headers=None))]
fn from_bearer_token(
_cls: &Bound<'_, PyType>,
bearer_token: Option<String>,
bearer_token_expiry: Option<&Bound<PyDate>>,
proxies: Option<HashMap<String, String>>,
fusion_e2e: Option<String>,
headers: Option<HashMap<String, String>>,
) -> PyResult<Self> {
Ok(Self {
resource: None,
auth_url: None,
proxies: proxies.unwrap_or_default(),
grant_type: "bearer".to_string(),
fusion_e2e,
headers: headers.unwrap_or_default(),
fusion_token: HashMap::new(),
bearer_token: Some(AuthToken {
token: match bearer_token {
Expand All @@ -491,7 +504,7 @@ impl FusionCredentials {

#[allow(clippy::too_many_arguments)]
#[new]
#[pyo3(signature = (client_id=None, client_secret=None, username=None, password=None, resource=None, auth_url=None, bearer_token=None, proxies=None, grant_type=None, fusion_e2e=None))]
#[pyo3(signature = (client_id=None, client_secret=None, username=None, password=None, resource=None, auth_url=None, bearer_token=None, proxies=None, grant_type=None, fusion_e2e=None, headers=None))]
fn new(
client_id: Option<String>,
client_secret: Option<String>,
Expand All @@ -503,6 +516,7 @@ impl FusionCredentials {
proxies: Option<HashMap<String, String>>,
grant_type: Option<String>,
fusion_e2e: Option<String>,
headers: Option<HashMap<String, String>>,
) -> PyResult<Self> {
Ok(FusionCredentials {
client_id,
Expand All @@ -516,6 +530,7 @@ impl FusionCredentials {
proxies: proxies.unwrap_or_default(),
grant_type: grant_type.unwrap_or_else(|| "client_credentials".to_string()),
fusion_e2e,
headers: headers.unwrap_or_default(),
http_client: None,
})
}
Expand Down Expand Up @@ -696,6 +711,12 @@ impl FusionCredentials {
self.fusion_e2e.as_ref().unwrap().clone(),
);
}
// if headers are not empty add each key value pair to the headers
if !self.headers.is_empty() {
for (key, value) in self.headers.iter() {
ret.insert(key.clone(), value.clone());
}
}

self._refresh_bearer_token(py, false, 30)?;
let bearer_token_tup = self
Expand Down Expand Up @@ -798,13 +819,15 @@ impl FusionCredentials {
credentials.auth_url,
Some(untyped_proxies(credentials.proxies)),
credentials.fusion_e2e,
credentials.headers,
)?,
"bearer" => FusionCredentials::from_bearer_token(
cls,
None,
None,
Some(untyped_proxies(credentials.proxies)),
credentials.fusion_e2e,
credentials.headers,
)?,
"password" => FusionCredentials::from_user_id(
cls,
Expand All @@ -815,6 +838,7 @@ impl FusionCredentials {
credentials.auth_url,
Some(untyped_proxies(credentials.proxies)),
credentials.fusion_e2e,
credentials.headers,
)?,
_ => {
return Err(pyo3::exceptions::PyValueError::new_err(
Expand Down Expand Up @@ -1025,6 +1049,7 @@ mod tests {
Some(HashMap::new()),
Some("grant_type".to_string()),
Some("fusion_e2e".to_string()),
Some(HashMap::new()),
)
.unwrap();

Expand All @@ -1051,6 +1076,7 @@ mod tests {
assert!(creds.fusion_e2e.is_none());
assert!(creds.bearer_token.is_none());
assert!(creds.proxies.is_empty());
assert!(creds.headers.is_empty());
assert!(creds.fusion_token.is_empty());
}

Expand All @@ -1067,6 +1093,7 @@ mod tests {
Some(HashMap::new()),
Some("grant_type".to_string()),
Some("fusion_e2e".to_string()),
Some(HashMap::new()),
)
.unwrap();

Expand Down Expand Up @@ -1096,6 +1123,7 @@ mod tests {
Some(HashMap::new()),
Some("grant_type".to_string()),
Some("fusion_e2e".to_string()),
Some(HashMap::new()),
)
.unwrap();

Expand Down Expand Up @@ -1126,6 +1154,7 @@ mod tests {
Some(HashMap::new()),
Some("grant_type".to_string()),
Some("fusion_e2e".to_string()),
Some(HashMap::new()),
)
.unwrap();

Expand All @@ -1140,6 +1169,7 @@ mod tests {
proxies,
grant_type,
fusion_e2e,
headers,
) = creds.__getnewargs__().unwrap();

assert_eq!(client_id, Some("client_id".to_string()));
Expand All @@ -1150,6 +1180,7 @@ mod tests {
assert_eq!(auth_url, Some("auth_url".to_string()));
assert!(bearer_token.is_none());
assert!(proxies.is_some());
assert!(headers.is_some());
assert_eq!(grant_type, Some("grant_type".to_string()));
assert_eq!(fusion_e2e, Some("fusion_e2e".to_string()));
}
Expand All @@ -1166,6 +1197,7 @@ mod tests {
Some("auth_url".to_string()),
None,
None,
None,
)
.unwrap();

Expand All @@ -1190,6 +1222,7 @@ mod tests {
Some("auth_url".to_string()),
None,
None,
None,
)
.unwrap();

Expand All @@ -1212,6 +1245,7 @@ mod tests {
Some(&expiry_date),
None,
None,
None,
)
.unwrap();

Expand Down

0 comments on commit afc5e4d

Please sign in to comment.