use std::collections::HashMap;
use std::path::PathBuf;
use redis::aio::MultiplexedConnection;
use serde::{Deserialize, Serialize};
use crate::error::Result;
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct TableState {
pub snapshot_id: Option<i64>,
pub watermark: Option<String>,
#[serde(default)]
pub done_files: Vec<String>,
pub synced_at: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cdc: Option<crate::cdc::CdcCapability>,
}
#[derive(Debug)]
pub enum SyncState {
File {
path: PathBuf,
tables: HashMap<String, TableState>,
dirty: bool,
},
Valkey {
conn: MultiplexedConnection,
key: String,
tables: HashMap<String, TableState>,
},
}
impl SyncState {
pub fn open(path: impl Into<PathBuf>) -> Result<Self> {
let path = path.into();
let tables = match std::fs::read_to_string(&path) {
Ok(s) => serde_json::from_str(&s)?,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => HashMap::new(),
Err(e) => return Err(e.into()),
};
Ok(SyncState::File {
path,
tables,
dirty: false,
})
}
pub async fn open_valkey(url: &str, key: impl Into<String>) -> Result<Self> {
let key = key.into();
let client = redis::Client::open(url)?;
let mut conn = client.get_multiplexed_async_connection().await?;
let raw: HashMap<String, String> = redis::cmd("HGETALL")
.arg(&key)
.query_async(&mut conn)
.await?;
let mut tables = HashMap::with_capacity(raw.len());
for (table, json) in raw {
tables.insert(table, serde_json::from_str(&json)?);
}
Ok(SyncState::Valkey { conn, key, tables })
}
fn cache(&self) -> &HashMap<String, TableState> {
match self {
SyncState::File { tables, .. } | SyncState::Valkey { tables, .. } => tables,
}
}
pub fn get(&self, table: &str) -> TableState {
self.cache().get(table).cloned().unwrap_or_default()
}
pub fn get_ref(&self, table: &str) -> Option<&TableState> {
self.cache().get(table)
}
pub async fn advance(&mut self, table: &str, mut state: TableState) -> Result<()> {
state.synced_at = Some(now_iso8601());
match self {
SyncState::File {
path,
tables,
dirty,
} => {
tables.insert(table.to_string(), state);
persist_file(path, tables)?;
*dirty = false;
Ok(())
}
SyncState::Valkey { conn, key, tables } => {
write_through_valkey(conn, key, tables, table, state).await
}
}
}
pub async fn advance_staged(&mut self, table: &str, mut state: TableState) -> Result<()> {
state.synced_at = Some(now_iso8601());
match self {
SyncState::File { tables, dirty, .. } => {
tables.insert(table.to_string(), state);
*dirty = true;
Ok(())
}
SyncState::Valkey { conn, key, tables } => {
write_through_valkey(conn, key, tables, table, state).await
}
}
}
pub fn flush(&mut self) -> Result<()> {
match self {
SyncState::File {
path,
tables,
dirty,
} => {
if *dirty {
persist_file(path, tables)?;
*dirty = false;
}
Ok(())
}
SyncState::Valkey { .. } => Ok(()),
}
}
}
async fn write_through_valkey(
conn: &mut MultiplexedConnection,
key: &str,
tables: &mut HashMap<String, TableState>,
table: &str,
state: TableState,
) -> Result<()> {
let json = serde_json::to_string(&state)?;
redis::cmd("HSET")
.arg(key)
.arg(table)
.arg(json)
.query_async::<()>(conn)
.await?;
tables.insert(table.to_string(), state);
Ok(())
}
fn persist_file(path: &PathBuf, tables: &HashMap<String, TableState>) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let tmp = path.with_extension("tmp");
std::fs::write(&tmp, serde_json::to_vec(tables)?)?;
std::fs::rename(&tmp, path)?;
Ok(())
}
fn now_iso8601() -> String {
let secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
chrono::DateTime::<chrono::Utc>::from_timestamp(secs as i64, 0)
.unwrap_or_else(|| chrono::DateTime::<chrono::Utc>::from_timestamp(0, 0).unwrap())
.to_rfc3339()
}
#[cfg(test)]
mod now_iso8601_tests {
use super::now_iso8601;
#[test]
fn synced_at_is_parsable_iso8601() {
let s = now_iso8601();
assert!(
!s.starts_with('@'),
"synced_at must not use the @epoch encoding, got {s:?}"
);
let parsed = chrono::DateTime::parse_from_rfc3339(&s);
assert!(parsed.is_ok(), "synced_at {s:?} must parse as RFC3339");
}
}