Skip to main content

log_full/
quickwit.rs

1//! Quickwit 集成模块
2//! 
3//! 提供将日志发送到 Quickwit 的功能
4
5use crate::error::{LogError, LogResult};
6use std::collections::HashMap;
7use std::sync::Arc;
8use std::time::{SystemTime, UNIX_EPOCH};
9
10/// Quickwit 配置
11#[derive(Debug, Clone)]
12pub struct QuickwitConfig {
13    /// Quickwit 服务器 URL
14    pub url: String,
15    /// 索引 ID
16    pub index_id: String,
17    /// 连接超时时间(秒)
18    pub timeout: u64,
19    /// 批量发送大小
20    pub batch_size: usize,
21    /// 是否启用
22    pub enabled: bool,
23    /// 认证 token(可选)
24    pub token: Option<String>,
25}
26
27impl Default for QuickwitConfig {
28    fn default() -> Self {
29        Self {
30            url: String::new(),
31            index_id: String::new(),
32            timeout: 30,
33            batch_size: 100,
34            enabled: false,
35            token: None,
36        }
37    }
38}
39
40impl QuickwitConfig {
41    /// 创建新的 Quickwit 配置
42    pub fn new(url: String, index_id: String) -> Self {
43        Self {
44            url,
45            index_id,
46            timeout: 30,
47            batch_size: 100,
48            enabled: true,
49            token: None,
50        }
51    }
52
53    /// 设置超时时间
54    pub fn with_timeout(mut self, timeout: u64) -> Self {
55        self.timeout = timeout;
56        self
57    }
58
59    /// 设置批量大小
60    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
61        self.batch_size = batch_size;
62        self
63    }
64
65    /// 设置认证 token
66    pub fn with_token<T: Into<String>>(mut self, token: T) -> Self {
67        self.token = Some(token.into());
68        self
69    }
70
71    /// 验证配置
72    pub fn validate(&self) -> LogResult<()> {
73        if !self.enabled {
74            return Ok(());
75        }
76        
77        if self.url.is_empty() {
78            return Err(LogError::config("Quickwit URL cannot be empty"));
79        }
80        
81        if self.index_id.is_empty() {
82            return Err(LogError::config("Quickwit index_id cannot be empty"));
83        }
84        
85        if !self.url.starts_with("http://") && !self.url.starts_with("https://") {
86            return Err(LogError::config("Quickwit URL must start with http:// or https://"));
87        }
88        
89        Ok(())
90    }
91}
92
93/// Quickwit 日志条目
94#[derive(Debug, Clone)]
95pub struct QuickwitLogEntry {
96    /// 时间戳
97    pub timestamp: u64,
98    /// 日志级别
99    pub level: String,
100    /// 日志消息
101    pub message: String,
102    /// 模块路径
103    pub module: Option<String>,
104    /// 文件名
105    pub file: Option<String>,
106    /// 行号
107    pub line: Option<u32>,
108    /// 进程 ID
109    pub process_id: Option<u32>,
110    /// 线程 ID
111    pub thread_id: Option<String>,
112    /// 自定义字段
113    pub custom_fields: HashMap<String, String>,
114}
115
116impl QuickwitLogEntry {
117    /// 创建新的日志条目
118    pub fn new(level: log::Level, message: String) -> Self {
119        let timestamp = SystemTime::now()
120            .duration_since(UNIX_EPOCH)
121            .unwrap_or_default()
122            .as_secs();
123            
124        Self {
125            timestamp,
126            level: level.to_string(),
127            message,
128            module: None,
129            file: None,
130            line: None,
131            process_id: None,
132            thread_id: None,
133            custom_fields: HashMap::new(),
134        }
135    }
136
137    /// 从 log::Record 创建日志条目
138    pub fn from_record(record: &log::Record) -> Self {
139        let timestamp = SystemTime::now()
140            .duration_since(UNIX_EPOCH)
141            .unwrap_or_default()
142            .as_secs();
143            
144        let mut entry = Self {
145            timestamp,
146            level: record.level().to_string(),
147            message: format!("{}", record.args()),
148            module: Some(record.target().to_string()),
149            file: record.file().map(|f| f.to_string()),
150            line: record.line(),
151            process_id: Some(std::process::id()),
152            thread_id: std::thread::current().name().map(|n| n.to_string()),
153            custom_fields: HashMap::new(),
154        };
155        
156        // 如果没有线程名,使用线程 ID
157        if entry.thread_id.is_none() {
158            entry.thread_id = Some(format!("{:?}", std::thread::current().id()));
159        }
160        
161        entry
162    }
163
164    /// 添加自定义字段
165    pub fn add_field<K: Into<String>, V: Into<String>>(mut self, key: K, value: V) -> Self {
166        self.custom_fields.insert(key.into(), value.into());
167        self
168    }
169
170    /// 转换为 JSON 字符串
171    pub fn to_json(&self) -> LogResult<String> {
172        let escaped_message = self.message.replace('"', "\\\"");
173        let mut json = format!(
174            "{{\"timestamp\":{},\"level\":\"{}\",\"message\":\"{}\"",
175            self.timestamp,
176            self.level,
177            escaped_message
178        );
179        
180        if let Some(ref module) = self.module {
181            json.push_str(&format!(",\"module\":\"{}\"", module));
182        }
183        
184        if let Some(ref file) = self.file {
185            json.push_str(&format!(",\"file\":\"{}\"", file));
186        }
187        
188        if let Some(line) = self.line {
189            json.push_str(&format!(",\"line\":{}", line));
190        }
191        
192        if let Some(process_id) = self.process_id {
193            json.push_str(&format!(",\"process_id\":{}", process_id));
194        }
195        
196        if let Some(ref thread_id) = self.thread_id {
197            json.push_str(&format!(",\"thread_id\":\"{}\"", thread_id));
198        }
199        
200        // 添加自定义字段
201        for (key, value) in &self.custom_fields {
202            let escaped_value = value.replace('"', "\\\"");
203            json.push_str(&format!(
204                ",\"{}\":\"{}\"",
205                key,
206                escaped_value
207            ));
208        }
209        
210        json.push('}');
211        Ok(json)
212    }
213}
214
215/// Quickwit 客户端
216pub struct QuickwitClient {
217    config: QuickwitConfig,
218    client: Option<Arc<dyn HttpClient + Send + Sync>>,
219}
220
221/// HTTP 客户端 trait
222pub trait HttpClient {
223    fn post(&self, url: &str, body: &str) -> LogResult<()>;
224    fn post_with_token(&self, url: &str, body: &str, token: Option<&str>) -> LogResult<()>;
225}
226
227/// 简单的 HTTP 客户端实现(使用标准库)
228pub struct SimpleHttpClient;
229
230impl HttpClient for SimpleHttpClient {
231    fn post(&self, url: &str, body: &str) -> LogResult<()> {
232        self.post_with_token(url, body, None)
233    }
234
235    fn post_with_token(&self, url: &str, body: &str, token: Option<&str>) -> LogResult<()> {
236        // 这里使用简单的 HTTP 实现
237        // 在实际项目中,建议使用 reqwest 或其他 HTTP 客户端库
238        use std::io::Write;
239        use std::net::TcpStream;
240        
241        let url_parts: Vec<&str> = url.splitn(3, '/').collect();
242        if url_parts.len() < 3 {
243            return Err(LogError::custom("Invalid URL format"));
244        }
245        
246        let host_port = url_parts[2].split('/').next().unwrap_or("");
247        let path = &url[url.find(host_port).unwrap() + host_port.len()..];
248        
249        let (host, port) = if let Some(colon_pos) = host_port.find(':') {
250            (&host_port[..colon_pos], host_port[colon_pos + 1..].parse().unwrap_or(80))
251        } else {
252            (host_port, if url.starts_with("https") { 443 } else { 80 })
253        };
254        
255        let mut stream = TcpStream::connect((host, port))
256            .map_err(|e| LogError::custom(format!("Failed to connect: {}", e)))?;
257        
258        let mut headers = format!(
259            "POST {} HTTP/1.1\r\nHost: {}\r\nContent-Type: application/json\r\nContent-Length: {}",
260            path, host, body.len()
261        );
262        
263        // 添加 Authorization header(如果提供了 token)
264        if let Some(token) = token {
265            headers.push_str(&format!("\r\nAuthorization: Bearer {}", token));
266        }
267        
268        headers.push_str("\r\n\r\n");
269        let request = format!("{}{}", headers, body);
270        
271        stream.write_all(request.as_bytes())
272            .map_err(|e| LogError::custom(format!("Failed to send request: {}", e)))?;
273        
274        Ok(())
275    }
276}
277
278impl QuickwitClient {
279    /// 创建新的 Quickwit 客户端
280    pub fn new(config: QuickwitConfig) -> LogResult<Self> {
281        config.validate()?;
282        
283        Ok(Self {
284            config,
285            client: Some(Arc::new(SimpleHttpClient)),
286        })
287    }
288
289    /// 设置自定义 HTTP 客户端
290    pub fn with_client(mut self, client: Arc<dyn HttpClient + Send + Sync>) -> Self {
291        self.client = Some(client);
292        self
293    }
294
295    /// 发送单个日志条目
296    pub fn send_log(&self, entry: &QuickwitLogEntry) -> LogResult<()> {
297        if !self.config.enabled {
298            return Ok(());
299        }
300        
301        let json = entry.to_json()?;
302        let url = format!("{}/api/v1/{}/ingest", self.config.url, self.config.index_id);
303        
304        if let Some(ref client) = self.client {
305            client.post_with_token(&url, &json, self.config.token.as_deref())?;
306        }
307        
308        Ok(())
309    }
310
311    /// 批量发送日志条目
312    pub fn send_logs(&self, entries: &[QuickwitLogEntry]) -> LogResult<()> {
313        if !self.config.enabled || entries.is_empty() {
314            return Ok(());
315        }
316        
317        let mut json_array = String::from("[");
318        for (i, entry) in entries.iter().enumerate() {
319            if i > 0 {
320                json_array.push(',');
321            }
322            json_array.push_str(&entry.to_json()?);
323        }
324        json_array.push(']');
325        
326        self.send_json(&json_array)
327    }
328
329    /// 发送 JSON 数据到 Quickwit
330    fn send_json(&self, json: &str) -> LogResult<()> {
331        let url = format!("{}/api/v1/{}/ingest", self.config.url, self.config.index_id);
332        
333        if let Some(ref client) = self.client {
334            client.post_with_token(&url, json, self.config.token.as_deref())?;
335        }
336        
337        Ok(())
338    }
339}
340
341#[cfg(test)]
342mod tests {
343    use super::*;
344
345    #[test]
346    fn test_quickwit_config() {
347        let config = QuickwitConfig::new(
348            "http://localhost:7280".to_string(),
349            "my-index".to_string(),
350        );
351        
352        assert!(config.validate().is_ok());
353        assert!(config.enabled);
354    }
355
356    #[test]
357    fn test_log_entry_json() {
358        let entry = QuickwitLogEntry::new(log::Level::Info, "Test message".to_string())
359            .add_field("custom", "value");
360        
361        let json = entry.to_json().unwrap();
362        assert!(json.contains("Test message"));
363        assert!(json.contains("custom"));
364    }
365
366    #[test]
367    fn test_invalid_config() {
368        let config = QuickwitConfig::new(
369            "invalid-url".to_string(),
370            "my-index".to_string(),
371        );
372        
373        assert!(config.validate().is_err());
374    }
375}