use std::collections::BTreeMap;
use std::future::Future;
use std::pin::Pin;
use std::time::Duration;
use crate::cmd_path::StandardCommand;
use crate::descriptor::{CommandDescriptor, FlagDescriptor};
use crate::model::error::{ErrorCode, ObzError};
use crate::provider::traits::{ExtensionProvider, LogProvider, MetricProvider, TraceProvider};
use crate::provider::ProviderConfig;
pub struct BuiltProvider {
pub name: &'static str,
pub metric_query_language: Option<&'static str>,
pub log_query_language: Option<&'static str>,
pub metric: Option<Box<dyn MetricProvider>>,
pub log: Option<Box<dyn LogProvider>>,
pub trace: Option<Box<dyn TraceProvider>>,
pub extension: Option<Box<dyn ExtensionProvider>>,
}
pub type ProviderFactory = fn(&ProviderConfig) -> Result<BuiltProvider, ObzError>;
#[derive(Debug, Clone)]
pub struct CheckResult {
pub severity: CheckSeverity,
pub message: String,
pub scope: CheckScope,
pub latency: Option<Duration>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CheckSeverity {
Ok,
Warn,
Fail,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CheckScope {
Connectivity,
ConnectivityAndAuth,
ConfiguredNotVerifiable,
}
pub type ProviderCheckFn =
fn(&ProviderConfig) -> Pin<Box<dyn Future<Output = CheckResult> + Send + '_>>;
#[derive(Debug, Clone, Copy, Default)]
pub struct SupportedCommands {
pub metric_query: bool,
pub metric_list: bool,
pub metric_info: bool,
pub metric_labels: bool,
pub metric_label_values: bool,
pub metric_series: bool,
pub log_search: bool,
pub trace_search: bool,
pub trace_get: bool,
}
pub struct ProviderMeta {
pub name: &'static str,
pub display_name: &'static str,
pub aliases: &'static [&'static str],
pub supported_commands: SupportedCommands,
pub build: ProviderFactory,
pub check: Option<ProviderCheckFn>,
pub command_flags: &'static [(StandardCommand, &'static [FlagDescriptor])],
pub extension_commands: &'static [(&'static str, CommandDescriptor)],
}
#[derive(Default)]
pub struct ProviderRegistry {
by_alias: BTreeMap<String, usize>,
metas: Vec<ProviderMeta>,
}
impl ProviderRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register(&mut self, meta: ProviderMeta) {
let idx = self.metas.len();
for &alias in meta.aliases {
assert!(
self.by_alias.insert(alias.to_string(), idx).is_none(),
"provider alias '{}' is already registered",
alias
);
}
self.metas.push(meta);
}
pub fn get(&self, name: &str) -> Result<&ProviderMeta, ObzError> {
self.by_alias
.get(name)
.map(|&i| &self.metas[i])
.ok_or_else(|| ObzError::InvalidArgument {
code: ErrorCode::InvalidFlag,
message: format!(
"unknown provider '{}' — run `obz provider list` to see available providers",
name
),
suggestion: None,
})
}
pub fn all(&self) -> &[ProviderMeta] {
&self.metas
}
pub fn all_aliases(&self) -> Vec<&str> {
self.by_alias.keys().map(String::as_str).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn dummy_meta(name: &'static str, aliases: &'static [&'static str]) -> ProviderMeta {
fn dummy_build(_: &ProviderConfig) -> Result<BuiltProvider, ObzError> {
Ok(BuiltProvider {
name: "dummy",
metric_query_language: None,
log_query_language: None,
metric: None,
log: None,
trace: None,
extension: None,
})
}
ProviderMeta {
name,
display_name: name,
aliases,
supported_commands: SupportedCommands::default(),
build: dummy_build,
check: None,
command_flags: &[],
extension_commands: &[],
}
}
#[test]
fn test_get_resolves_provider_by_name_and_alias() {
let mut registry = ProviderRegistry::new();
registry.register(dummy_meta("victoriametrics", &["vm", "victoriametrics"]));
let meta = registry.get("vm").unwrap();
assert_eq!(meta.name, "victoriametrics");
let meta2 = registry.get("victoriametrics").unwrap();
assert_eq!(meta2.name, "victoriametrics");
}
#[test]
fn test_get_returns_error_for_unknown_provider() {
let registry = ProviderRegistry::new();
match registry.get("nonexistent") {
Err(ObzError::InvalidArgument { code, message, .. }) => {
assert_eq!(code, ErrorCode::InvalidFlag);
assert!(message.contains("nonexistent"));
}
Err(other) => panic!("expected InvalidArgument, got {other:?}"),
Ok(_) => panic!("expected error, got Ok"),
}
}
#[test]
fn test_all_preserves_insertion_order() {
let mut registry = ProviderRegistry::new();
registry.register(dummy_meta("alpha", &["alpha"]));
registry.register(dummy_meta("beta", &["beta"]));
registry.register(dummy_meta("gamma", &["gamma"]));
let names: Vec<&str> = registry.all().iter().map(|m| m.name).collect();
assert_eq!(names, vec!["alpha", "beta", "gamma"]);
}
#[test]
fn test_all_aliases_returns_sorted_list() {
let mut registry = ProviderRegistry::new();
registry.register(dummy_meta("victoriametrics", &["vm", "victoriametrics"]));
registry.register(dummy_meta("victorialogs", &["vl", "victorialogs"]));
let aliases = registry.all_aliases();
assert_eq!(aliases, vec!["victorialogs", "victoriametrics", "vl", "vm"]);
}
#[test]
#[should_panic(expected = "provider alias 'vm' is already registered")]
fn test_register_panics_on_duplicate_alias() {
let mut registry = ProviderRegistry::new();
registry.register(dummy_meta("provider1", &["vm"]));
registry.register(dummy_meta("provider2", &["vm"])); }
#[test]
fn test_new_registry_is_empty() {
let registry = ProviderRegistry::new();
assert!(registry.all().is_empty());
assert!(registry.all_aliases().is_empty());
}
#[test]
fn test_registry_with_extension_commands() {
static EXT_CMDS: &[(&str, CommandDescriptor)] = &[
(
"trace",
CommandDescriptor {
name: "services",
description: "List services",
flags: &[],
},
),
(
"metric",
CommandDescriptor {
name: "top-queries",
description: "Top queries",
flags: &[],
},
),
];
fn build_with_ext(_: &ProviderConfig) -> Result<BuiltProvider, ObzError> {
Ok(BuiltProvider {
name: "test",
metric_query_language: None,
log_query_language: None,
metric: None,
log: None,
trace: None,
extension: None,
})
}
let meta = ProviderMeta {
name: "testprov",
display_name: "TestProvider",
aliases: &["tp"],
supported_commands: SupportedCommands::default(),
build: build_with_ext,
check: None,
command_flags: &[],
extension_commands: EXT_CMDS,
};
let mut registry = ProviderRegistry::new();
registry.register(meta);
let retrieved = registry.get("tp").unwrap();
assert_eq!(retrieved.extension_commands.len(), 2);
assert_eq!(retrieved.extension_commands[0].0, "trace");
assert_eq!(retrieved.extension_commands[0].1.name, "services");
assert_eq!(retrieved.extension_commands[1].0, "metric");
assert_eq!(retrieved.extension_commands[1].1.name, "top-queries");
}
}