storage: use synthetic probes for Yugabyte sources

MaterializeInc · Nov 8, 2024 · b73a85e · b73a85e
1 parent 5acbe52
commit b73a85e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/storage/src/source/postgres.rs b/src/storage/src/source/postgres.rs
@@ -245,7 +245,7 @@ impl SourceRender for PostgresSourceConnection {
             Some(uppers),
             health,
             stats_stream,
-            Some(probe_stream),
+            probe_stream,
             vec![snapshot_token, repl_token],
         )
     }

diff --git a/src/storage/src/source/postgres/replication.rs b/src/storage/src/source/postgres/replication.rs
@@ -85,6 +85,7 @@ use mz_ore::cast::CastFrom;
 use mz_ore::collections::HashSet;
 use mz_ore::future::InTask;
 use mz_ore::iter::IteratorExt;
+use mz_postgres_util::tunnel::PostgresFlavor;
 use mz_postgres_util::{simple_query_opt, Client};
 use mz_repr::{Datum, DatumVec, Diff, Row};
 use mz_sql_parser::ast::{display::AstDisplay, Ident};
@@ -150,7 +151,7 @@ pub(crate) fn render<G: Scope<Timestamp = MzOffset>>(
     StackedCollection<G, (usize, Result<SourceMessage, DataflowError>)>,
     Stream<G, Infallible>,
     Stream<G, ProgressStatisticsUpdate>,
-    Stream<G, Probe<MzOffset>>,
+    Option<Stream<G, Probe<MzOffset>>>,
     Stream<G, ReplicationError>,
     PressOnDropButton,
 ) {
@@ -165,6 +166,12 @@ pub(crate) fn render<G: Scope<Timestamp = MzOffset>>(
     let (stats_output, stats_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
     let (probe_output, probe_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
 
+    // Yugabyte doesn't support LSN probing currently.
+    let probe_stream = match connection.connection.flavor {
+        PostgresFlavor::Vanilla => Some(probe_stream),
+        PostgresFlavor::Yugabyte => None,
+    };
+
     let mut rewind_input = builder.new_input_for(
         rewind_stream,
         Exchange::new(move |_| slot_reader),