1use std::collections::BTreeMap;
83use std::convert::Infallible;
84use std::rc::Rc;
85use std::time::Duration;
86
87use differential_dataflow::AsCollection;
88use itertools::Itertools as _;
89use mz_expr::{EvalError, MirScalarExpr};
90use mz_ore::cast::CastFrom;
91use mz_ore::error::ErrorExt;
92use mz_postgres_util::desc::PostgresTableDesc;
93use mz_postgres_util::{Client, PostgresError, simple_query_opt};
94use mz_repr::{Datum, Diff, GlobalId, Row};
95use mz_sql_parser::ast::Ident;
96use mz_sql_parser::ast::display::AstDisplay;
97use mz_storage_types::errors::{DataflowError, SourceError, SourceErrorDetails};
98use mz_storage_types::sources::postgres::CastType;
99use mz_storage_types::sources::{
100 MzOffset, PostgresSourceConnection, SourceExport, SourceExportDetails, SourceTimestamp,
101};
102use mz_timely_util::builder_async::PressOnDropButton;
103use serde::{Deserialize, Serialize};
104use timely::container::CapacityContainerBuilder;
105use timely::dataflow::operators::core::Partition;
106use timely::dataflow::operators::{Concat, Map, ToStream};
107use timely::dataflow::{Scope, Stream};
108use timely::progress::Antichain;
109use tokio_postgres::error::SqlState;
110use tokio_postgres::types::PgLsn;
111
112use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
113use crate::source::types::{Probe, SourceRender, StackedCollection};
114use crate::source::{RawSourceCreationConfig, SourceMessage};
115
116mod replication;
117mod snapshot;
118
119impl SourceRender for PostgresSourceConnection {
120 type Time = MzOffset;
121
122 const STATUS_NAMESPACE: StatusNamespace = StatusNamespace::Postgres;
123
124 fn render<G: Scope<Timestamp = MzOffset>>(
127 self,
128 scope: &mut G,
129 config: &RawSourceCreationConfig,
130 resume_uppers: impl futures::Stream<Item = Antichain<MzOffset>> + 'static,
131 _start_signal: impl std::future::Future<Output = ()> + 'static,
132 ) -> (
133 BTreeMap<GlobalId, StackedCollection<G, Result<SourceMessage, DataflowError>>>,
134 Stream<G, Infallible>,
135 Stream<G, HealthStatusMessage>,
136 Option<Stream<G, Probe<MzOffset>>>,
137 Vec<PressOnDropButton>,
138 ) {
139 let mut table_info = BTreeMap::new();
141 for (idx, (id, export)) in config.source_exports.iter().enumerate() {
142 let SourceExport {
143 details,
144 storage_metadata: _,
145 data_config: _,
146 } = export;
147 let details = match details {
148 SourceExportDetails::Postgres(details) => details,
149 SourceExportDetails::None => continue,
151 _ => panic!("unexpected source export details: {:?}", details),
152 };
153 let desc = details.table.clone();
154 let casts = details.column_casts.clone();
155 let resume_upper = Antichain::from_iter(
156 config
157 .source_resume_uppers
158 .get(id)
159 .expect("all source exports must be present in source resume uppers")
160 .iter()
161 .map(MzOffset::decode_row),
162 );
163 let output = SourceOutputInfo {
164 desc,
165 casts,
166 resume_upper,
167 export_id: id.clone(),
168 };
169 table_info
170 .entry(output.desc.oid)
171 .or_insert_with(BTreeMap::new)
172 .insert(idx, output);
173 }
174
175 let metrics = config.metrics.get_postgres_source_metrics(config.id);
176
177 let (snapshot_updates, rewinds, slot_ready, snapshot_err, snapshot_token) =
178 snapshot::render(
179 scope.clone(),
180 config.clone(),
181 self.clone(),
182 table_info.clone(),
183 metrics.snapshot_metrics.clone(),
184 );
185
186 let (repl_updates, uppers, probe_stream, repl_err, repl_token) = replication::render(
187 scope.clone(),
188 config.clone(),
189 self,
190 table_info,
191 &rewinds,
192 &slot_ready,
193 resume_uppers,
194 metrics,
195 );
196
197 let updates = snapshot_updates.concat(&repl_updates);
198 let partition_count = u64::cast_from(config.source_exports.len());
199 let data_streams: Vec<_> = updates
200 .inner
201 .partition::<CapacityContainerBuilder<_>, _, _>(
202 partition_count,
203 |((output, data), time, diff): &(
204 (usize, Result<SourceMessage, DataflowError>),
205 MzOffset,
206 Diff,
207 )| {
208 let output = u64::cast_from(*output);
209 (output, (data.clone(), time.clone(), diff.clone()))
210 },
211 );
212 let mut data_collections = BTreeMap::new();
213 for (id, data_stream) in config.source_exports.keys().zip_eq(data_streams) {
214 data_collections.insert(*id, data_stream.as_collection());
215 }
216
217 let init = std::iter::once(HealthStatusMessage {
218 id: None,
219 namespace: Self::STATUS_NAMESPACE,
220 update: HealthStatusUpdate::Running,
221 })
222 .to_stream(scope);
223
224 let errs = snapshot_err.concat(&repl_err).map(move |err| {
228 let err_string = err.display_with_causes().to_string();
230 let update = HealthStatusUpdate::halting(err_string.clone(), None);
231
232 let namespace = match err {
233 ReplicationError::Transient(err)
234 if matches!(
235 &*err,
236 TransientError::PostgresError(PostgresError::Ssh(_))
237 | TransientError::PostgresError(PostgresError::SshIo(_))
238 ) =>
239 {
240 StatusNamespace::Ssh
241 }
242 _ => Self::STATUS_NAMESPACE,
243 };
244
245 HealthStatusMessage {
246 id: None,
247 namespace: namespace.clone(),
248 update,
249 }
250 });
251
252 let health = init.concat(&errs);
253
254 (
255 data_collections,
256 uppers,
257 health,
258 probe_stream,
259 vec![snapshot_token, repl_token],
260 )
261 }
262}
263
264#[derive(Clone, Debug)]
265struct SourceOutputInfo {
266 desc: PostgresTableDesc,
267 casts: Vec<(CastType, MirScalarExpr)>,
268 resume_upper: Antichain<MzOffset>,
269 export_id: GlobalId,
270}
271
272#[derive(Clone, Debug, thiserror::Error)]
273pub enum ReplicationError {
274 #[error(transparent)]
275 Transient(#[from] Rc<TransientError>),
276 #[error(transparent)]
277 Definite(#[from] Rc<DefiniteError>),
278}
279
280#[derive(Debug, thiserror::Error)]
282pub enum TransientError {
283 #[error("replication slot mysteriously missing")]
284 MissingReplicationSlot,
285 #[error(
286 "slot overcompacted. Requested LSN {requested_lsn} but only LSNs >= {available_lsn} are available"
287 )]
288 OvercompactedReplicationSlot {
289 requested_lsn: MzOffset,
290 available_lsn: MzOffset,
291 },
292 #[error("replication slot already exists")]
293 ReplicationSlotAlreadyExists,
294 #[error("stream ended prematurely")]
295 ReplicationEOF,
296 #[error("unexpected replication message")]
297 UnknownReplicationMessage,
298 #[error("unexpected logical replication message")]
299 UnknownLogicalReplicationMessage,
300 #[error("received replication event outside of transaction")]
301 BareTransactionEvent,
302 #[error("lsn mismatch between BEGIN and COMMIT")]
303 InvalidTransaction,
304 #[error("BEGIN within existing BEGIN stream")]
305 NestedTransaction,
306 #[error("recoverable errors should crash the process during snapshots")]
307 SyntheticError,
308 #[error("sql client error")]
309 SQLClient(#[from] tokio_postgres::Error),
310 #[error(transparent)]
311 PostgresError(#[from] PostgresError),
312 #[error(transparent)]
313 Generic(#[from] anyhow::Error),
314}
315
316#[derive(Debug, Clone, Serialize, Deserialize, thiserror::Error)]
318pub enum DefiniteError {
319 #[error("slot compacted past snapshot point. snapshot consistent point={0} resume_lsn={1}")]
320 SlotCompactedPastResumePoint(MzOffset, MzOffset),
321 #[error("table was truncated")]
322 TableTruncated,
323 #[error("table was dropped")]
324 TableDropped,
325 #[error("publication {0:?} does not exist")]
326 PublicationDropped(String),
327 #[error("replication slot has been invalidated because it exceeded the maximum reserved size")]
328 InvalidReplicationSlot,
329 #[error("unexpected number of columns while parsing COPY output")]
330 MissingColumn,
331 #[error("failed to parse COPY protocol")]
332 InvalidCopyInput,
333 #[error(
334 "unsupported action: database restored from point-in-time backup. Expected timeline ID {expected} but got {actual}"
335 )]
336 InvalidTimelineId { expected: u64, actual: u64 },
337 #[error(
338 "TOASTed value missing from old row. Did you forget to set REPLICA IDENTITY to FULL for your table?"
339 )]
340 MissingToast,
341 #[error(
342 "old row missing from replication stream. Did you forget to set REPLICA IDENTITY to FULL for your table?"
343 )]
344 DefaultReplicaIdentity,
345 #[error("incompatible schema change: {0}")]
346 IncompatibleSchema(String),
348 #[error("invalid UTF8 string: {0:?}")]
349 InvalidUTF8(Vec<u8>),
350 #[error("failed to cast raw column: {0}")]
351 CastError(#[source] EvalError),
352 #[error("unexpected binary data in replication stream")]
353 UnexpectedBinaryData,
354}
355
356impl From<DefiniteError> for DataflowError {
357 fn from(err: DefiniteError) -> Self {
358 let m = err.to_string().into();
359 DataflowError::SourceError(Box::new(SourceError {
360 error: match &err {
361 DefiniteError::SlotCompactedPastResumePoint(_, _) => SourceErrorDetails::Other(m),
362 DefiniteError::TableTruncated => SourceErrorDetails::Other(m),
363 DefiniteError::TableDropped => SourceErrorDetails::Other(m),
364 DefiniteError::PublicationDropped(_) => SourceErrorDetails::Initialization(m),
365 DefiniteError::InvalidReplicationSlot => SourceErrorDetails::Initialization(m),
366 DefiniteError::MissingColumn => SourceErrorDetails::Other(m),
367 DefiniteError::InvalidCopyInput => SourceErrorDetails::Other(m),
368 DefiniteError::InvalidTimelineId { .. } => SourceErrorDetails::Initialization(m),
369 DefiniteError::MissingToast => SourceErrorDetails::Other(m),
370 DefiniteError::DefaultReplicaIdentity => SourceErrorDetails::Other(m),
371 DefiniteError::IncompatibleSchema(_) => SourceErrorDetails::Other(m),
372 DefiniteError::InvalidUTF8(_) => SourceErrorDetails::Other(m),
373 DefiniteError::CastError(_) => SourceErrorDetails::Other(m),
374 DefiniteError::UnexpectedBinaryData => SourceErrorDetails::Other(m),
375 },
376 }))
377 }
378}
379
380async fn ensure_replication_slot(client: &Client, slot: &str) -> Result<(), TransientError> {
381 let slot = Ident::new_unchecked(slot).to_ast_string_simple();
383 let query = format!("CREATE_REPLICATION_SLOT {slot} LOGICAL \"pgoutput\" NOEXPORT_SNAPSHOT");
384 match simple_query_opt(client, &query).await {
385 Ok(_) => Ok(()),
386 Err(PostgresError::Postgres(err)) if err.code() == Some(&SqlState::DUPLICATE_OBJECT) => {
388 tracing::trace!("replication slot {slot} already existed");
389 Ok(())
390 }
391 Err(err) => Err(TransientError::PostgresError(err)),
392 }
393}
394
395struct SlotMetadata {
397 active_pid: Option<i32>,
400 confirmed_flush_lsn: MzOffset,
403}
404
405async fn fetch_slot_metadata(
407 client: &Client,
408 slot: &str,
409 interval: Duration,
410) -> Result<SlotMetadata, TransientError> {
411 loop {
412 let query = "SELECT active_pid, confirmed_flush_lsn
413 FROM pg_replication_slots WHERE slot_name = $1";
414 let Some(row) = client.query_opt(query, &[&slot]).await? else {
415 return Err(TransientError::MissingReplicationSlot);
416 };
417
418 match row.get::<_, Option<PgLsn>>("confirmed_flush_lsn") {
419 Some(lsn) => {
423 return Ok(SlotMetadata {
424 confirmed_flush_lsn: MzOffset::from(lsn),
425 active_pid: row.get("active_pid"),
426 });
427 }
428 None => tokio::time::sleep(interval).await,
432 };
433 }
434}
435
436async fn fetch_max_lsn(client: &Client) -> Result<MzOffset, TransientError> {
438 let query = "SELECT pg_current_wal_lsn()";
439 let row = simple_query_opt(client, query).await?;
440
441 match row.and_then(|row| {
442 row.get("pg_current_wal_lsn")
443 .map(|lsn| lsn.parse::<PgLsn>().unwrap())
444 }) {
445 Some(lsn) => Ok(MzOffset::from(lsn)),
450 None => Err(TransientError::Generic(anyhow::anyhow!(
451 "pg_current_wal_lsn() mysteriously has no value"
452 ))),
453 }
454}
455
456fn verify_schema(
459 oid: u32,
460 expected_desc: &PostgresTableDesc,
461 upstream_info: &BTreeMap<u32, PostgresTableDesc>,
462 casts: &[(CastType, MirScalarExpr)],
463) -> Result<(), DefiniteError> {
464 let current_desc = upstream_info.get(&oid).ok_or(DefiniteError::TableDropped)?;
465
466 let allow_oids_to_change_by_col_num = expected_desc
467 .columns
468 .iter()
469 .zip_eq(casts.iter())
470 .flat_map(|(col, (cast_type, _))| match cast_type {
471 CastType::Text => Some(col.col_num),
472 CastType::Natural => None,
473 })
474 .collect();
475
476 match expected_desc.determine_compatibility(current_desc, &allow_oids_to_change_by_col_num) {
477 Ok(()) => Ok(()),
478 Err(err) => Err(DefiniteError::IncompatibleSchema(err.to_string())),
479 }
480}
481
482fn cast_row(
484 casts: &[(CastType, MirScalarExpr)],
485 datums: &[Datum<'_>],
486 row: &mut Row,
487) -> Result<(), DefiniteError> {
488 let arena = mz_repr::RowArena::new();
489 let mut packer = row.packer();
490 for (_, column_cast) in casts {
491 let datum = column_cast
492 .eval(datums, &arena)
493 .map_err(DefiniteError::CastError)?;
494 packer.push(datum);
495 }
496 Ok(())
497}
498
499fn decode_utf8_text(bytes: &[u8]) -> Result<Datum<'_>, DefiniteError> {
501 match std::str::from_utf8(bytes) {
502 Ok(text) => Ok(Datum::String(text)),
503 Err(_) => Err(DefiniteError::InvalidUTF8(bytes.to_vec())),
504 }
505}