1use super::{NodeId, ProxyError, Result};
7use std::collections::HashMap;
8use std::sync::Arc;
9use tokio::sync::RwLock;
10use uuid::Uuid;
11
12#[derive(Debug, Clone)]
14pub struct JournalEntry {
15 pub sequence: u64,
17 pub statement: String,
19 pub parameters: Vec<JournalValue>,
21 pub result_checksum: Option<u64>,
23 pub rows_affected: Option<u64>,
25 pub timestamp: chrono::DateTime<chrono::Utc>,
27 pub statement_type: StatementType,
29 pub duration_ms: u64,
31}
32
33#[derive(Debug, Clone)]
35pub enum JournalValue {
36 Null,
37 Bool(bool),
38 Int64(i64),
39 Float64(f64),
40 Text(String),
41 Bytes(Vec<u8>),
42 Array(Vec<JournalValue>),
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum StatementType {
48 Select,
50 Insert,
52 Update,
54 Delete,
56 Ddl,
58 Transaction,
60 Set,
62 Other,
64}
65
66impl StatementType {
67 pub fn from_sql(sql: &str) -> Self {
69 let upper = sql.trim().to_uppercase();
70 if upper.starts_with("SELECT") {
71 StatementType::Select
72 } else if upper.starts_with("INSERT") {
73 StatementType::Insert
74 } else if upper.starts_with("UPDATE") {
75 StatementType::Update
76 } else if upper.starts_with("DELETE") {
77 StatementType::Delete
78 } else if upper.starts_with("CREATE")
79 || upper.starts_with("ALTER")
80 || upper.starts_with("DROP")
81 {
82 StatementType::Ddl
83 } else if upper.starts_with("BEGIN")
84 || upper.starts_with("COMMIT")
85 || upper.starts_with("ROLLBACK")
86 || upper.starts_with("SAVEPOINT")
87 {
88 StatementType::Transaction
89 } else if upper.starts_with("SET") {
90 StatementType::Set
91 } else {
92 StatementType::Other
93 }
94 }
95
96 pub fn is_read_only(&self) -> bool {
98 matches!(self, StatementType::Select)
99 }
100
101 pub fn is_mutation(&self) -> bool {
103 matches!(
104 self,
105 StatementType::Insert
106 | StatementType::Update
107 | StatementType::Delete
108 | StatementType::Ddl
109 )
110 }
111}
112
113#[derive(Debug, Clone)]
115pub struct TransactionJournalEntry {
116 pub tx_id: Uuid,
118 pub session_id: Uuid,
120 pub node_id: NodeId,
122 pub started_at: chrono::DateTime<chrono::Utc>,
124 pub start_lsn: u64,
126 pub entries: Vec<JournalEntry>,
128 pub current_sequence: u64,
130 pub active: bool,
132 pub has_mutations: bool,
134 pub savepoints: Vec<Savepoint>,
136}
137
138#[derive(Debug, Clone)]
140pub struct Savepoint {
141 pub name: String,
143 pub sequence: u64,
145 pub created_at: chrono::DateTime<chrono::Utc>,
147}
148
149impl TransactionJournalEntry {
150 pub fn new(tx_id: Uuid, session_id: Uuid, node_id: NodeId, start_lsn: u64) -> Self {
152 Self {
153 tx_id,
154 session_id,
155 node_id,
156 started_at: chrono::Utc::now(),
157 start_lsn,
158 entries: Vec::new(),
159 current_sequence: 0,
160 active: true,
161 has_mutations: false,
162 savepoints: Vec::new(),
163 }
164 }
165
166 pub fn add_entry(&mut self, entry: JournalEntry) {
168 if entry.statement_type.is_mutation() {
169 self.has_mutations = true;
170 }
171 self.current_sequence = entry.sequence;
172 self.entries.push(entry);
173 }
174
175 pub fn create_savepoint(&mut self, name: String) {
177 self.savepoints.push(Savepoint {
178 name,
179 sequence: self.current_sequence,
180 created_at: chrono::Utc::now(),
181 });
182 }
183
184 pub fn rollback_to_savepoint(&mut self, name: &str) -> Option<u64> {
186 if let Some(idx) = self.savepoints.iter().position(|s| s.name == name) {
187 let savepoint = &self.savepoints[idx];
188 let sequence = savepoint.sequence;
189
190 self.entries.retain(|e| e.sequence <= sequence);
192
193 self.savepoints.truncate(idx + 1);
195
196 Some(sequence)
197 } else {
198 None
199 }
200 }
201
202 pub fn entries_for_replay(&self) -> Vec<&JournalEntry> {
204 self.entries.iter().collect()
205 }
206
207 pub fn mutation_entries(&self) -> Vec<&JournalEntry> {
209 self.entries
210 .iter()
211 .filter(|e| e.statement_type.is_mutation())
212 .collect()
213 }
214
215 pub fn total_size(&self) -> usize {
217 self.entries
218 .iter()
219 .map(|e| e.statement.len() + estimate_params_size(&e.parameters))
220 .sum()
221 }
222}
223
224fn estimate_params_size(params: &[JournalValue]) -> usize {
225 params
226 .iter()
227 .map(|p| match p {
228 JournalValue::Null => 1,
229 JournalValue::Bool(_) => 1,
230 JournalValue::Int64(_) => 8,
231 JournalValue::Float64(_) => 8,
232 JournalValue::Text(s) => s.len(),
233 JournalValue::Bytes(b) => b.len(),
234 JournalValue::Array(a) => estimate_params_size(a),
235 })
236 .sum()
237}
238
239pub struct TransactionJournal {
241 journals: Arc<RwLock<HashMap<Uuid, TransactionJournalEntry>>>,
243 max_entries: usize,
245 max_size: usize,
247 max_journals: usize,
254 enabled: bool,
256}
257
258impl TransactionJournal {
259 pub fn new() -> Self {
261 Self {
262 journals: Arc::new(RwLock::new(HashMap::new())),
263 max_entries: 10000,
264 max_size: 64 * 1024 * 1024, max_journals: 50_000,
266 enabled: true,
267 }
268 }
269
270 pub fn with_max_entries(mut self, max: usize) -> Self {
272 self.max_entries = max;
273 self
274 }
275
276 pub fn with_max_journals(mut self, max: usize) -> Self {
278 self.max_journals = max.max(1);
279 self
280 }
281
282 fn evict_oldest_locked(
286 journals: &mut HashMap<Uuid, TransactionJournalEntry>,
287 target_len: usize,
288 ) {
289 if journals.len() <= target_len {
290 return;
291 }
292 let remove_count = journals.len() - target_len;
293 let mut by_time: Vec<(Uuid, chrono::DateTime<chrono::Utc>)> =
294 journals.iter().map(|(k, v)| (*k, v.started_at)).collect();
295 by_time.sort_by_key(|(_, t)| *t);
298 for (k, _) in by_time.into_iter().take(remove_count) {
299 journals.remove(&k);
300 }
301 }
302
303 pub fn with_max_size(mut self, max: usize) -> Self {
305 self.max_size = max;
306 self
307 }
308
309 pub fn set_enabled(&mut self, enabled: bool) {
311 self.enabled = enabled;
312 }
313
314 pub async fn entries_in_window(
324 &self,
325 from: chrono::DateTime<chrono::Utc>,
326 to: chrono::DateTime<chrono::Utc>,
327 ) -> Vec<(Uuid, JournalEntry)> {
328 let journals = self.journals.read().await;
329 let mut out: Vec<(Uuid, JournalEntry)> = Vec::new();
330 for (tx_id, j) in journals.iter() {
331 for entry in &j.entries {
332 if entry.timestamp >= from && entry.timestamp <= to {
333 out.push((*tx_id, entry.clone()));
334 }
335 }
336 }
337 out.sort_by_key(|(_, e)| e.timestamp);
338 out
339 }
340
341 pub async fn begin_transaction(
343 &self,
344 tx_id: Uuid,
345 session_id: Uuid,
346 node_id: NodeId,
347 start_lsn: u64,
348 ) -> Result<()> {
349 if !self.enabled {
350 return Ok(());
351 }
352
353 let journal = TransactionJournalEntry::new(tx_id, session_id, node_id, start_lsn);
354 let mut journals = self.journals.write().await;
355 if journals.len() >= self.max_journals {
359 let target = (self.max_journals * 9 / 10).max(1);
360 Self::evict_oldest_locked(&mut journals, target);
361 }
362 journals.insert(tx_id, journal);
363 drop(journals);
364
365 tracing::debug!("Started journaling transaction {:?}", tx_id);
366 Ok(())
367 }
368
369 pub async fn log_statement(
371 &self,
372 tx_id: Uuid,
373 statement: String,
374 parameters: Vec<JournalValue>,
375 result_checksum: Option<u64>,
376 rows_affected: Option<u64>,
377 duration_ms: u64,
378 ) -> Result<()> {
379 if !self.enabled {
380 return Ok(());
381 }
382
383 let mut journals = self.journals.write().await;
384 let journal = journals.get_mut(&tx_id).ok_or_else(|| {
385 ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
386 })?;
387
388 if journal.entries.len() >= self.max_entries {
390 return Err(ProxyError::Internal(
391 "Transaction journal entries limit exceeded".to_string(),
392 ));
393 }
394
395 if journal.total_size() >= self.max_size {
396 return Err(ProxyError::Internal(
397 "Transaction journal size limit exceeded".to_string(),
398 ));
399 }
400
401 let sequence = journal.current_sequence + 1;
402 let statement_type = StatementType::from_sql(&statement);
403
404 let entry = JournalEntry {
405 sequence,
406 statement,
407 parameters,
408 result_checksum,
409 rows_affected,
410 timestamp: chrono::Utc::now(),
411 statement_type,
412 duration_ms,
413 };
414
415 journal.add_entry(entry);
416
417 Ok(())
418 }
419
420 pub async fn create_savepoint(&self, tx_id: Uuid, name: String) -> Result<()> {
422 if !self.enabled {
423 return Ok(());
424 }
425
426 let mut journals = self.journals.write().await;
427 let journal = journals.get_mut(&tx_id).ok_or_else(|| {
428 ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
429 })?;
430
431 journal.create_savepoint(name);
432 Ok(())
433 }
434
435 pub async fn rollback_to_savepoint(&self, tx_id: Uuid, name: &str) -> Result<()> {
437 if !self.enabled {
438 return Ok(());
439 }
440
441 let mut journals = self.journals.write().await;
442 let journal = journals.get_mut(&tx_id).ok_or_else(|| {
443 ProxyError::Internal(format!("No journal for transaction {:?}", tx_id))
444 })?;
445
446 journal
447 .rollback_to_savepoint(name)
448 .ok_or_else(|| ProxyError::Internal(format!("Savepoint '{}' not found", name)))?;
449
450 Ok(())
451 }
452
453 pub async fn commit_transaction(&self, tx_id: Uuid) -> Result<()> {
455 self.journals.write().await.remove(&tx_id);
456 tracing::debug!("Committed and cleared journal for transaction {:?}", tx_id);
457 Ok(())
458 }
459
460 pub async fn rollback_transaction(&self, tx_id: Uuid) -> Result<()> {
462 self.journals.write().await.remove(&tx_id);
463 tracing::debug!(
464 "Rolled back and cleared journal for transaction {:?}",
465 tx_id
466 );
467 Ok(())
468 }
469
470 pub async fn get_journal(&self, tx_id: &Uuid) -> Option<TransactionJournalEntry> {
472 self.journals.read().await.get(tx_id).cloned()
473 }
474
475 pub async fn active_count(&self) -> usize {
477 self.journals.read().await.len()
478 }
479
480 pub async fn stats(&self) -> JournalStats {
482 let journals = self.journals.read().await;
483 let total_entries: usize = journals.values().map(|j| j.entries.len()).sum();
484 let total_size: usize = journals.values().map(|j| j.total_size()).sum();
485
486 JournalStats {
487 active_transactions: journals.len(),
488 total_entries,
489 total_size_bytes: total_size,
490 enabled: self.enabled,
491 }
492 }
493
494 pub async fn get_all_active(&self) -> Vec<TransactionJournalEntry> {
496 self.journals.read().await.values().cloned().collect()
497 }
498
499 pub async fn get_max_start_lsn(&self) -> Option<u64> {
502 let journals = self.journals.read().await;
503 journals.values().map(|j| j.start_lsn).max()
504 }
505
506 pub async fn get_transactions_for_node(&self, node_id: NodeId) -> Vec<TransactionJournalEntry> {
509 self.journals
510 .read()
511 .await
512 .values()
513 .filter(|j| j.node_id == node_id)
514 .cloned()
515 .collect()
516 }
517}
518
519impl Default for TransactionJournal {
520 fn default() -> Self {
521 Self::new()
522 }
523}
524
525#[derive(Debug, Clone)]
527pub struct JournalStats {
528 pub active_transactions: usize,
530 pub total_entries: usize,
532 pub total_size_bytes: usize,
534 pub enabled: bool,
536}
537
538#[cfg(test)]
539mod tests {
540 use super::*;
541
542 #[test]
543 fn test_statement_type_detection() {
544 assert_eq!(
545 StatementType::from_sql("SELECT * FROM users"),
546 StatementType::Select
547 );
548 assert_eq!(
549 StatementType::from_sql("INSERT INTO users VALUES (1)"),
550 StatementType::Insert
551 );
552 assert_eq!(
553 StatementType::from_sql("UPDATE users SET name = 'x'"),
554 StatementType::Update
555 );
556 assert_eq!(
557 StatementType::from_sql("DELETE FROM users"),
558 StatementType::Delete
559 );
560 assert_eq!(
561 StatementType::from_sql("CREATE TABLE foo (id INT)"),
562 StatementType::Ddl
563 );
564 assert_eq!(StatementType::from_sql("BEGIN"), StatementType::Transaction);
565 assert_eq!(
566 StatementType::from_sql("SET search_path = public"),
567 StatementType::Set
568 );
569 }
570
571 #[test]
572 fn test_statement_type_properties() {
573 assert!(StatementType::Select.is_read_only());
574 assert!(!StatementType::Insert.is_read_only());
575
576 assert!(StatementType::Insert.is_mutation());
577 assert!(StatementType::Update.is_mutation());
578 assert!(!StatementType::Select.is_mutation());
579 }
580
581 #[tokio::test]
582 async fn test_journal_lifecycle() {
583 let journal = TransactionJournal::new();
584 let tx_id = Uuid::new_v4();
585 let session_id = Uuid::new_v4();
586 let node_id = NodeId::new();
587
588 journal
590 .begin_transaction(tx_id, session_id, node_id, 0)
591 .await
592 .unwrap();
593
594 journal
596 .log_statement(
597 tx_id,
598 "SELECT * FROM users".to_string(),
599 vec![],
600 Some(12345),
601 None,
602 10,
603 )
604 .await
605 .unwrap();
606
607 journal
608 .log_statement(
609 tx_id,
610 "INSERT INTO users (name) VALUES ($1)".to_string(),
611 vec![JournalValue::Text("test".to_string())],
612 None,
613 Some(1),
614 5,
615 )
616 .await
617 .unwrap();
618
619 let j = journal.get_journal(&tx_id).await.unwrap();
621 assert_eq!(j.entries.len(), 2);
622 assert!(j.has_mutations);
623
624 journal.commit_transaction(tx_id).await.unwrap();
626 assert!(journal.get_journal(&tx_id).await.is_none());
627 }
628
629 #[tokio::test]
630 async fn test_savepoints() {
631 let journal = TransactionJournal::new();
632 let tx_id = Uuid::new_v4();
633 let session_id = Uuid::new_v4();
634 let node_id = NodeId::new();
635
636 journal
637 .begin_transaction(tx_id, session_id, node_id, 0)
638 .await
639 .unwrap();
640
641 for i in 0..3 {
643 journal
644 .log_statement(
645 tx_id,
646 format!("INSERT INTO t VALUES ({})", i),
647 vec![],
648 None,
649 Some(1),
650 1,
651 )
652 .await
653 .unwrap();
654 }
655
656 journal
658 .create_savepoint(tx_id, "sp1".to_string())
659 .await
660 .unwrap();
661
662 for i in 3..5 {
664 journal
665 .log_statement(
666 tx_id,
667 format!("INSERT INTO t VALUES ({})", i),
668 vec![],
669 None,
670 Some(1),
671 1,
672 )
673 .await
674 .unwrap();
675 }
676
677 let j = journal.get_journal(&tx_id).await.unwrap();
678 assert_eq!(j.entries.len(), 5);
679
680 journal.rollback_to_savepoint(tx_id, "sp1").await.unwrap();
682
683 let j = journal.get_journal(&tx_id).await.unwrap();
684 assert_eq!(j.entries.len(), 3);
685 }
686
687 #[tokio::test]
688 async fn test_stats() {
689 let journal = TransactionJournal::new();
690 let tx_id = Uuid::new_v4();
691 let session_id = Uuid::new_v4();
692 let node_id = NodeId::new();
693
694 journal
695 .begin_transaction(tx_id, session_id, node_id, 0)
696 .await
697 .unwrap();
698 journal
699 .log_statement(tx_id, "SELECT 1".to_string(), vec![], None, None, 1)
700 .await
701 .unwrap();
702
703 let stats = journal.stats().await;
704 assert_eq!(stats.active_transactions, 1);
705 assert_eq!(stats.total_entries, 1);
706 assert!(stats.enabled);
707 }
708
709 #[tokio::test]
713 async fn global_cap_evicts_oldest_journals() {
714 let journal = TransactionJournal::new().with_max_journals(10);
715 let node_id = NodeId::new();
716 for _ in 0..100 {
719 let tx = Uuid::new_v4();
720 journal
721 .begin_transaction(tx, Uuid::new_v4(), node_id, 0)
722 .await
723 .unwrap();
724 journal
725 .log_statement(
726 tx,
727 "INSERT INTO t VALUES (1)".to_string(),
728 vec![],
729 None,
730 None,
731 1,
732 )
733 .await
734 .unwrap();
735 }
736 let stats = journal.stats().await;
738 assert!(
739 stats.active_transactions <= 10,
740 "journal map must stay within the cap, got {}",
741 stats.active_transactions
742 );
743 }
744}