use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;
use async_trait::async_trait;
use parking_lot::Mutex;
use rskit_bootstrap::{Component, Health};
use rskit_errors::{AppError, AppResult, ErrorCode};
use serde::{Deserialize, Serialize};
use crate::config::{DatabaseConfig, MemoryDatabaseConfig};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatabaseQuery {
pub statement: String,
#[serde(default)]
pub parameters: Vec<serde_json::Value>,
}
impl DatabaseQuery {
#[must_use]
pub fn new(statement: impl Into<String>) -> Self {
Self {
statement: statement.into(),
parameters: Vec::new(),
}
}
#[must_use]
pub fn with_parameter(mut self, value: impl Into<serde_json::Value>) -> Self {
self.parameters.push(value.into());
self
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct DatabaseResult {
pub rows_affected: u64,
}
#[async_trait]
pub trait DatabaseTransaction: Send + Sync {
async fn execute(&self, query: DatabaseQuery) -> AppResult<DatabaseResult>;
async fn commit(self: Box<Self>) -> AppResult<()>;
async fn rollback(self: Box<Self>) -> AppResult<()>;
}
#[async_trait]
pub trait DatabaseClient: Send + Sync {
async fn execute(&self, query: DatabaseQuery) -> AppResult<DatabaseResult>;
async fn begin(&self) -> AppResult<Box<dyn DatabaseTransaction>>;
async fn ping(&self) -> AppResult<()>;
}
#[derive(Debug)]
pub struct InMemoryDatabase {
config: MemoryDatabaseConfig,
connected: AtomicBool,
history: Mutex<VecDeque<DatabaseQuery>>,
}
impl InMemoryDatabase {
#[must_use]
pub fn new(config: MemoryDatabaseConfig) -> Self {
Self {
config,
connected: AtomicBool::new(true),
history: Mutex::new(VecDeque::new()),
}
}
#[must_use]
pub fn recorded_queries(&self) -> Vec<DatabaseQuery> {
self.history.lock().iter().cloned().collect()
}
fn record(&self, query: DatabaseQuery) {
if self.config.statement_history == 0 {
return;
}
let mut history = self.history.lock();
if history.len() == self.config.statement_history {
history.pop_front();
}
history.push_back(query);
}
}
impl Default for InMemoryDatabase {
fn default() -> Self {
Self::new(MemoryDatabaseConfig::default())
}
}
#[async_trait]
impl DatabaseClient for InMemoryDatabase {
async fn execute(&self, query: DatabaseQuery) -> AppResult<DatabaseResult> {
if !self.connected.load(Ordering::SeqCst) {
return Err(AppError::new(
ErrorCode::ConnectionFailed,
"in-memory database is stopped",
));
}
if query.statement.trim().is_empty() {
return Err(AppError::new(
ErrorCode::InvalidInput,
"database query statement is required",
));
}
self.record(query);
Ok(DatabaseResult { rows_affected: 1 })
}
async fn begin(&self) -> AppResult<Box<dyn DatabaseTransaction>> {
self.ping().await?;
Ok(Box::new(InMemoryTransaction {
started_at: Instant::now(),
queries: Mutex::new(Vec::new()),
committed: AtomicBool::new(false),
rolled_back: AtomicBool::new(false),
}))
}
async fn ping(&self) -> AppResult<()> {
if self.connected.load(Ordering::SeqCst) {
Ok(())
} else {
Err(AppError::new(
ErrorCode::ConnectionFailed,
"in-memory database is stopped",
))
}
}
}
#[async_trait]
impl Component for InMemoryDatabase {
fn name(&self) -> &str {
"database"
}
async fn start(&self) -> AppResult<()> {
self.connected.store(true, Ordering::SeqCst);
Ok(())
}
async fn stop(&self) -> AppResult<()> {
self.connected.store(false, Ordering::SeqCst);
Ok(())
}
fn health(&self) -> Health {
if self.connected.load(Ordering::SeqCst) {
Health::healthy("database")
} else {
Health::unhealthy("database", "backend is stopped")
}
}
}
#[derive(Debug)]
struct InMemoryTransaction {
started_at: Instant,
queries: Mutex<Vec<DatabaseQuery>>,
committed: AtomicBool,
rolled_back: AtomicBool,
}
#[async_trait]
impl DatabaseTransaction for InMemoryTransaction {
async fn execute(&self, query: DatabaseQuery) -> AppResult<DatabaseResult> {
if query.statement.trim().is_empty() {
return Err(AppError::new(
ErrorCode::InvalidInput,
"database query statement is required",
));
}
self.queries.lock().push(query);
Ok(DatabaseResult { rows_affected: 1 })
}
async fn commit(self: Box<Self>) -> AppResult<()> {
if self.rolled_back.load(Ordering::SeqCst) {
return Err(AppError::new(
ErrorCode::InvalidInput,
"database transaction was already rolled back",
));
}
self.committed.store(true, Ordering::SeqCst);
let _elapsed = self.started_at.elapsed();
Ok(())
}
async fn rollback(self: Box<Self>) -> AppResult<()> {
if self.committed.load(Ordering::SeqCst) {
return Err(AppError::new(
ErrorCode::InvalidInput,
"database transaction was already committed",
));
}
self.rolled_back.store(true, Ordering::SeqCst);
Ok(())
}
}
pub(crate) fn memory_from_config(config: &DatabaseConfig) -> InMemoryDatabase {
InMemoryDatabase::new(config.memory.clone())
}