301 lines
8.3 KiB
Rust
301 lines
8.3 KiB
Rust
//! State management and persistence
|
|
|
|
use crate::{EtlError, Result};
|
|
use chrono::{DateTime, Utc};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use tokio::sync::RwLock;
|
|
|
|
/// State management configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StateConfig {
|
|
pub backend: StateBackendType,
|
|
pub checkpoint_interval_ms: u64,
|
|
pub cleanup_interval_ms: u64,
|
|
}
|
|
|
|
impl Default for StateConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
backend: StateBackendType::InMemory,
|
|
checkpoint_interval_ms: 30000,
|
|
cleanup_interval_ms: 300000,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// State backend types
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub enum StateBackendType {
|
|
InMemory,
|
|
Redis { connection_string: String },
|
|
PostgreSQL { connection_string: String },
|
|
}
|
|
|
|
/// State manager
|
|
#[derive(Debug)]
|
|
pub struct StateManager {
|
|
config: StateConfig,
|
|
backend: Arc<dyn StateBackend>,
|
|
state_cache: Arc<RwLock<HashMap<String, StateEntry>>>,
|
|
}
|
|
|
|
/// State entry
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StateEntry {
|
|
pub key: String,
|
|
pub value: serde_json::Value,
|
|
pub timestamp: DateTime<Utc>,
|
|
pub ttl: Option<u64>,
|
|
}
|
|
|
|
/// State backend trait
|
|
pub trait StateBackend: Send + Sync + std::fmt::Debug {
|
|
fn get(
|
|
&self,
|
|
key: &str,
|
|
) -> std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
|
|
>;
|
|
fn set(
|
|
&self,
|
|
key: &str,
|
|
value: serde_json::Value,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>>;
|
|
fn delete(
|
|
&self,
|
|
key: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>>;
|
|
fn list_keys(
|
|
&self,
|
|
prefix: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>>;
|
|
}
|
|
|
|
/// State snapshot
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct StateSnapshot {
|
|
pub snapshot_id: String,
|
|
pub timestamp: DateTime<Utc>,
|
|
pub entries: HashMap<String, serde_json::Value>,
|
|
}
|
|
|
|
/// State recovery
|
|
#[derive(Debug)]
|
|
pub struct StateRecovery {
|
|
state_manager: Arc<StateManager>,
|
|
}
|
|
|
|
/// Redis state manager
|
|
#[derive(Debug)]
|
|
pub struct RedisStateManager {
|
|
connection_string: String,
|
|
}
|
|
|
|
/// `PostgreSQL` state manager
|
|
#[derive(Debug)]
|
|
pub struct PostgresStateManager {
|
|
connection_string: String,
|
|
}
|
|
|
|
/// In-memory state backend
|
|
#[derive(Debug)]
|
|
struct InMemoryBackend {
|
|
data: Arc<RwLock<HashMap<String, serde_json::Value>>>,
|
|
}
|
|
|
|
impl StateManager {
|
|
pub async fn new(config: StateConfig) -> Result<Self> {
|
|
let backend: Arc<dyn StateBackend> = match &config.backend {
|
|
StateBackendType::InMemory => Arc::new(InMemoryBackend {
|
|
data: Arc::new(RwLock::new(HashMap::new())),
|
|
}),
|
|
StateBackendType::Redis { connection_string } => Arc::new(RedisStateManager {
|
|
connection_string: connection_string.clone(),
|
|
}),
|
|
StateBackendType::PostgreSQL { connection_string } => Arc::new(PostgresStateManager {
|
|
connection_string: connection_string.clone(),
|
|
}),
|
|
};
|
|
|
|
Ok(Self {
|
|
config,
|
|
backend,
|
|
state_cache: Arc::new(RwLock::new(HashMap::new())),
|
|
})
|
|
}
|
|
|
|
pub async fn get_state(&self, key: &str) -> Result<Option<serde_json::Value>> {
|
|
self.backend.get(key).await
|
|
}
|
|
|
|
pub async fn set_state(&self, key: &str, value: serde_json::Value) -> Result<()> {
|
|
self.backend.set(key, value).await
|
|
}
|
|
|
|
pub async fn update_task_state(&self, task_id: &str, state_value: &str) -> Result<()> {
|
|
let key = format!("task_state_{task_id}");
|
|
let value = serde_json::json!(state_value);
|
|
self.set_state(&key, value).await
|
|
}
|
|
|
|
pub async fn create_snapshot(&self) -> Result<StateSnapshot> {
|
|
let keys = self.backend.list_keys("").await?;
|
|
let mut entries = HashMap::new();
|
|
|
|
for key in keys {
|
|
if let Some(value) = self.backend.get(&key).await? {
|
|
entries.insert(key, value);
|
|
}
|
|
}
|
|
|
|
Ok(StateSnapshot {
|
|
snapshot_id: uuid::Uuid::new_v4().to_string(),
|
|
timestamp: Utc::now(),
|
|
entries,
|
|
})
|
|
}
|
|
}
|
|
|
|
impl StateBackend for InMemoryBackend {
|
|
fn get(
|
|
&self,
|
|
key: &str,
|
|
) -> std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
|
|
> {
|
|
let key = key.to_string();
|
|
let data = self.data.clone();
|
|
Box::pin(async move {
|
|
let data = data.read().await;
|
|
Ok(data.get(&key).cloned())
|
|
})
|
|
}
|
|
|
|
fn set(
|
|
&self,
|
|
key: &str,
|
|
value: serde_json::Value,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
let key = key.to_string();
|
|
let data = self.data.clone();
|
|
Box::pin(async move {
|
|
let mut data = data.write().await;
|
|
data.insert(key, value);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
fn delete(
|
|
&self,
|
|
key: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
let key = key.to_string();
|
|
let data = self.data.clone();
|
|
Box::pin(async move {
|
|
let mut data = data.write().await;
|
|
data.remove(&key);
|
|
Ok(())
|
|
})
|
|
}
|
|
|
|
fn list_keys(
|
|
&self,
|
|
prefix: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
|
|
let prefix = prefix.to_string();
|
|
let data = self.data.clone();
|
|
Box::pin(async move {
|
|
let data = data.read().await;
|
|
Ok(data
|
|
.keys()
|
|
.filter(|key| key.starts_with(&prefix))
|
|
.cloned()
|
|
.collect())
|
|
})
|
|
}
|
|
}
|
|
|
|
impl StateBackend for RedisStateManager {
|
|
fn get(
|
|
&self,
|
|
_key: &str,
|
|
) -> std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
|
|
> {
|
|
Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
|
|
}
|
|
|
|
fn set(
|
|
&self,
|
|
_key: &str,
|
|
_value: serde_json::Value,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
|
|
}
|
|
|
|
fn delete(
|
|
&self,
|
|
_key: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
|
|
}
|
|
|
|
fn list_keys(
|
|
&self,
|
|
_prefix: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
|
|
Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
|
|
}
|
|
}
|
|
|
|
impl StateBackend for PostgresStateManager {
|
|
fn get(
|
|
&self,
|
|
_key: &str,
|
|
) -> std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
|
|
> {
|
|
Box::pin(async move {
|
|
Err(EtlError::Other(
|
|
"PostgreSQL backend not implemented".to_string(),
|
|
))
|
|
})
|
|
}
|
|
|
|
fn set(
|
|
&self,
|
|
_key: &str,
|
|
_value: serde_json::Value,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
Box::pin(async move {
|
|
Err(EtlError::Other(
|
|
"PostgreSQL backend not implemented".to_string(),
|
|
))
|
|
})
|
|
}
|
|
|
|
fn delete(
|
|
&self,
|
|
_key: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
|
|
Box::pin(async move {
|
|
Err(EtlError::Other(
|
|
"PostgreSQL backend not implemented".to_string(),
|
|
))
|
|
})
|
|
}
|
|
|
|
fn list_keys(
|
|
&self,
|
|
_prefix: &str,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
|
|
Box::pin(async move {
|
|
Err(EtlError::Other(
|
|
"PostgreSQL backend not implemented".to_string(),
|
|
))
|
|
})
|
|
}
|
|
}
|