use prost_reflect::{MethodDescriptor, ServiceDescriptor};
use std::collections::HashMap;
use tonic::Status;
#[derive(Debug, Clone)]
pub struct ServiceDescriptorCache {
services: HashMap<String, ServiceDescriptor>,
methods: HashMap<(String, String), MethodDescriptor>,
}
impl Default for ServiceDescriptorCache {
fn default() -> Self {
Self::new()
}
}
impl ServiceDescriptorCache {
pub fn new() -> Self {
Self {
services: HashMap::new(),
methods: HashMap::new(),
}
}
pub fn add_service(&mut self, service: ServiceDescriptor) {
let service_name = service.full_name().to_string();
self.services.insert(service_name.clone(), service.clone());
for method in service.methods() {
let method_name = method.name().to_string();
self.methods.insert((service_name.clone(), method_name), method);
}
}
pub fn get_service(&self, service_name: &str) -> Option<&ServiceDescriptor> {
self.services.get(service_name)
}
pub fn get_service_with_error(&self, service_name: &str) -> Result<&ServiceDescriptor, Status> {
self.services.get(service_name).ok_or_else(|| {
Status::not_found(format!("Service '{}' not found in descriptor cache", service_name))
})
}
pub fn get_method(
&self,
service_name: &str,
method_name: &str,
) -> Result<&MethodDescriptor, Status> {
self.methods
.get(&(service_name.to_string(), method_name.to_string()))
.ok_or_else(|| {
Status::not_found(format!(
"Method {} not found in service {}",
method_name, service_name
))
})
}
pub fn contains_service(&self, service_name: &str) -> bool {
self.services.contains_key(service_name)
}
pub fn contains_method(&self, service_name: &str, method_name: &str) -> bool {
self.methods.contains_key(&(service_name.to_string(), method_name.to_string()))
}
pub fn service_count(&self) -> usize {
self.services.len()
}
pub fn method_count(&self) -> usize {
self.methods.len()
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_module_compiles() {
}
}