1use crate::backend::{BackendClient, BackendConfig, ParamValue};
15use crate::transaction_journal::{JournalValue, TransactionJournal};
16use crate::{ProxyError, Result};
17use chrono::{DateTime, Utc};
18use std::sync::Arc;
19
20#[derive(Debug, Clone)]
22pub struct TimeTravelRequest {
23 pub from: DateTime<Utc>,
25 pub to: DateTime<Utc>,
27 pub target_host: String,
29 pub target_port: u16,
31 pub target_user: Option<String>,
35 pub target_password: Option<String>,
39 pub target_database: Option<String>,
41}
42
43#[derive(Debug, Clone, serde::Serialize)]
45pub struct ReplaySummary {
46 pub statements_replayed: u64,
48 pub failures: u64,
50 pub elapsed_ms: u64,
52 #[serde(with = "chrono::serde::ts_seconds")]
54 pub from: DateTime<Utc>,
55 #[serde(with = "chrono::serde::ts_seconds")]
56 pub to: DateTime<Utc>,
57 pub first_error: Option<String>,
60}
61
62pub struct ReplayEngine {
64 journal: Arc<TransactionJournal>,
65 backend_template: BackendConfig,
67}
68
69impl ReplayEngine {
70 pub fn new(journal: Arc<TransactionJournal>, backend_template: BackendConfig) -> Self {
71 Self {
72 journal,
73 backend_template,
74 }
75 }
76
77 pub async fn replay_window(&self, req: &TimeTravelRequest) -> Result<ReplaySummary> {
86 if req.from > req.to {
87 return Err(ProxyError::Internal("replay window: from > to".to_string()));
88 }
89
90 let entries = self.journal.entries_in_window(req.from, req.to).await;
91 let total = entries.len();
92 tracing::info!(
93 total_entries = total,
94 from = %req.from,
95 to = %req.to,
96 target = %format!("{}:{}", req.target_host, req.target_port),
97 "starting time-travel replay"
98 );
99
100 let mut cfg = self.backend_template.clone();
101 cfg.host = req.target_host.clone();
102 cfg.port = req.target_port;
103 if let Some(ref u) = req.target_user {
104 cfg.user = u.clone();
105 }
106 if let Some(ref p) = req.target_password {
107 cfg.password = Some(p.clone());
108 }
109 if let Some(ref d) = req.target_database {
110 cfg.database = Some(d.clone());
111 }
112
113 let start = std::time::Instant::now();
114 let mut client = BackendClient::connect(&cfg)
115 .await
116 .map_err(|e| ProxyError::ReplayFailed(format!("connect to target: {}", e)))?;
117
118 let mut statements_replayed: u64 = 0;
119 let mut failures: u64 = 0;
120 let mut first_error: Option<String> = None;
121
122 for (tx_id, entry) in entries {
123 let params: Vec<ParamValue> = entry
124 .parameters
125 .iter()
126 .map(journal_value_to_param)
127 .collect();
128
129 let outcome = if params.is_empty() {
130 client.simple_query(&entry.statement).await
131 } else {
132 client.query_with_params(&entry.statement, ¶ms).await
133 };
134
135 match outcome {
136 Ok(_) => {
137 statements_replayed += 1;
138 }
139 Err(e) => {
140 failures += 1;
141 if first_error.is_none() {
142 first_error = Some(format!("tx {} seq {}: {}", tx_id, entry.sequence, e));
143 }
144 tracing::warn!(
145 tx = %tx_id,
146 sequence = entry.sequence,
147 error = %e,
148 "replay statement failed"
149 );
150 }
151 }
152 }
153
154 client.close().await;
155
156 Ok(ReplaySummary {
157 statements_replayed,
158 failures,
159 elapsed_ms: start.elapsed().as_millis() as u64,
160 from: req.from,
161 to: req.to,
162 first_error,
163 })
164 }
165}
166
167fn journal_value_to_param(v: &JournalValue) -> ParamValue {
171 match v {
172 JournalValue::Null => ParamValue::Null,
173 JournalValue::Bool(b) => ParamValue::Bool(*b),
174 JournalValue::Int64(i) => ParamValue::Int(*i),
175 JournalValue::Float64(f) => ParamValue::Float(*f),
176 JournalValue::Text(s) => ParamValue::Text(s.clone()),
177 JournalValue::Bytes(b) => {
178 let mut s = String::with_capacity(2 + b.len() * 2);
179 s.push_str("\\x");
180 for byte in b {
181 s.push_str(&format!("{:02x}", byte));
182 }
183 ParamValue::Text(s)
184 }
185 JournalValue::Array(_) => ParamValue::Null,
186 }
187}
188
189#[cfg(test)]
190mod tests {
191 use super::*;
192 use crate::backend::{tls::default_client_config, TlsMode};
193 use crate::NodeId;
194 use std::time::Duration;
195 use uuid::Uuid;
196
197 fn test_template() -> BackendConfig {
198 BackendConfig {
199 host: "placeholder".into(),
200 port: 0,
201 user: "postgres".into(),
202 password: None,
203 database: None,
204 application_name: Some("helios-replay".into()),
205 tls_mode: TlsMode::Disable,
206 connect_timeout: Duration::from_millis(200),
207 query_timeout: Duration::from_millis(200),
208 tls_config: default_client_config(),
209 }
210 }
211
212 #[tokio::test]
213 async fn test_replay_rejects_inverted_window() {
214 let journal = Arc::new(TransactionJournal::new());
215 let engine = ReplayEngine::new(journal, test_template());
216 let now = Utc::now();
217 let req = TimeTravelRequest {
218 from: now,
219 to: now - chrono::Duration::seconds(1),
220 target_host: "127.0.0.1".into(),
221 target_port: 1,
222 target_user: None,
223 target_password: None,
224 target_database: None,
225 };
226 let err = engine.replay_window(&req).await.unwrap_err();
227 assert!(matches!(err, ProxyError::Internal(_)));
228 }
229
230 #[tokio::test]
235 async fn test_replay_empty_window_still_connects() {
236 let journal = Arc::new(TransactionJournal::new());
237 let engine = ReplayEngine::new(journal, test_template());
238 let now = Utc::now();
239 let req = TimeTravelRequest {
240 from: now - chrono::Duration::hours(1),
241 to: now,
242 target_host: "127.0.0.1".into(),
243 target_port: 1, target_user: None,
245 target_password: None,
246 target_database: None,
247 };
248 let err = engine.replay_window(&req).await.unwrap_err();
249 match err {
250 ProxyError::ReplayFailed(msg) => assert!(msg.contains("connect")),
251 other => panic!("expected ReplayFailed, got {:?}", other),
252 }
253 }
254
255 #[tokio::test]
259 async fn test_entries_in_window_filters_correctly() {
260 let journal = Arc::new(TransactionJournal::new());
261 let tx_id = Uuid::new_v4();
262 let session = Uuid::new_v4();
263 let node = NodeId::new();
264
265 let base = Utc::now();
266 journal
267 .begin_transaction(tx_id, session, node, 0)
268 .await
269 .unwrap();
270
271 let _ = base; journal
277 .log_statement(tx_id, "SELECT 1".to_string(), vec![], None, None, 1)
278 .await
279 .unwrap();
280
281 let from = Utc::now() - chrono::Duration::seconds(5);
282 let to = Utc::now() + chrono::Duration::seconds(5);
283 let entries = journal.entries_in_window(from, to).await;
284 assert_eq!(entries.len(), 1, "single in-window entry");
285
286 let far_past_to = Utc::now() - chrono::Duration::hours(1);
287 let far_past_from = far_past_to - chrono::Duration::hours(1);
288 let entries = journal.entries_in_window(far_past_from, far_past_to).await;
289 assert!(entries.is_empty(), "no entries in far-past window");
290 }
291
292 #[test]
293 fn test_journal_value_to_param_matches_failover_shape() {
294 assert!(matches!(
299 journal_value_to_param(&JournalValue::Null),
300 ParamValue::Null
301 ));
302 assert!(matches!(
303 journal_value_to_param(&JournalValue::Bool(true)),
304 ParamValue::Bool(true)
305 ));
306 assert!(matches!(
307 journal_value_to_param(&JournalValue::Int64(-7)),
308 ParamValue::Int(-7)
309 ));
310 }
311
312 #[test]
318 fn test_credential_overrides_replace_template_fields() {
319 let mut cfg = test_template();
320 cfg.user = "default_user".into();
321 cfg.password = None;
322 cfg.database = None;
323
324 let req = TimeTravelRequest {
325 from: Utc::now(),
326 to: Utc::now(),
327 target_host: "h".into(),
328 target_port: 5432,
329 target_user: Some("override_user".into()),
330 target_password: Some("secret".into()),
331 target_database: Some("staging".into()),
332 };
333
334 if let Some(ref u) = req.target_user {
339 cfg.user = u.clone();
340 }
341 if let Some(ref p) = req.target_password {
342 cfg.password = Some(p.clone());
343 }
344 if let Some(ref d) = req.target_database {
345 cfg.database = Some(d.clone());
346 }
347
348 assert_eq!(cfg.user, "override_user");
349 assert_eq!(cfg.password.as_deref(), Some("secret"));
350 assert_eq!(cfg.database.as_deref(), Some("staging"));
351 }
352
353 #[test]
354 fn test_credential_overrides_none_keeps_template_fields() {
355 let mut cfg = test_template();
356 cfg.user = "default_user".into();
357 cfg.password = Some("template_pw".into());
358 cfg.database = Some("default_db".into());
359
360 let req = TimeTravelRequest {
361 from: Utc::now(),
362 to: Utc::now(),
363 target_host: "h".into(),
364 target_port: 5432,
365 target_user: None,
366 target_password: None,
367 target_database: None,
368 };
369
370 if let Some(ref u) = req.target_user {
371 cfg.user = u.clone();
372 }
373 let _ = req;
375
376 assert_eq!(cfg.user, "default_user");
377 assert_eq!(cfg.password.as_deref(), Some("template_pw"));
378 assert_eq!(cfg.database.as_deref(), Some("default_db"));
379 }
380
381 #[test]
384 fn test_replay_summary_serializes() {
385 let s = ReplaySummary {
386 statements_replayed: 5,
387 failures: 1,
388 elapsed_ms: 42,
389 from: Utc::now(),
390 to: Utc::now(),
391 first_error: Some("oops".into()),
392 };
393 let j = serde_json::to_string(&s).unwrap();
394 assert!(j.contains("\"statements_replayed\":5"));
395 assert!(j.contains("\"failures\":1"));
396 assert!(j.contains("oops"));
397 }
398}