mz_storage/source/
postgres.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Code to render the ingestion dataflow of a [`PostgresSourceConnection`]. The dataflow consists
11//! of multiple operators in order to take advantage of all the available workers.
12//!
13//! # Snapshot
14//!
15//! One part of the dataflow deals with snapshotting the tables involved in the ingestion. Each
16//! table that needs a snapshot is assigned to a specific worker which performs a `COPY` query
17//! and distributes the raw COPY bytes to all workers to decode the text encoded rows.
18//!
19//! For all tables that ended up being snapshotted the snapshot reader also emits a rewind request
20//! to the replication reader which will ensure that the requested portion of the replication
21//! stream is subtracted from the snapshot.
22//!
23//! See the [snapshot] module for more information on the snapshot strategy.
24//!
25//! # Replication
26//!
27//! The other part of the dataflow deals with reading the logical replication slot, which must
28//! happen from a single worker. The minimum amount of processing is performed from that worker
29//! and the data is then distributed among all workers for decoding.
30//!
31//! See the [replication] module for more information on the replication strategy.
32//!
33//! # Error handling
34//!
35//! There are two kinds of errors that can happen during ingestion that are represented as two
36//! separate error types:
37//!
38//! [`DefiniteError`]s are errors that happen during processing of a specific
39//! collection record at a specific LSN. These are the only errors that can ever end up in the
40//! error collection of a subsource.
41//!
42//! Transient errors are any errors that can happen for reasons that are unrelated to the data
43//! itself. This could be authentication failures, connection failures, etc. The only operators
44//! that can emit such errors are the `TableReader` and the `ReplicationReader` operators, which
45//! are the ones that talk to the external world. Both of these operators are built with the
46//! `AsyncOperatorBuilder::build_fallible` method which allows transient errors to be propagated
47//! upwards with the standard `?` operator without risking downgrading the capability and producing
48//! bogus frontiers.
49//!
50//! The error streams from both of those operators are published to the source status and also
51//! trigger a restart of the dataflow.
52//!
53//! ```text
54//!    ┏━━━━━━━━━━━━━━┓
55//!    ┃    table     ┃
56//!    ┃    reader    ┃
57//!    ┗━┯━━━━━━━━━━┯━┛
58//!      │          │rewind
59//!      │          │requests
60//!      │          ╰────╮
61//!      │             ┏━v━━━━━━━━━━━┓
62//!      │             ┃ replication ┃
63//!      │             ┃   reader    ┃
64//!      │             ┗━┯━━━━━━━━━┯━┛
65//!  COPY│           slot│         │
66//!  data│           data│         │
67//! ┏━━━━v━━━━━┓ ┏━━━━━━━v━━━━━┓   │
68//! ┃  COPY    ┃ ┃ replication ┃   │
69//! ┃ decoder  ┃ ┃   decoder   ┃   │
70//! ┗━━━━┯━━━━━┛ ┗━━━━━┯━━━━━━━┛   │
71//!      │snapshot     │replication│
72//!      │updates      │updates    │
73//!      ╰────╮    ╭───╯           │
74//!          ╭┴────┴╮              │
75//!          │concat│              │
76//!          ╰──┬───╯              │
77//!             │ data             │progress
78//!             │ output           │output
79//!             v                  v
80//! ```
81
82use 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    /// Render the ingestion dataflow. This function only connects things together and contains no
125    /// actual processing logic.
126    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        // Collect the source outputs that we will be exporting into a per-table map.
140        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                // This is an export that doesn't need any data output to it.
150                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        // N.B. Note that we don't check ssh tunnel statuses here. We could, but immediately on
225        // restart we are going to set the status to an ssh error correctly, so we don't do this
226        // extra work.
227        let errs = snapshot_err.concat(&repl_err).map(move |err| {
228            // This update will cause the dataflow to restart
229            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/// A transient error that never ends up in the collection of a specific table.
281#[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/// A definite error that always ends up in the collection of a specific table.
317#[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    // TODO: proper error variants for all the expected schema violations
347    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    // Note: Using unchecked here is okay because we're using it in a SQL query.
382    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        // If the slot already exists that's still ok
387        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
395/// The state of a replication slot.
396struct SlotMetadata {
397    /// The process ID of the session using this slot if the slot is currently actively being used.
398    /// None if inactive.
399    active_pid: Option<i32>,
400    /// The address (LSN) up to which the logical slot's consumer has confirmed receiving data.
401    /// Data corresponding to the transactions committed before this LSN is not available anymore.
402    confirmed_flush_lsn: MzOffset,
403}
404
405/// Fetches the minimum LSN at which this slot can safely resume.
406async 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            // For postgres, `confirmed_flush_lsn` means that the slot is able to produce
420            // all transactions that happen at tx_lsn >= confirmed_flush_lsn. Therefore this value
421            // already has "upper" semantics.
422            Some(lsn) => {
423                return Ok(SlotMetadata {
424                    confirmed_flush_lsn: MzOffset::from(lsn),
425                    active_pid: row.get("active_pid"),
426                });
427            }
428            // It can happen that confirmed_flush_lsn is NULL as the slot initializes
429            // This could probably be a `tokio::time::interval`, but its only is called twice,
430            // so its fine like this.
431            None => tokio::time::sleep(interval).await,
432        };
433    }
434}
435
436/// Fetch the `pg_current_wal_lsn`, used to report metrics.
437async 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        // Based on the documentation, it appears that `pg_current_wal_lsn` has
446        // the same "upper" semantics of `confirmed_flush_lsn`:
447        // <https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-ADMIN-BACKUP>
448        // We may need to revisit this and use `pg_current_wal_flush_lsn`.
449        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
456// Ensures that the table with oid `oid` and expected schema `expected_schema` is still compatible
457// with the current upstream schema `upstream_info`.
458fn 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
482/// Casts a text row into the target types
483fn 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
499/// Converts raw bytes that are expected to be UTF8 encoded into a `Datum::String`
500fn 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}