1use anyhow::{anyhow, Result};
16use chrono::Utc;
17use log::{debug, error, info, trace, warn};
18use std::collections::HashMap;
19use std::sync::atomic::Ordering;
20use std::sync::Arc;
21use std::time::Duration;
22use tokio::sync::{oneshot, RwLock};
23use tokio::time::{interval, sleep};
24
25use super::connection::{parse_lsn, ReplicationConnection};
26use super::decoder::PgOutputDecoder;
27use super::protocol::BackendMessage;
28use super::types::{StandbyStatusUpdate, WalMessage};
29use super::{PostgresSourceConfig, ReplayState};
30use drasi_core::models::{Element, ElementMetadata, ElementReference, SourceChange};
31use drasi_lib::channels::{ComponentStatus, SourceEvent, SourceEventWrapper};
32use drasi_lib::component_graph::ComponentStatusHandle;
33use drasi_lib::sources::base::SourceBase;
34
35pub struct ReplicationStream {
36 config: PostgresSourceConfig,
37 source_id: String,
38 connection: Option<ReplicationConnection>,
39 decoder: PgOutputDecoder,
40 #[allow(dead_code)]
41 status_handle: ComponentStatusHandle,
42 base: SourceBase,
43 replay_state: Arc<ReplayState>,
44 read_lsn: u64,
45 start_lsn: Option<u64>,
46 last_feedback_time: std::time::Instant,
47 pending_transaction: Option<Vec<(SourceChange, u64)>>,
48 relations: HashMap<u32, RelationMapping>,
49 table_primary_keys: Arc<RwLock<HashMap<String, Vec<String>>>>,
50}
51
52struct RelationMapping {
53 #[allow(dead_code)]
54 table_name: String,
55 #[allow(dead_code)]
56 schema_name: String,
57 label: String,
58}
59
60impl ReplicationStream {
61 pub(crate) fn new(
62 config: PostgresSourceConfig,
63 source_id: String,
64 status_handle: ComponentStatusHandle,
65 base: SourceBase,
66 replay_state: Arc<ReplayState>,
67 start_lsn: Option<u64>,
68 ) -> Self {
69 Self {
70 config,
71 source_id,
72 connection: None,
73 decoder: PgOutputDecoder::new(),
74 status_handle,
75 base,
76 replay_state,
77 read_lsn: 0,
78 start_lsn,
79 last_feedback_time: std::time::Instant::now(),
80 pending_transaction: None,
81 relations: HashMap::new(),
82 table_primary_keys: Arc::new(RwLock::new(HashMap::new())),
83 }
84 }
85
86 pub async fn run(
91 &mut self,
92 startup_tx: Option<oneshot::Sender<std::result::Result<(), String>>>,
93 ) -> Result<()> {
94 info!("Starting replication stream for source {}", self.source_id);
95
96 if let Err(error) = self.connect_and_setup().await {
98 if let Some(tx) = startup_tx {
99 let _ = tx.send(Err(format!("{error:#}")));
100 }
101 return Err(error);
102 }
103 if let Some(tx) = startup_tx {
104 let _ = tx.send(Ok(()));
105 }
106
107 let mut keepalive_interval = interval(Duration::from_secs(10));
109
110 loop {
111 {
113 let status = self.status_handle.get_status().await;
114 if status == ComponentStatus::Stopping || status == ComponentStatus::Stopped {
115 info!("Received stop signal, shutting down replication");
116 break;
117 }
118 }
119
120 tokio::select! {
121 result = self.read_next_message() => {
123 match result {
124 Ok(Some(msg)) => {
125 if let Err(e) = self.handle_message(msg).await {
126 error!("Error handling message: {e}");
127 if let Err(e) = self.recover_connection().await {
129 error!("Failed to recover connection: {e}");
130 return Err(e);
131 }
132 }
133 }
134 Ok(None) => {
135 }
137 Err(e) => {
138 error!("Error reading message: {e}");
139 if let Err(e) = self.recover_connection().await {
141 error!("Failed to recover connection: {e}");
142 return Err(e);
143 }
144 }
145 }
146 }
147
148 _ = keepalive_interval.tick() => {
150 if let Err(e) = self.send_feedback(false).await {
151 warn!("Failed to send keepalive: {e}");
152 }
153 }
154 }
155 }
156
157 self.shutdown().await?;
159 Ok(())
160 }
161
162 async fn connect_and_setup(&mut self) -> Result<()> {
163 info!("Connecting to PostgreSQL for replication");
164
165 let mut conn = ReplicationConnection::connect(
167 &self.config.host,
168 self.config.port,
169 &self.config.database,
170 &self.config.user,
171 &self.config.password,
172 )
173 .await?;
174
175 let system_info = conn.identify_system().await?;
177 info!("Connected to PostgreSQL system: {system_info:?}");
178
179 let slot_info = conn
181 .create_replication_slot(&self.config.slot_name, false)
182 .await?;
183 info!("Using replication slot: {slot_info:?}");
184
185 let slot_lsn =
188 if !slot_info.consistent_point.is_empty() && slot_info.consistent_point != "0/0" {
189 parse_lsn(&slot_info.consistent_point)?
190 } else {
191 0
192 };
193 self.read_lsn = self.start_lsn.unwrap_or(slot_lsn);
194 self.replay_state
195 .read_lsn
196 .store(self.read_lsn, Ordering::Release);
197
198 let mut options = HashMap::new();
200 options.insert("proto_version".to_string(), "1".to_string());
201 options.insert(
202 "publication_names".to_string(),
203 self.config.publication_name.clone(),
204 );
205
206 conn.start_replication(&self.config.slot_name, Some(self.read_lsn), options)
208 .await?;
209
210 self.connection = Some(conn);
211 info!(
212 "Replication started from read LSN {:x} (slot watermark {:x})",
213 self.read_lsn, slot_lsn
214 );
215
216 Ok(())
217 }
218
219 async fn read_next_message(&mut self) -> Result<Option<BackendMessage>> {
220 if let Some(conn) = &mut self.connection {
221 match tokio::time::timeout(Duration::from_millis(100), conn.read_replication_message())
223 .await
224 {
225 Ok(Ok(msg)) => Ok(Some(msg)),
226 Ok(Err(e)) => Err(e),
227 Err(_) => Ok(None), }
229 } else {
230 Err(anyhow!("No connection available"))
231 }
232 }
233
234 async fn handle_message(&mut self, msg: BackendMessage) -> Result<()> {
235 match msg {
236 BackendMessage::CopyData(data) => {
237 self.handle_copy_data(&data).await?;
238 }
239 BackendMessage::PrimaryKeepaliveMessage {
240 wal_end,
241 timestamp: _,
242 reply,
243 } => {
244 self.read_lsn = wal_end;
245 self.replay_state
246 .read_lsn
247 .store(self.read_lsn, Ordering::Release);
248 if reply == 1 {
249 self.send_feedback(true).await?;
250 }
251 }
252 BackendMessage::ErrorResponse(err) => {
253 error!("Server error: {}", err.message);
254 return Err(anyhow!("Server error: {}", err.message));
255 }
256 _ => {
257 trace!("Ignoring message: {msg:?}");
258 }
259 }
260 Ok(())
261 }
262
263 async fn handle_copy_data(&mut self, data: &[u8]) -> Result<()> {
264 if data.is_empty() {
265 return Ok(());
266 }
267
268 let msg_type = data[0];
270
271 match msg_type {
272 b'w' => {
273 self.handle_xlog_data(&data[1..]).await?;
275 }
276 b'k' => {
277 self.handle_keepalive(&data[1..]).await?;
279 }
280 _ => {
281 warn!("Unknown copy data message type: 0x{msg_type:02x}");
282 }
283 }
284
285 Ok(())
286 }
287
288 async fn handle_xlog_data(&mut self, data: &[u8]) -> Result<()> {
289 if data.len() < 24 {
290 return Err(anyhow!("XLogData message too short: {} bytes", data.len()));
291 }
292
293 let _start_lsn = u64::from_be_bytes([
295 data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
296 ]);
297 let end_lsn = u64::from_be_bytes([
298 data[8], data[9], data[10], data[11], data[12], data[13], data[14], data[15],
299 ]);
300 let _timestamp = i64::from_be_bytes([
301 data[16], data[17], data[18], data[19], data[20], data[21], data[22], data[23],
302 ]);
303
304 self.read_lsn = end_lsn;
306 self.replay_state
307 .read_lsn
308 .store(self.read_lsn, Ordering::Release);
309
310 let wal_data = &data[24..];
312
313 if !wal_data.is_empty() {
315 let msg_type = wal_data[0];
316 debug!(
317 "Attempting to decode WAL message type: {} ({}), data length: {}",
318 msg_type as char,
319 msg_type,
320 wal_data.len()
321 );
322 }
323
324 match self.decoder.decode_message(wal_data) {
325 Ok(Some(wal_msg)) => {
326 self.process_wal_message(wal_msg).await?;
327 }
328 Ok(None) => {
329 }
331 Err(e) => {
332 if !wal_data.is_empty() {
334 debug!(
335 "Failed to decode WAL message type {} ({}): {}, data length: {}",
336 wal_data[0] as char,
337 wal_data[0],
338 e,
339 wal_data.len()
340 );
341 }
342 }
344 }
345
346 if self.last_feedback_time.elapsed() > Duration::from_secs(5) {
348 self.send_feedback(false).await?;
349 }
350
351 Ok(())
352 }
353
354 async fn handle_keepalive(&mut self, data: &[u8]) -> Result<()> {
355 if data.len() < 17 {
356 return Err(anyhow!("Keepalive message too short"));
357 }
358
359 let wal_end = u64::from_be_bytes([
360 data[0], data[1], data[2], data[3], data[4], data[5], data[6], data[7],
361 ]);
362 let reply = data[16];
363
364 self.read_lsn = wal_end;
365 self.replay_state
366 .read_lsn
367 .store(self.read_lsn, Ordering::Release);
368
369 if reply == 1 {
370 self.send_feedback(true).await?;
371 }
372
373 Ok(())
374 }
375
376 async fn process_wal_message(&mut self, msg: WalMessage) -> Result<()> {
377 match msg {
378 WalMessage::Begin(_) => {
379 self.pending_transaction = Some(Vec::new());
381 }
382 WalMessage::Commit(tx_info) => {
383 if let Some(changes) = self.pending_transaction.take() {
388 let change_count = changes.len();
389 for (offset, (change, _)) in changes.into_iter().enumerate() {
390 let position = super::connection::commit_position_bytes(
391 tx_info.commit_lsn,
392 offset as u64,
393 );
394 self.dispatch_change(change, position).await;
395 }
396 debug!(
397 "Committed transaction {} with LSN {:x} ({} change(s))",
398 tx_info.xid, tx_info.commit_lsn, change_count
399 );
400 }
401 }
402 WalMessage::Relation(relation) => {
403 let label = relation.name.clone();
406 self.relations.insert(
407 relation.id,
408 RelationMapping {
409 table_name: relation.name.clone(),
410 schema_name: relation.namespace.clone(),
411 label,
412 },
413 );
414
415 }
418 WalMessage::Insert { relation_id, tuple } => {
419 if let Some(change) = self.convert_insert(relation_id, tuple).await? {
420 if let Some(tx) = &mut self.pending_transaction {
421 tx.push((change, self.read_lsn));
422 } else {
423 let position = super::connection::commit_position_bytes(self.read_lsn, 0);
426 self.dispatch_change(change, position).await;
427 }
428 }
429 }
430 WalMessage::Update {
431 relation_id,
432 old_tuple,
433 new_tuple,
434 } => {
435 if let Some(change) = self
436 .convert_update(relation_id, old_tuple, new_tuple)
437 .await?
438 {
439 if let Some(tx) = &mut self.pending_transaction {
440 tx.push((change, self.read_lsn));
441 } else {
442 let position = super::connection::commit_position_bytes(self.read_lsn, 0);
443 self.dispatch_change(change, position).await;
444 }
445 }
446 }
447 WalMessage::Delete {
448 relation_id,
449 old_tuple,
450 } => {
451 if let Some(change) = self.convert_delete(relation_id, old_tuple).await? {
452 if let Some(tx) = &mut self.pending_transaction {
453 tx.push((change, self.read_lsn));
454 } else {
455 let position = super::connection::commit_position_bytes(self.read_lsn, 0);
456 self.dispatch_change(change, position).await;
457 }
458 }
459 }
460 WalMessage::Truncate { relation_ids } => {
461 warn!("Truncate not yet implemented for relations: {relation_ids:?}");
462 }
463 }
464 Ok(())
465 }
466
467 async fn convert_insert(
468 &self,
469 relation_id: u32,
470 tuple: Vec<super::types::PostgresValue>,
471 ) -> Result<Option<SourceChange>> {
472 let relation = self
474 .decoder
475 .get_relation(relation_id)
476 .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;
477
478 let mapping = self
479 .relations
480 .get(&relation_id)
481 .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;
482
483 let mut properties = drasi_core::models::ElementPropertyMap::new();
485 for (i, value) in tuple.iter().enumerate() {
486 if let Some(column) = relation.columns.get(i) {
487 let json_value = value.to_json();
488 if !json_value.is_null() {
489 properties.insert(
490 &column.name,
491 drasi_lib::sources::manager::convert_json_to_element_value(&json_value),
492 );
493 }
494 }
495 }
496
497 let element_id = self.generate_element_id(relation, &tuple).await?;
499
500 let element = Element::Node {
502 metadata: ElementMetadata {
503 reference: ElementReference::new(&self.source_id, &element_id),
504 labels: Arc::from([Arc::from(mapping.label.as_str())]),
505 effective_from: Utc::now().timestamp_millis() as u64,
506 },
507 properties,
508 };
509
510 Ok(Some(SourceChange::Insert { element }))
511 }
512
513 async fn convert_update(
514 &self,
515 relation_id: u32,
516 old_tuple: Option<Vec<super::types::PostgresValue>>,
517 new_tuple: Vec<super::types::PostgresValue>,
518 ) -> Result<Option<SourceChange>> {
519 let relation = self
520 .decoder
521 .get_relation(relation_id)
522 .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;
523
524 let mapping = self
525 .relations
526 .get(&relation_id)
527 .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;
528
529 let element_id = self.generate_element_id(relation, &new_tuple).await?;
531
532 if old_tuple.is_none() {
533 warn!("UPDATE without old tuple for relation {relation_id}, preserving UPDATE");
534 }
535
536 let mut after_properties = drasi_core::models::ElementPropertyMap::new();
539
540 for (i, column) in relation.columns.iter().enumerate() {
542 if let Some(value) = new_tuple.get(i) {
543 let json_value = value.to_json();
544 if !json_value.is_null() {
545 after_properties.insert(
546 &column.name,
547 drasi_lib::sources::manager::convert_json_to_element_value(&json_value),
548 );
549 }
550 }
551 }
552
553 let after_element = Element::Node {
554 metadata: ElementMetadata {
555 reference: ElementReference::new(&self.source_id, &element_id),
556 labels: Arc::from([Arc::from(mapping.label.as_str())]),
557 effective_from: Utc::now().timestamp_millis() as u64,
558 },
559 properties: after_properties,
560 };
561
562 Ok(Some(SourceChange::Update {
563 element: after_element,
564 }))
565 }
566
567 async fn convert_delete(
568 &self,
569 relation_id: u32,
570 old_tuple: Vec<super::types::PostgresValue>,
571 ) -> Result<Option<SourceChange>> {
572 let relation = self
573 .decoder
574 .get_relation(relation_id)
575 .ok_or_else(|| anyhow!("Unknown relation {relation_id}"))?;
576
577 let mapping = self
578 .relations
579 .get(&relation_id)
580 .ok_or_else(|| anyhow!("No mapping for relation {relation_id}"))?;
581
582 let element_id = self.generate_element_id(relation, &old_tuple).await?;
583
584 Ok(Some(SourceChange::Delete {
585 metadata: ElementMetadata {
586 reference: ElementReference::new(&self.source_id, &element_id),
587 labels: Arc::from([Arc::from(mapping.label.as_str())]),
588 effective_from: Utc::now().timestamp_millis() as u64,
589 },
590 }))
591 }
592
593 async fn generate_element_id(
605 &self,
606 relation: &super::types::RelationInfo,
607 tuple: &[super::types::PostgresValue],
608 ) -> Result<String> {
609 let table_name = if relation.namespace == "public" {
611 relation.name.clone()
612 } else {
613 format!("{}.{}", relation.namespace, relation.name)
614 };
615
616 let primary_keys = self.table_primary_keys.read().await;
618 let pk_columns = primary_keys.get(&table_name);
619
620 let configured_keys = self
622 .config
623 .table_keys
624 .iter()
625 .find(|tk| tk.table == table_name)
626 .map(|tk| &tk.key_columns);
627
628 let key_columns = configured_keys.or(pk_columns);
630
631 if let Some(keys) = key_columns {
632 let mut key_parts = Vec::new();
633
634 for (i, column) in relation.columns.iter().enumerate() {
635 if keys.contains(&column.name) {
636 if let Some(value) = tuple.get(i) {
637 let json_val = value.to_json();
638 if !json_val.is_null() {
639 let val_str = json_val.to_string();
641 let cleaned = val_str.trim_matches('"');
642 key_parts.push(cleaned.to_string());
643 }
644 }
645 }
646 }
647
648 if !key_parts.is_empty() {
649 return Ok(format!("{}:{}", table_name, key_parts.join("_")));
651 }
652 }
653
654 warn!("No primary key value found for table '{table_name}'. Consider adding 'table_keys' configuration.");
656 Ok(format!("{}:{}", table_name, uuid::Uuid::new_v4()))
658 }
659
660 async fn send_feedback(&mut self, reply_requested: bool) -> Result<()> {
661 if let Some(conn) = &mut self.connection {
662 let confirmed_lsn = match self.base.compute_confirmed_source_position().await {
673 Some(bytes) => match super::connection::position_bytes_to_lsn(&bytes) {
674 Ok(lsn) => lsn,
675 Err(e) => {
676 warn!(
677 "[{}] Confirmed source position could not be decoded ({}); \
678 not advancing flush_lsn",
679 self.source_id, e
680 );
681 0
682 }
683 },
684 None => 0, };
686
687 let fence = self.replay_state.effective_flush_fence();
691 let (effective_lsn, was_clamped) = if fence < u64::MAX && confirmed_lsn > fence {
692 (fence, true)
693 } else {
694 (confirmed_lsn, false)
695 };
696
697 let status = StandbyStatusUpdate {
698 write_lsn: self.read_lsn,
699 flush_lsn: effective_lsn,
700 apply_lsn: effective_lsn,
701 reply_requested,
702 };
703
704 conn.send_standby_status(status).await?;
705 self.last_feedback_time = std::time::Instant::now();
706
707 if !was_clamped && effective_lsn > 0 {
713 if let Some(confirmed_seq) = self.base.compute_confirmed_position().await {
714 self.base.prune_position_map(confirmed_seq).await;
715 }
716 }
717
718 trace!(
719 "[{}] Sent feedback: write_lsn={:x}, flush_lsn={:x}{}",
720 self.source_id,
721 self.read_lsn,
722 effective_lsn,
723 if was_clamped { " (fenced)" } else { "" }
724 );
725 }
726
727 Ok(())
728 }
729
730 async fn dispatch_change(&self, change: SourceChange, position: bytes::Bytes) {
739 let mut profiling = drasi_lib::profiling::ProfilingMetadata::new();
740 profiling.source_send_ns = Some(drasi_lib::profiling::timestamp_ns());
741
742 let mut wrapper = SourceEventWrapper::with_profiling(
743 self.source_id.clone(),
744 SourceEvent::Change(change),
745 chrono::Utc::now(),
746 profiling,
747 );
748
749 wrapper.set_source_position(position);
751
752 if let Err(e) = self.base.dispatch_event(wrapper).await {
754 debug!(
755 "[{}] Failed to dispatch change (no subscribers): {}",
756 self.source_id, e
757 );
758 }
759 }
760
761 #[allow(dead_code)]
762 async fn check_stop_signal(&self) -> bool {
763 let status = self.status_handle.get_status().await;
764 status == ComponentStatus::Stopping || status == ComponentStatus::Stopped
765 }
766
767 async fn recover_connection(&mut self) -> Result<()> {
768 warn!("Attempting to recover connection");
769
770 if let Some(conn) = self.connection.take() {
772 let _ = conn.close().await;
773 }
774
775 sleep(Duration::from_secs(5)).await;
777
778 self.connect_and_setup().await?;
780
781 info!("Connection recovered successfully");
782 Ok(())
783 }
784
785 async fn shutdown(&mut self) -> Result<()> {
786 info!("Shutting down replication stream");
787
788 let _ = self.send_feedback(false).await;
790
791 if let Some(conn) = self.connection.take() {
793 conn.close().await?;
794 }
795
796 Ok(())
797 }
798}
799
800#[cfg(test)]
801mod tests {
802 use chrono::Utc;
803 use drasi_core::models::validate_effective_from;
804
805 #[test]
808 fn effective_from_uses_milliseconds() {
809 let effective_from = Utc::now().timestamp_millis() as u64;
810 assert!(
811 validate_effective_from(effective_from).is_ok(),
812 "Postgres CDC effective_from ({effective_from}) should be in millisecond range"
813 );
814 }
815
816 #[test]
818 fn effective_from_rejects_nanoseconds_pattern() {
819 let bad_effective_from = Utc::now().timestamp_nanos_opt().unwrap() as u64;
820 assert!(
821 validate_effective_from(bad_effective_from).is_err(),
822 "Nanosecond timestamp ({bad_effective_from}) should be rejected"
823 );
824 }
825}