Skip to main content

drasi_source_postgres/
stream.rs

1// Copyright 2025 The Drasi Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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    // Note: table_primary_keys is initialized empty and remains so.
87    // Element IDs are generated from configured table_keys (in config.table_keys),
88    // or fall back to using all column values if no keys are configured.
89
90    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        // Connect and setup replication
97        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        // Start streaming loop
108        let mut keepalive_interval = interval(Duration::from_secs(10));
109
110        loop {
111            // Check if we should stop
112            {
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                // Check for replication messages
122                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                                // Attempt recovery
128                                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                            // No message available
136                        }
137                        Err(e) => {
138                            error!("Error reading message: {e}");
139                            // Attempt recovery
140                            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                // Send periodic keepalives
149                _ = 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        // Clean shutdown
158        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        // Create connection
166        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        // Identify system
176        let system_info = conn.identify_system().await?;
177        info!("Connected to PostgreSQL system: {system_info:?}");
178
179        // Create or verify replication slot
180        let slot_info = conn
181            .create_replication_slot(&self.config.slot_name, false)
182            .await?;
183        info!("Using replication slot: {slot_info:?}");
184
185        // Parse starting LSN — use explicit start_lsn if provided (replay),
186        // otherwise use the slot's consistent_point.
187        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        // Build replication options
199        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        // Start replication
207        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            // Try to read with a short timeout to avoid blocking forever
222            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), // Timeout, no message available
228            }
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        // First byte indicates the message type
269        let msg_type = data[0];
270
271        match msg_type {
272            b'w' => {
273                // XLogData message
274                self.handle_xlog_data(&data[1..]).await?;
275            }
276            b'k' => {
277                // Primary keepalive
278                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        // Parse XLogData header
294        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        // Update read LSN
305        self.read_lsn = end_lsn;
306        self.replay_state
307            .read_lsn
308            .store(self.read_lsn, Ordering::Release);
309
310        // Decode the actual WAL message
311        let wal_data = &data[24..];
312
313        // Try to decode, but don't fail on partial messages
314        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                // No message or skipped message type
330            }
331            Err(e) => {
332                // Log but don't fail on decode errors - might be partial message
333                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                // We'll get the rest in the next message
343            }
344        }
345
346        // Send feedback periodically
347        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                // Start a new transaction — track changes with their LSN
380                self.pending_transaction = Some(Vec::new());
381            }
382            WalMessage::Commit(tx_info) => {
383                // Commit the transaction — stamp all changes with the commit LSN
384                // (the atomic ordering key) plus a distinct in-transaction offset
385                // so the per-subscriber high-water-mark filter does not collapse
386                // multiple changes that share the same commit LSN (issue #599).
387                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                // Store relation mapping - use table name as-is for label (no uppercase)
404                // This ensures consistency with bootstrap which uses the actual table name case
405                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                // Update decoder's relation info
416                // The decoder already handles this internally
417            }
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                        // Defensive: pgoutput always wraps changes in Begin/Commit,
424                        // but keep the encoding uniform (16-byte position) if not.
425                        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        // Get relation info
473        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        // Convert tuple to properties
484        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        // Create element ID (use primary key if available, otherwise generate)
498        let element_id = self.generate_element_id(relation, &tuple).await?;
499
500        // Create the element
501        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        // Generate element ID (should be the same for both old and new tuples)
530        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        // Create properties for after state
537        // Note: We allow UPDATE without old_tuple to avoid converting to INSERT.
538        let mut after_properties = drasi_core::models::ElementPropertyMap::new();
539
540        // Process new tuple (after state)
541        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    /// Generate a stable element ID for a tuple based on primary key values.
594    ///
595    /// Priority order:
596    /// 1. User-configured key columns (from table_keys config)
597    /// 2. Automatically detected primary keys from PostgreSQL system catalogs
598    /// 3. UUID fallback if no keys are available
599    ///
600    /// Element ID format:
601    /// - Single key: Table name prefix with value (e.g., "stocks:AAPL")
602    /// - Composite key: Table name prefix with values joined (e.g., "portfolio:tenant1_user2")
603    /// - No key: Table name prefix with UUID (e.g., "orders:550e8400-e29b-41d4-a716-446655440000")
604    async fn generate_element_id(
605        &self,
606        relation: &super::types::RelationInfo,
607        tuple: &[super::types::PostgresValue],
608    ) -> Result<String> {
609        // Get the table name (use schema-qualified if not in public schema)
610        let table_name = if relation.namespace == "public" {
611            relation.name.clone()
612        } else {
613            format!("{}.{}", relation.namespace, relation.name)
614        };
615
616        // Get primary key columns for this table
617        let primary_keys = self.table_primary_keys.read().await;
618        let pk_columns = primary_keys.get(&table_name);
619
620        // Check configured table_keys first
621        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        // Use configured keys if available, otherwise use detected primary keys
629        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                            // Remove quotes from JSON string representation
640                            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                // Include table name as prefix to ensure uniqueness across tables
650                return Ok(format!("{}:{}", table_name, key_parts.join("_")));
651            }
652        }
653
654        // No primary key found or all key values are NULL
655        warn!("No primary key value found for table '{table_name}'. Consider adding 'table_keys' configuration.");
656        // Still include table name prefix for consistency
657        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            // Use the confirmed source position (min LSN across all query
663            // position handles) for flush_lsn / apply_lsn.  This tells
664            // Postgres that WAL up to this LSN has been durably processed
665            // by all subscribers.  Using read_lsn here would advance the
666            // slot watermark past un-checkpointed query positions, causing
667            // PositionUnavailable on crash+restart.
668            //
669            // Positions are 16 bytes (`[ commit_lsn | offset ]`); only the
670            // leading commit_lsn drives WAL feedback. The in-transaction offset
671            // is irrelevant to the slot watermark.
672            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, // No confirmed position yet — don't advance
685            };
686
687            // Apply the flush fence: during the subscription window after a
688            // replay restart, cap flush_lsn to prevent the slot's restart_lsn
689            // from advancing past positions that pending subscribers still need.
690            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            // Prune the sequence→position map up to the confirmed sequence
708            // now that feedback was successfully sent. Skip pruning when the
709            // fence clamped flush_lsn — Postgres hasn't actually acknowledged
710            // the full confirmed position, so we must retain map entries for
711            // accurate re-computation after the fence lifts.
712            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    /// Dispatch a single change event through the framework's `dispatch_event()`.
731    ///
732    /// The `position` parameter is the opaque 16-byte source-position
733    /// (`[ commit_lsn | in-transaction offset ]`) the checkpoint framework
734    /// persists for recovery/replay and the per-subscriber high-water-mark
735    /// filter uses for dedup. Each change in a transaction shares the same
736    /// `commit_lsn` but carries a distinct offset so none are suppressed
737    /// (issue #599).
738    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        // Attach the opaque source_position bytes for checkpoint/recovery and replay dedup
750        wrapper.set_source_position(position);
751
752        // Use dispatch_event() which stamps the monotonic sequence
753        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        // Close existing connection if any
771        if let Some(conn) = self.connection.take() {
772            let _ = conn.close().await;
773        }
774
775        // Wait a bit before reconnecting
776        sleep(Duration::from_secs(5)).await;
777
778        // Try to reconnect
779        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        // Send final feedback
789        let _ = self.send_feedback(false).await;
790
791        // Close connection
792        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    /// Validates that the timestamp pattern used in convert_insert/convert_update/convert_delete
806    /// produces a value in the millisecond range.
807    #[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    /// Verifies that using nanoseconds would be caught by the validator.
817    #[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}