use std::borrow::Cow;
use std::sync::{Arc, OnceLock};
use autumn_web::actuator::{MetricFamily, MetricsSource};
use autumn_web::app::AppBuilder;
use autumn_web::plugin::Plugin;
use autumn_web::{AppState, AutumnError};
use axum::extract::FromRef;
use crate::client::ClickHouse;
use crate::config::{ClickHouseConfig, ConfigError, DEFAULT_SECTION};
use crate::error::ClickHouseError;
use crate::health::ClickHouseCheck;
pub const PLUGIN_NAME: &str = "autumn-plugin-clickhouse";
#[derive(Default)]
pub(crate) struct Shared {
pub(crate) handle: OnceLock<ClickHouse>,
}
struct Source(Arc<Shared>);
impl MetricsSource for Source {
fn collect(&self) -> Vec<MetricFamily> {
self.0
.handle
.get()
.map_or_else(Vec::new, |db| db.metrics().families())
}
}
enum ConfigSource {
Section(String),
Explicit(Box<ClickHouseConfig>),
}
type Change = Box<dyn FnOnce(&mut ClickHouseConfig) + Send>;
#[must_use]
pub struct ClickHousePlugin {
source: ConfigSource,
changes: Vec<Change>,
}
impl Default for ClickHousePlugin {
fn default() -> Self {
Self::new()
}
}
impl ClickHousePlugin {
pub fn new() -> Self {
Self {
source: ConfigSource::Section(DEFAULT_SECTION.to_owned()),
changes: Vec::new(),
}
}
pub fn config_section(mut self, section: impl Into<String>) -> Self {
self.source = ConfigSource::Section(section.into());
self
}
pub fn config(mut self, config: ClickHouseConfig) -> Self {
self.source = ConfigSource::Explicit(Box::new(config));
self
}
pub fn configure(
mut self,
change: impl FnOnce(&mut ClickHouseConfig) + Send + 'static,
) -> Self {
self.changes.push(Box::new(change));
self
}
fn resolve(
source: &ConfigSource,
changes: Vec<Change>,
) -> Result<ClickHouseConfig, ConfigError> {
let mut config = match source {
ConfigSource::Section(section) => ClickHouseConfig::resolve(section)?,
ConfigSource::Explicit(config) => (**config).clone(),
};
for change in changes {
change(&mut config);
}
match source {
ConfigSource::Section(section) => config.validate_section(section)?,
ConfigSource::Explicit(_) => config.validate()?,
}
Ok(config)
}
}
impl Plugin for ClickHousePlugin {
fn name(&self) -> Cow<'static, str> {
Cow::Borrowed(PLUGIN_NAME)
}
fn build(self, app: AppBuilder) -> AppBuilder {
let Self { source, changes } = self;
let mut app = app;
if let ConfigSource::Section(section) = &source {
app = app.config_section(section.clone());
}
let resolved = Self::resolve(&source, changes);
let shared = Arc::new(Shared::default());
app = app.metrics_source("clickhouse", Arc::new(Source(Arc::clone(&shared))));
if resolved.as_ref().is_ok_and(|config| config.health_check) {
app = app.health_indicator(
"clickhouse",
Arc::new(ClickHouseCheck::new(Arc::clone(&shared))),
);
}
let on_start = Arc::clone(&shared);
let resolved = Arc::new(resolved);
app.on_startup(move |state| {
let shared = Arc::clone(&on_start);
let resolved = Arc::clone(&resolved);
async move {
let config = resolved.as_ref().clone().map_err(|err| boot_error(&err))?;
let db = ClickHouse::new(&config);
db.ping().await.map_err(|err| boot_error(&err))?;
state.insert_extension(db.clone());
let _ = shared.handle.set(db);
tracing::info!("the ClickHouse plugin is ready");
Ok(())
}
})
}
}
fn boot_error(err: &dyn std::fmt::Display) -> AutumnError {
AutumnError::internal_server_error_msg(format!("{PLUGIN_NAME}: {err}"))
}
impl std::fmt::Debug for ClickHousePlugin {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let source = match &self.source {
ConfigSource::Section(section) => section.as_str(),
ConfigSource::Explicit(_) => "(explicit)",
};
f.debug_struct("ClickHousePlugin")
.field("config", &source)
.field("changes", &self.changes.len())
.finish()
}
}
impl ClickHouse {
#[must_use]
pub fn from_state(state: &AppState) -> Option<Self> {
state.extension::<Self>().map(|db| (*db).clone())
}
}
impl<S> axum::extract::FromRequestParts<S> for ClickHouse
where
AppState: FromRef<S>,
S: Send + Sync,
{
type Rejection = AutumnError;
async fn from_request_parts(
_parts: &mut http::request::Parts,
state: &S,
) -> Result<Self, Self::Rejection> {
let app_state = AppState::from_ref(state);
Self::from_state(&app_state).ok_or_else(|| ClickHouseError::NotInstalled.into_autumn())
}
}
#[cfg(test)]
mod tests;