use crate::segment::{AwsOperation, Http, Request, Response, Subsegment};
pub trait Namespace {
fn name(&self, prefix: &str) -> String;
fn update_subsegment(&self, subsegment: &mut Subsegment);
}
#[derive(Debug)]
pub struct AwsNamespace {
service: String,
operation: String,
request_id: Option<String>,
response_status: Option<u16>,
}
impl AwsNamespace {
pub fn new(service: impl Into<String>, operation: impl Into<String>) -> Self {
Self {
service: service.into(),
operation: operation.into(),
request_id: None,
response_status: None,
}
}
pub fn request_id(&mut self, request_id: impl Into<String>) -> &mut Self {
self.request_id = Some(request_id.into());
self
}
pub fn response_status(&mut self, status: u16) -> &mut Self {
self.response_status = Some(status);
self
}
}
impl Namespace for AwsNamespace {
fn name(&self, _prefix: &str) -> String {
self.service.clone()
}
fn update_subsegment(&self, subsegment: &mut Subsegment) {
if subsegment.namespace.is_none() {
subsegment.namespace = Some("aws".to_string());
}
if let Some(aws) = subsegment.aws.as_mut() {
if aws.operation.is_none() {
aws.operation = Some(self.operation.clone());
}
if aws.request_id.is_none() {
aws.request_id = self.request_id.clone();
}
} else {
subsegment.aws = Some(AwsOperation {
operation: Some(self.operation.clone()),
request_id: self.request_id.clone(),
..AwsOperation::default()
});
}
if let Some(response_status) = self.response_status {
if let Some(http) = subsegment.http.as_mut() {
if let Some(response) = http.response.as_mut() {
if response.status.is_none() {
response.status = Some(response_status);
}
} else {
http.response = Some(Response {
status: Some(response_status),
..Response::default()
});
}
} else {
subsegment.http = Some(Http {
response: Some(Response {
status: Some(response_status),
..Response::default()
}),
..Http::default()
});
}
}
}
}
#[derive(Debug)]
pub struct RemoteNamespace {
name: String,
method: String,
url: String,
response_status: Option<u16>,
}
impl RemoteNamespace {
pub fn new(name: impl Into<String>, method: impl Into<String>, url: impl Into<String>) -> Self {
Self {
name: name.into(),
method: method.into(),
url: url.into(),
response_status: None,
}
}
pub fn response_status(&mut self, status: u16) -> &mut Self {
self.response_status = Some(status);
self
}
}
impl Namespace for RemoteNamespace {
fn name(&self, _prefix: &str) -> String {
self.name.clone()
}
fn update_subsegment(&self, subsegment: &mut Subsegment) {
if subsegment.namespace.is_none() {
subsegment.namespace = Some("remote".to_string());
}
if let Some(http) = subsegment.http.as_mut() {
if let Some(request) = http.request.as_mut() {
if request.method.is_none() {
request.method = Some(self.method.clone());
}
if request.url.is_none() {
request.url = Some(self.url.clone());
}
} else {
http.request = Some(Request {
url: Some(self.url.clone()),
method: Some(self.method.clone()),
..Request::default()
});
}
} else {
subsegment.http = Some(Http {
request: Some(Request {
url: Some(self.url.clone()),
method: Some(self.method.clone()),
..Request::default()
}),
..Http::default()
});
}
if let Some(response_status) = self.response_status {
let http = subsegment.http.as_mut().expect("http must have been set");
if let Some(response) = http.response.as_mut() {
if response.status.is_none() {
response.status = Some(response_status);
}
} else {
http.response = Some(Response {
status: Some(response_status),
..Response::default()
});
}
}
}
}
#[derive(Debug)]
pub struct CustomNamespace {
name: String,
}
impl CustomNamespace {
pub fn new(name: impl Into<String>) -> Self {
Self { name: name.into() }
}
}
impl Namespace for CustomNamespace {
fn name(&self, prefix: &str) -> String {
format!("{}{}", prefix, self.name)
}
fn update_subsegment(&self, _subsegment: &mut Subsegment) {}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn aws_namespace_should_have_service_name_as_name() {
let namespace = AwsNamespace::new("S3", "GetObject");
assert_eq!(namespace.name(""), "S3");
assert_eq!(namespace.name("prefix"), "S3");
}
#[test]
fn aws_namespace_should_update_subsegment_with_aws_operation() {
let namespace = AwsNamespace::new("S3", "GetObject");
let mut subsegment = Subsegment::default();
namespace.update_subsegment(&mut subsegment);
assert_eq!(subsegment.namespace.unwrap(), "aws");
assert_eq!(subsegment.aws.unwrap().operation.unwrap(), "GetObject");
}
#[test]
fn aws_namespace_should_update_subsegment_with_request_id() {
let mut namespace = AwsNamespace::new("S3", "GetObject");
namespace.request_id("12345");
let mut subsegment = Subsegment::default();
namespace.update_subsegment(&mut subsegment);
assert_eq!(subsegment.aws.unwrap().request_id.unwrap(), "12345");
}
#[test]
fn aws_namespace_should_update_subsegment_with_response_status() {
let mut namespace = AwsNamespace::new("S3", "GetObject");
namespace.response_status(200);
let mut subsegment = Subsegment::default();
namespace.update_subsegment(&mut subsegment);
assert_eq!(
subsegment
.http
.expect("http")
.response
.expect("response")
.status
.expect("status"),
200,
);
}
#[test]
fn remote_namespace_should_have_name_as_name() {
let namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
assert_eq!(namespace.name(""), "codemonger.io");
assert_eq!(namespace.name("prefix"), "codemonger.io");
}
#[test]
fn remote_namespace_should_update_subsegment_with_remote_service() {
let namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
let mut subsegment = Subsegment::default();
namespace.update_subsegment(&mut subsegment);
assert_eq!(subsegment.namespace.unwrap(), "remote");
let request = subsegment.http.expect("http").request.expect("request");
assert_eq!(request.method.unwrap(), "GET");
assert_eq!(request.url.unwrap(), "https://codemonger.io/");
}
#[test]
fn remote_namespace_should_update_subsegment_with_response_status() {
let mut namespace = RemoteNamespace::new("codemonger.io", "GET", "https://codemonger.io/");
namespace.response_status(200);
let mut subsegment = Subsegment::default();
namespace.update_subsegment(&mut subsegment);
assert_eq!(
subsegment
.http
.expect("http")
.response
.expect("response")
.status
.expect("status"),
200,
);
}
#[test]
fn custom_namespace_should_have_prefixed_name() {
let namespace = CustomNamespace::new("TestSubsegment");
assert_eq!(namespace.name(""), "TestSubsegment");
assert_eq!(namespace.name("prefix"), "prefixTestSubsegment");
}
}