#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use std::sync::Arc;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use std::sync::atomic::AtomicBool;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use std::time::{Duration, Instant};
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use parking_lot::Mutex;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use super::CircuitBreaker;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use super::FileSink;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use crate::DataMasker;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use crate::LogRecord;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
use crate::Metrics;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
mod database_impl;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
pub use database_impl::convert_logs_to_parquet;
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
pub(crate) use database_impl::effective_db_worker_limit;
#[cfg(all(
any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
),
test
))]
use database_impl::{ADAPTIVE_WINDOW_SIZE, MAX_BATCH_SIZE, MIN_BATCH_SIZE};
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
struct DatabaseSinkInner {
buffer: Vec<LogRecord>,
flush_buffer: Vec<LogRecord>,
last_flush: Instant,
fallback_sink: Option<FileSink>,
circuit_breaker: CircuitBreaker,
current_batch_size: usize,
write_latencies: Vec<Duration>,
success_count: usize,
failure_count: usize,
metrics: Option<Arc<Metrics>>,
}
#[cfg(any(
feature = "sqlite",
feature = "postgres",
feature = "mysql",
feature = "duckdb"
))]
pub struct DatabaseSink {
inner: Mutex<DatabaseSinkInner>,
database: Arc<dyn crate::integrations::infra::Database>,
masker: Arc<DataMasker>,
stop: Arc<AtomicBool>,
}
#[cfg(test)]
mod tests {
use super::super::LogSink;
use super::*;
use crate::DatabaseSinkConfig;
use crate::InklogError;
use crate::LogRecord;
use crate::Metrics;
use crate::integrations::MockDatabaseAdapter;
use std::sync::Arc;
#[test]
#[cfg(feature = "parquet")]
fn test_convert_logs_to_parquet_empty() {
let logs: Vec<LogRecord> = vec![];
let config = DatabaseSinkConfig::default().parquet_config;
let result = convert_logs_to_parquet(&logs, &config);
assert!(result.is_ok());
let bytes = result.unwrap();
assert!(!bytes.is_empty()); }
#[test]
#[cfg(feature = "parquet")]
fn test_convert_logs_to_parquet_non_empty() {
let log1 = LogRecord::default();
let log2 = LogRecord {
message: "warn message".into(),
..Default::default()
};
let config = DatabaseSinkConfig::default().parquet_config;
let result = convert_logs_to_parquet(&[log1, log2], &config);
assert!(result.is_ok());
assert!(!result.unwrap().is_empty());
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_write_with_mock_db() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let config = DatabaseSinkConfig::default();
let sink = DatabaseSink::new_with_config(mock_db.clone(), Some(config)).unwrap();
let metrics = Arc::new(Metrics::new());
sink.set_metrics(metrics.clone()).await;
let record = LogRecord::default();
let result = sink.write(&record).await;
assert!(result.is_ok());
let flush_result = sink.flush().await;
assert!(flush_result.is_ok());
assert_eq!(mock_db.stored_count(), 1);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_new_without_config() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db.clone()).unwrap();
let record = LogRecord::default();
let result = sink.write(&record).await;
assert!(result.is_ok());
let flush_result = sink.flush().await;
assert!(flush_result.is_ok());
assert_eq!(mock_db.stored_count(), 1);
}
#[test]
fn test_database_sink_display() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db).unwrap();
let s = format!("{}", sink);
assert_eq!(s, "DatabaseSink");
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_flush_empty_buffer() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db).unwrap();
let result = sink.flush().await;
assert!(result.is_ok());
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_shutdown() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db.clone()).unwrap();
let record = LogRecord::default();
let _ = sink.write(&record).await;
let result = sink.shutdown().await;
assert!(result.is_ok());
assert_eq!(mock_db.stored_count(), 1);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_buffer_full_triggers_flush() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(mock_db.clone(), Some(config)).unwrap();
for i in 0..10 {
let record = LogRecord {
message: format!("message {}", i),
..Default::default()
};
let result = sink.write(&record).await;
assert!(result.is_ok(), "Write {} failed: {:?}", i, result.err());
}
assert_eq!(mock_db.stored_count(), 10);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_flush_timeout_triggers_flush() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let config = DatabaseSinkConfig {
batch_size: 1000,
..Default::default()
};
let sink = DatabaseSink::new_with_config(mock_db.clone(), Some(config)).unwrap();
let record = LogRecord::default();
let _ = sink.write(&record).await;
assert_eq!(mock_db.stored_count(), 0);
tokio::time::sleep(Duration::from_millis(600)).await;
let record2 = LogRecord::default();
let result = sink.write(&record2).await;
assert!(result.is_ok());
assert_eq!(mock_db.stored_count(), 2);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_applies_masking() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db.clone()).unwrap();
let record = LogRecord {
message: "User email: test@example.com".to_string(),
..Default::default()
};
let _ = sink.write(&record).await;
let _ = sink.flush().await;
let records = mock_db.get_records();
assert_eq!(records.len(), 1);
assert!(
!records[0].message.contains("test@example.com"),
"Message should be masked, got: {}",
records[0].message
);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_set_metrics() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let sink = DatabaseSink::new(mock_db.clone()).unwrap();
let metrics = Arc::new(Metrics::new());
sink.set_metrics(metrics.clone()).await;
let record = LogRecord::default();
let _ = sink.write(&record).await;
let _ = sink.flush().await;
assert!(metrics.db_batch_records_total() > 0);
}
struct FailingDatabase;
#[async_trait::async_trait]
impl crate::integrations::infra::Database for FailingDatabase {
async fn insert_batch(&self, _records: &[LogRecord]) -> Result<usize, InklogError> {
Err(InklogError::DatabaseError {
message: "Simulated database failure".to_string(),
source: None,
})
}
async fn is_healthy(&self) -> bool {
false
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_flush_failure_returns_error() {
let failing_db = Arc::new(FailingDatabase);
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(failing_db, Some(config)).unwrap();
for _ in 0..10 {
let record = LogRecord::default();
let _ = sink.write(&record).await;
}
let result = sink.flush().await;
let _ = result;
}
#[tokio::test(flavor = "multi_thread")]
async fn test_database_sink_shutdown_with_failing_db() {
let failing_db = Arc::new(FailingDatabase);
let sink = DatabaseSink::new(failing_db).unwrap();
let record = LogRecord::default();
let _ = sink.write(&record).await;
let result = sink.shutdown().await;
assert!(result.is_ok());
}
#[tokio::test(flavor = "multi_thread")]
async fn test_adjust_batch_size_grows_on_high_success_and_low_latency() {
let mock_db = Arc::new(MockDatabaseAdapter::new());
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(mock_db.clone(), Some(config)).unwrap();
let metrics = Arc::new(Metrics::new());
sink.set_metrics(metrics.clone()).await;
for i in 0..100 {
let record = LogRecord {
message: format!("adjust-grow-{}", i),
..Default::default()
};
let result = sink.write(&record).await;
assert!(result.is_ok(), "write {} failed: {:?}", i, result.err());
}
assert_eq!(
mock_db.stored_count(),
100,
"all 100 records should be flushed to db"
);
assert!(metrics.db_batch_records_total() > 0);
}
struct SlowDatabase;
#[async_trait::async_trait]
impl crate::integrations::infra::Database for SlowDatabase {
async fn insert_batch(&self, _records: &[LogRecord]) -> Result<usize, InklogError> {
tokio::time::sleep(Duration::from_millis(250)).await;
Ok(_records.len())
}
async fn is_healthy(&self) -> bool {
true
}
}
#[tokio::test(flavor = "multi_thread")]
async fn test_adjust_batch_size_shrinks_on_high_latency() {
let slow_db = Arc::new(SlowDatabase);
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(slow_db, Some(config)).unwrap();
for i in 0..100 {
let record = LogRecord {
message: format!("adjust-shrink-{}", i),
..Default::default()
};
let result = sink.write(&record).await;
assert!(result.is_ok(), "write {} failed: {:?}", i, result.err());
}
let before = sink.flush().await;
let _ = before;
}
#[tokio::test(flavor = "multi_thread")]
async fn test_circuit_breaker_open_routes_to_fallback_sink() {
let failing_db = Arc::new(FailingDatabase);
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(failing_db, Some(config)).unwrap();
for i in 0..12 {
let record = LogRecord {
message: format!("cb-open-{}", i),
..Default::default()
};
let _ = sink.write(&record).await;
}
let record = LogRecord {
message: "after circuit open".to_string(),
..Default::default()
};
let result = sink.write(&record).await;
assert!(
result.is_ok(),
"write after circuit open should route to fallback and return Ok, got: {:?}",
result
);
}
#[tokio::test(flavor = "multi_thread")]
async fn test_flush_failure_updates_sink_error_metrics() {
let failing_db = Arc::new(FailingDatabase);
let config = DatabaseSinkConfig {
batch_size: 10,
..Default::default()
};
let sink = DatabaseSink::new_with_config(failing_db, Some(config)).unwrap();
let metrics = Arc::new(Metrics::new());
sink.set_metrics(metrics.clone()).await;
for _ in 0..10 {
let record = LogRecord::default();
let _ = sink.write(&record).await;
}
assert!(
metrics.sink_errors() >= 1,
"flush failure should increment sink_errors, got: {}",
metrics.sink_errors()
);
}
fn make_test_inner(batch_size: usize) -> DatabaseSinkInner {
DatabaseSinkInner {
buffer: Vec::new(),
flush_buffer: Vec::new(),
last_flush: Instant::now(),
fallback_sink: None,
circuit_breaker: CircuitBreaker::new(3, Duration::from_secs(30), 3),
current_batch_size: batch_size,
write_latencies: Vec::new(),
success_count: 0,
failure_count: 0,
metrics: None,
}
}
#[test]
fn test_adjust_batch_size_early_return_when_insufficient_latencies() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(10); 5];
inner.success_count = 5;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 100,
"batch_size should not change on early return"
);
assert_eq!(
inner.write_latencies.len(),
5,
"latencies should not be cleared on early return"
);
assert_eq!(
inner.success_count, 5,
"success_count should not be cleared on early return"
);
assert_eq!(
inner.failure_count, 0,
"failure_count should not be cleared on early return"
);
}
#[test]
fn test_adjust_batch_size_grows_on_high_success_low_latency() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(10); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 10;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 200,
"batch_size should double from 100 to 200"
);
assert!(
inner.write_latencies.is_empty(),
"latencies should be cleared"
);
assert_eq!(inner.success_count, 0, "success_count should be cleared");
assert_eq!(inner.failure_count, 0, "failure_count should be cleared");
}
#[test]
fn test_adjust_batch_size_shrinks_on_high_latency_direct() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(300); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 10;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 50,
"batch_size should halve from 100 to 50"
);
assert!(
inner.write_latencies.is_empty(),
"latencies should be cleared"
);
}
#[test]
fn test_adjust_batch_size_shrinks_on_low_success_rate() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(10); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 5;
inner.failure_count = 5;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 50,
"batch_size should halve when success_rate < 0.8"
);
assert!(
inner.write_latencies.is_empty(),
"latencies should be cleared"
);
}
#[test]
fn test_adjust_batch_size_grow_capped_at_max() {
let mut inner = make_test_inner(600);
inner.write_latencies = vec![Duration::from_millis(10); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 10;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, MAX_BATCH_SIZE,
"batch_size should be capped at MAX_BATCH_SIZE"
);
}
#[test]
fn test_adjust_batch_size_shrink_floored_at_min() {
let mut inner = make_test_inner(MIN_BATCH_SIZE);
inner.write_latencies = vec![Duration::from_millis(300); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 10;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, MIN_BATCH_SIZE,
"batch_size should be floored at MIN_BATCH_SIZE"
);
}
#[test]
fn test_adjust_batch_size_no_change_in_middle_zone() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(100); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 9;
inner.failure_count = 1;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 100,
"batch_size should not change in middle zone"
);
assert!(
inner.write_latencies.is_empty(),
"latencies should be cleared even without adjustment"
);
assert_eq!(inner.success_count, 0, "success_count should be cleared");
assert_eq!(inner.failure_count, 0, "failure_count should be cleared");
}
#[test]
fn test_adjust_batch_size_total_ops_zero_uses_default_success_rate() {
let mut inner = make_test_inner(100);
inner.write_latencies = vec![Duration::from_millis(10); ADAPTIVE_WINDOW_SIZE];
inner.success_count = 0;
inner.failure_count = 0;
DatabaseSink::adjust_batch_size(&mut inner);
assert_eq!(
inner.current_batch_size, 200,
"should grow when total_ops=0 (success_rate defaults to 1.0)"
);
}
#[test]
fn test_effective_db_worker_limit() {
let limit = effective_db_worker_limit();
assert!(limit >= 1, "limit should be at least 1, got {}", limit);
assert!(
limit <= database_impl::MAX_DB_WORKER_LIMIT,
"limit should be at most {}, got {}",
database_impl::MAX_DB_WORKER_LIMIT,
limit
);
}
}