axiom_log.rs

use log::{Level, Metadata, Record};
use serde_json::{json, Value as JsonValue};
use std::sync::{Arc, Mutex};

pub struct Logger {
    axiom_token: String,
    axiom_dataset: String,
    client: reqwest::blocking::Client,

    buffer: Arc<Mutex<Vec<JsonValue>>>,
}

impl Logger {
    pub fn new(axiom_token: impl Into<String>, axiom_dataset: impl Into<String>) -> Self {
        Self {
            axiom_token: axiom_token.into(),
            axiom_dataset: axiom_dataset.into(),
            client: reqwest::blocking::Client::new(),
            buffer: Arc::new(Mutex::new(Vec::new())),
        }
    }
}

impl log::Log for Logger {
    fn enabled(&self, metadata: &Metadata) -> bool {
        metadata.level() <= Level::Info
    }

    fn log(&self, record: &Record) {
        if self.enabled(record.metadata()) {
            let event = json!({
                "level": record.level(),
                "target": record.target(),
                "message": record.args().to_string(),
            });

            if let Ok(mut buffer) = self.buffer.lock() {
                buffer.push(event);
            } else {
                eprintln!("Failed to add event to buffer");
            }
        }
    }

    fn flush(&self) {
        if let Ok(mut buffer) = self.buffer.lock() {
            if buffer.is_empty() {
                return;
            }

            match self
                .client
                .post(format!(
                    "https://api.axiom.co/v1/datasets/{}/ingest",
                    &self.axiom_dataset
                ))
                .bearer_auth(&self.axiom_token)
                .header("Content-Type", "application/json")
                .json(&**buffer)
                .send()
            {
                Ok(res) => {
                    if res.status().is_success() {
                        *buffer = Vec::new(); // Clear buffer
                    } else {
                        eprintln!("Failed to send events to Axiom: {}", res.status());
                    }
                }
                Err(e) => {
                    eprintln!("Failed to send events to Axiom: {}", e);
                }
            }
        }
    }
}