mz_storage/source/postgres/snapshot.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//! Renders the table snapshot side of the [`PostgresSourceConnection`] ingestion dataflow.
11//!
12//! # Snapshot reading
13//!
14//! Depending on the resumption LSNs the table reader decides which tables need to be snapshotted
15//! and performs a simple `COPY` query on them in order to get a snapshot. There are a few subtle
16//! points about this operation, described in the following sections.
17//!
18//! ## Consistent LSN point for snapshot transactions
19//!
20//! Given that all our ingestion is based on correctly timestamping updates with the LSN they
21//! happened at it is important that we run the `COPY` query at a specific LSN point that is
22//! relatable with the LSN numbers we receive from the replication stream. Such point does not
23//! necessarily exist for a normal SQL transaction. To achieve this we must force postgres to
24//! produce a consistent point and let us know of the LSN number of that by creating a replication
25//! slot as the first statement in a transaction.
26//!
27//! This is a temporary dummy slot that is only used to put our snapshot transaction on a
28//! consistent LSN point. Unfortunately no lighterweight method exists for doing this. See this
29//! [postgres thread] for more details.
30//!
31//! One might wonder why we don't use the actual real slot to provide us with the snapshot point
32//! which would automatically be at the correct LSN. The answer is that it's possible that we crash
33//! and restart after having already created the slot but before having finished the snapshot. In
34//! that case the restarting process will have lost its opportunity to run queries at the slot's
35//! consistent point as that opportunity only exists in the ephemeral transaction that created the
36//! slot and that is long gone. Additionally there are good reasons of why we'd like to move the
37//! slot creation much earlier, e.g during purification, in which case the slot will always be
38//! pre-created.
39//!
40//! [postgres thread]: https://www.postgresql.org/message-id/flat/CAMN0T-vzzNy6TV1Jvh4xzNQdAvCLBQK_kh6_U7kAXgGU3ZFg-Q%40mail.gmail.com
41//!
42//! ## Reusing the consistent point among all workers
43//!
44//! Creating replication slots is potentially expensive so the code makes is such that all workers
45//! cooperate and reuse one consistent snapshot among them. In order to do so we make use the
46//! "export transaction" feature of postgres. This feature allows one SQL session to create an
47//! identifier for the transaction (a string identifier) it is currently in, which can be used by
48//! other sessions to enter the same "snapshot".
49//!
50//! We accomplish this by picking one worker at random to function as the transaction leader. The
51//! transaction leader is responsible for starting a SQL session, creating a temporary replication
52//! slot in a transaction, exporting the transaction id, and broadcasting the transaction
53//! information to all other workers via a broadcasted feedback edge.
54//!
55//! During this phase the follower workers are simply waiting to hear on the feedback edge,
56//! effectively synchronizing with the leader. Once all workers have received the snapshot
57//! information they can all start to perform their assigned COPY queries.
58//!
59//! The leader and follower steps described above are accomplished by the [`export_snapshot`] and
60//! [`use_snapshot`] functions respectively.
61//!
62//! ## Coordinated transaction COMMIT
63//!
64//! When follower workers are done with snapshotting they commit their transaction, close their
65//! session, and then drop their snapshot feedback capability. When the leader worker is done with
66//! snapshotting it drops its snapshot feedback capability and waits until it observes the
67//! snapshot input advancing to the empty frontier. This allows the leader to COMMIT its
68//! transaction last, which is the transaction that exported the snapshot.
69//!
70//! It's unclear if this is strictly necessary, but having the frontiers made it easy enough that I
71//! added the synchronization.
72//!
73//! ## Snapshot rewinding
74//!
75//! Ingestion dataflows must produce definite data, including the snapshot. What this means
76//! practically is that whenever we deem it necessary to snapshot a table we must do so at the same
77//! LSN. However, the method for running a transaction described above doesn't let us choose the
78//! LSN, it could be an LSN in the future chosen by PostgresSQL while it creates the temporary
79//! replication slot.
80//!
81//! The definition of differential collections states that a collection at some time `t_snapshot`
82//! is defined to be the accumulation of all updates that happen at `t <= t_snapshot`, where `<=`
83//! is the partial order. In this case we are faced with the problem of knowing the state of a
84//! table at `t_snapshot` but actually wanting to know the snapshot at `t_slot <= t_snapshot`.
85//!
86//! From the definition we can see that the snapshot at `t_slot` is related to the snapshot at
87//! `t_snapshot` with the following equations:
88//!
89//!```text
90//! sum(update: t <= t_snapshot) = sum(update: t <= t_slot) + sum(update: t_slot <= t <= t_snapshot)
91//! |
92//! V
93//! sum(update: t <= t_slot) = sum(update: t <= snapshot) - sum(update: t_slot <= t <= t_snapshot)
94//! ```
95//!
96//! Therefore, if we manage to recover the `sum(update: t_slot <= t <= t_snapshot)` term we will be
97//! able to "rewind" the snapshot we obtained at `t_snapshot` to `t_slot` by emitting all updates
98//! that happen between these two points with their diffs negated.
99//!
100//! It turns out that this term is exactly what the main replication slot provides us with and we
101//! can rewind snapshot at arbitrary points! In order to do this the snapshot dataflow emits rewind
102//! requests to the replication reader which informs it that a certain range of updates must be
103//! emitted at LSN 0 (by convention) with their diffs negated. These negated diffs are consolidated
104//! with the diffs taken at `t_snapshot` that were also emitted at LSN 0 (by convention) and we end
105//! up with a TVC that at LSN 0 contains the snapshot at `t_slot`.
106//!
107//! # Snapshot decoding
108//!
109//! The expectation is that tables will most likely be skewed on the number of rows they contain so
110//! while a `COPY` query for any given table runs on a single worker the decoding of the COPY
111//! stream is distributed to all workers.
112//!
113//!
114//! ```text
115//! ╭──────────────────╮
116//! ┏━━━━━━━━━━━━v━┓ │ exported
117//! ┃ table ┃ ╭─────────╮ │ snapshot id
118//! ┃ reader ┠─>─┤broadcast├──╯
119//! ┗━┯━━━━━━━━━━┯━┛ ╰─────────╯
120//! raw│ │
121//! COPY│ │
122//! data│ │
123//! ╭────┴─────╮ │
124//! │distribute│ │
125//! ╰────┬─────╯ │
126//! ┏━━━━┷━━━━┓ │
127//! ┃ COPY ┃ │
128//! ┃ decoder ┃ │
129//! ┗━━━━┯━━━━┛ │
130//! │ snapshot │rewind
131//! │ updates │requests
132//! v v
133//! ```
134
135use std::collections::BTreeMap;
136use std::convert::Infallible;
137use std::pin::pin;
138use std::rc::Rc;
139use std::sync::Arc;
140use std::time::Duration;
141
142use anyhow::bail;
143use differential_dataflow::AsCollection;
144use futures::{StreamExt as _, TryStreamExt};
145use mz_ore::cast::CastFrom;
146use mz_ore::future::InTask;
147use mz_postgres_util::tunnel::PostgresFlavor;
148use mz_postgres_util::{Client, PostgresError, simple_query_opt};
149use mz_repr::{Datum, DatumVec, Diff, Row};
150use mz_sql_parser::ast::{Ident, display::AstDisplay};
151use mz_storage_types::errors::DataflowError;
152use mz_storage_types::parameters::PgSourceSnapshotConfig;
153use mz_storage_types::sources::{MzOffset, PostgresSourceConnection};
154use mz_timely_util::builder_async::{
155 Event as AsyncEvent, OperatorBuilder as AsyncOperatorBuilder, PressOnDropButton,
156};
157use timely::container::CapacityContainerBuilder;
158use timely::dataflow::channels::pact::{Exchange, Pipeline};
159use timely::dataflow::operators::core::Map;
160use timely::dataflow::operators::{
161 Broadcast, CapabilitySet, Concat, ConnectLoop, Feedback, Operator,
162};
163use timely::dataflow::{Scope, Stream};
164use timely::progress::Timestamp;
165use tokio_postgres::error::SqlState;
166use tokio_postgres::types::{Oid, PgLsn};
167use tracing::trace;
168
169use crate::metrics::source::postgres::PgSnapshotMetrics;
170use crate::source::RawSourceCreationConfig;
171use crate::source::postgres::replication::RewindRequest;
172use crate::source::postgres::{
173 DefiniteError, ReplicationError, SourceOutputInfo, TransientError, verify_schema,
174};
175use crate::source::types::{
176 ProgressStatisticsUpdate, SignaledFuture, SourceMessage, StackedCollection,
177};
178
179/// Renders the snapshot dataflow. See the module documentation for more information.
180pub(crate) fn render<G: Scope<Timestamp = MzOffset>>(
181 mut scope: G,
182 config: RawSourceCreationConfig,
183 connection: PostgresSourceConnection,
184 table_info: BTreeMap<u32, BTreeMap<usize, SourceOutputInfo>>,
185 metrics: PgSnapshotMetrics,
186) -> (
187 StackedCollection<G, (usize, Result<SourceMessage, DataflowError>)>,
188 Stream<G, RewindRequest>,
189 Stream<G, Infallible>,
190 Stream<G, ProgressStatisticsUpdate>,
191 Stream<G, ReplicationError>,
192 PressOnDropButton,
193) {
194 let op_name = format!("TableReader({})", config.id);
195 let mut builder = AsyncOperatorBuilder::new(op_name, scope.clone());
196
197 let (feedback_handle, feedback_data) = scope.feedback(Default::default());
198
199 let (raw_handle, raw_data) = builder.new_output();
200 let (rewinds_handle, rewinds) = builder.new_output();
201 // This output is used to signal to the replication operator that the replication slot has been
202 // created. With the current state of execution serialization there isn't a lot of benefit
203 // of splitting the snapshot and replication phases into two operators.
204 // TODO(petrosagg): merge the two operators in one (while still maintaining separation as
205 // functions/modules)
206 let (_, slot_ready) = builder.new_output::<CapacityContainerBuilder<_>>();
207 let (snapshot_handle, snapshot) = builder.new_output();
208 let (definite_error_handle, definite_errors) = builder.new_output();
209
210 let (stats_output, stats_stream) = builder.new_output();
211
212 // This operator needs to broadcast data to itself in order to synchronize the transaction
213 // snapshot. However, none of the feedback capabilities result in output messages and for the
214 // feedback edge specifically having a default conncetion would result in a loop.
215 let mut snapshot_input = builder.new_disconnected_input(&feedback_data, Pipeline);
216
217 // The export id must be sent to all workers, so we broadcast the feedback connection
218 snapshot.broadcast().connect_loop(feedback_handle);
219
220 let is_snapshot_leader = config.responsible_for("snapshot_leader");
221
222 // A global view of all outputs that will be snapshot by all workers.
223 let mut all_outputs = vec![];
224 // A filtered table info containing only the tables that this worker should snapshot.
225 let mut reader_table_info = BTreeMap::new();
226 for (table, outputs) in table_info.iter() {
227 for (&output_index, output) in outputs {
228 if *output.resume_upper != [MzOffset::minimum()] {
229 // Already has been snapshotted.
230 continue;
231 }
232 all_outputs.push(output_index);
233 if config.responsible_for(*table) {
234 reader_table_info
235 .entry(*table)
236 .or_insert_with(BTreeMap::new)
237 .insert(output_index, (output.desc.clone(), output.casts.clone()));
238 }
239 }
240 }
241
242 let (button, transient_errors) = builder.build_fallible(move |caps| {
243 let busy_signal = Arc::clone(&config.busy_signal);
244 Box::pin(SignaledFuture::new(busy_signal, async move {
245 let id = config.id;
246 let worker_id = config.worker_id;
247
248 let [
249 data_cap_set,
250 rewind_cap_set,
251 slot_ready_cap_set,
252 snapshot_cap_set,
253 definite_error_cap_set,
254 stats_cap,
255 ]: &mut [_; 6] = caps.try_into().unwrap();
256
257 trace!(
258 %id,
259 "timely-{worker_id} initializing table reader \
260 with {} tables to snapshot",
261 reader_table_info.len()
262 );
263
264 // Nothing needs to be snapshot.
265 if all_outputs.is_empty() {
266 trace!(%id, "no exports to snapshot");
267 // Note we do not emit a `ProgressStatisticsUpdate::Snapshot` update here,
268 // as we do not want to attempt to override the current value with 0. We
269 // just leave it null.
270 return Ok(());
271 }
272
273 let connection_config = connection
274 .connection
275 .config(
276 &config.config.connection_context.secrets_reader,
277 &config.config,
278 InTask::Yes,
279 )
280 .await?;
281 let task_name = format!("timely-{worker_id} PG snapshotter");
282
283 let client = if is_snapshot_leader {
284 let client = connection_config
285 .connect_replication(&config.config.connection_context.ssh_tunnel_manager)
286 .await?;
287
288 // Attempt to export the snapshot by creating the main replication slot. If that
289 // succeeds then there is no need for creating additional temporary slots.
290 let main_slot = &connection.publication_details.slot;
291 let snapshot_info = match export_snapshot(&client, main_slot, false).await {
292 Ok(info) => info,
293 Err(err @ TransientError::ReplicationSlotAlreadyExists) => {
294 match connection.connection.flavor {
295 // If we're connecting to a vanilla we have the option of exporting a
296 // snapshot via a temporary slot
297 PostgresFlavor::Vanilla => {
298 let tmp_slot = format!(
299 "mzsnapshot_{}",
300 uuid::Uuid::new_v4()).replace('-', ""
301 );
302 export_snapshot(&client, &tmp_slot, true).await?
303 }
304 // No salvation for Yugabyte
305 PostgresFlavor::Yugabyte => return Err(err),
306 }
307 }
308 Err(err) => return Err(err),
309 };
310 trace!(
311 %id,
312 "timely-{worker_id} exporting snapshot info {snapshot_info:?}");
313 snapshot_handle.give(&snapshot_cap_set[0], snapshot_info);
314
315 client
316 } else {
317 // Only the snapshot leader needs a replication connection.
318 connection_config
319 .connect(
320 &task_name,
321 &config.config.connection_context.ssh_tunnel_manager,
322 )
323 .await?
324 };
325 *slot_ready_cap_set = CapabilitySet::new();
326
327 // Configure statement_timeout based on param. We want to be able to
328 // override the server value here in case it's set too low,
329 // respective to the size of the data we need to copy.
330 set_statement_timeout(
331 &client,
332 config
333 .config
334 .parameters
335 .pg_source_snapshot_statement_timeout,
336 )
337 .await?;
338
339 let (snapshot, snapshot_lsn) = loop {
340 match snapshot_input.next().await {
341 Some(AsyncEvent::Data(_, mut data)) => {
342 break data.pop().expect("snapshot sent above")
343 }
344 Some(AsyncEvent::Progress(_)) => continue,
345 None => panic!(
346 "feedback closed \
347 before sending snapshot info"
348 ),
349 }
350 };
351 // Snapshot leader is already in identified transaction but all other workers need to enter it.
352 if !is_snapshot_leader {
353 trace!(%id, "timely-{worker_id} using snapshot id {snapshot:?}");
354 use_snapshot(&client, &snapshot).await?;
355 }
356
357 let upstream_info = {
358 let schema_client = connection_config
359 .connect(
360 "snapshot schema info",
361 &config.config.connection_context.ssh_tunnel_manager,
362 )
363 .await?;
364 match mz_postgres_util::publication_info(&schema_client, &connection.publication, Some(&reader_table_info.keys().copied().collect::<Vec<_>>()))
365 .await
366 {
367 // If the replication stream cannot be obtained in a definite way there is
368 // nothing else to do. These errors are not retractable.
369 Err(PostgresError::PublicationMissing(publication)) => {
370 let err = DefiniteError::PublicationDropped(publication);
371 for (oid, outputs) in reader_table_info.iter() {
372 // Produce a definite error here and then exit to ensure
373 // a missing publication doesn't generate a transient
374 // error and restart this dataflow indefinitely.
375 //
376 // We pick `u64::MAX` as the LSN which will (in
377 // practice) never conflict any previously revealed
378 // portions of the TVC.
379 for output_index in outputs.keys() {
380 let update = (
381 (*oid, *output_index, Err(err.clone().into())),
382 MzOffset::from(u64::MAX),
383 Diff::ONE,
384 );
385 raw_handle.give_fueled(&data_cap_set[0], update).await;
386 }
387 }
388
389 definite_error_handle.give(
390 &definite_error_cap_set[0],
391 ReplicationError::Definite(Rc::new(err)),
392 );
393 return Ok(());
394 }
395 Err(e) => Err(TransientError::from(e))?,
396 Ok(i) => i,
397 }
398 };
399
400 let worker_tables = reader_table_info
401 .iter()
402 .map(|(_, outputs)| {
403 // just use the first output's desc since the fields accessed here should
404 // be the same for all outputs
405 let desc = &outputs.values().next().expect("at least 1").0;
406 (
407 format!(
408 "{}.{}",
409 Ident::new_unchecked(desc.namespace.clone()).to_ast_string_simple(),
410 Ident::new_unchecked(desc.name.clone()).to_ast_string_simple()
411 ),
412 desc.oid.clone(),
413 outputs.len(),
414 )
415 })
416 .collect();
417
418 let snapshot_total =
419 fetch_snapshot_size(&client, worker_tables, metrics, &config).await?;
420
421 stats_output.give(
422 &stats_cap[0],
423 ProgressStatisticsUpdate::Snapshot {
424 records_known: snapshot_total,
425 records_staged: 0,
426 },
427 );
428
429 let mut snapshot_staged = 0;
430 for (&oid, outputs) in reader_table_info.iter() {
431 let mut table_name = None;
432 let mut output_indexes = vec![];
433 for (output_index, (expected_desc, casts)) in outputs.iter() {
434 match verify_schema(oid, expected_desc, &upstream_info, casts) {
435 Ok(()) => {
436 if table_name.is_none() {
437 table_name = Some((
438 expected_desc.namespace.clone(),
439 expected_desc.name.clone(),
440 ));
441 }
442 output_indexes.push(output_index);
443 }
444 Err(err) => {
445 raw_handle
446 .give_fueled(
447 &data_cap_set[0],
448 (
449 (oid, *output_index, Err(err.into())),
450 MzOffset::minimum(),
451 Diff::ONE,
452 ),
453 )
454 .await;
455 continue;
456 }
457 };
458 }
459
460 let (namespace, table) = match table_name {
461 Some(t) => t,
462 None => {
463 // all outputs errored for this table
464 continue;
465 }
466 };
467
468 trace!(
469 %id,
470 "timely-{worker_id} snapshotting table {:?}({oid}) @ {snapshot_lsn}",
471 table
472 );
473
474 // To handle quoted/keyword names, we can use `Ident`'s AST printing, which
475 // emulate's PG's rules for name formatting.
476 let query = format!(
477 "COPY {}.{} TO STDOUT (FORMAT TEXT, DELIMITER '\t')",
478 Ident::new_unchecked(namespace).to_ast_string_simple(),
479 Ident::new_unchecked(table).to_ast_string_simple(),
480 );
481 let mut stream = pin!(client.copy_out_simple(&query).await?);
482
483 let mut update = ((oid, 0, Ok(vec![])), MzOffset::minimum(), Diff::ONE);
484 while let Some(bytes) = stream.try_next().await? {
485 let data = update.0 .2.as_mut().unwrap();
486 data.clear();
487 data.extend_from_slice(&bytes);
488 for output_index in &output_indexes {
489 update.0 .1 = **output_index;
490 raw_handle.give_fueled(&data_cap_set[0], &update).await;
491 snapshot_staged += 1;
492 // TODO(guswynn): does this 1000 need to be configurable?
493 if snapshot_staged % 1000 == 0 {
494 stats_output.give(
495 &stats_cap[0],
496 ProgressStatisticsUpdate::Snapshot {
497 records_known: snapshot_total,
498 records_staged: snapshot_staged,
499 },
500 );
501 }
502 }
503 }
504 }
505
506 // We are done with the snapshot so now we will emit rewind requests. It is important
507 // that this happens after the snapshot has finished because this is what unblocks the
508 // replication operator and we want this to happen serially. It might seem like a good
509 // idea to read the replication stream concurrently with the snapshot but it actually
510 // leads to a lot of data being staged for the future, which needlesly consumed memory
511 // in the cluster.
512 for output in reader_table_info.values() {
513 for (output_index, (desc, _)) in output {
514 trace!(%id, "timely-{worker_id} producing rewind request for table {} output {output_index}", desc.name);
515 let req = RewindRequest { output_index: *output_index, snapshot_lsn };
516 rewinds_handle.give(&rewind_cap_set[0], req);
517 }
518 }
519 *rewind_cap_set = CapabilitySet::new();
520
521 // Report the same known and staged records to signify that the snapshot is complete.
522 stats_output.give(
523 &stats_cap[0],
524 ProgressStatisticsUpdate::Snapshot {
525 records_known: snapshot_staged,
526 records_staged: snapshot_staged,
527 },
528 );
529
530 // Failure scenario after we have produced the snapshot, but before a successful COMMIT
531 fail::fail_point!("pg_snapshot_failure", |_| Err(
532 TransientError::SyntheticError
533 ));
534
535 // The exporting worker should wait for all the other workers to commit before dropping
536 // its client since this is what holds the exported transaction alive.
537 if is_snapshot_leader {
538 trace!(%id, "timely-{worker_id} waiting for all workers to finish");
539 *snapshot_cap_set = CapabilitySet::new();
540 while snapshot_input.next().await.is_some() {}
541 trace!(%id, "timely-{worker_id} (leader) comitting COPY transaction");
542 client.simple_query("COMMIT").await?;
543 } else {
544 trace!(%id, "timely-{worker_id} comitting COPY transaction");
545 client.simple_query("COMMIT").await?;
546 *snapshot_cap_set = CapabilitySet::new();
547 }
548 drop(client);
549 Ok(())
550 }))
551 });
552
553 // We now decode the COPY protocol and apply the cast expressions
554 let mut text_row = Row::default();
555 let mut final_row = Row::default();
556 let mut datum_vec = DatumVec::new();
557 let mut next_worker = (0..u64::cast_from(scope.peers()))
558 // Round robin on 1000-records basis to avoid creating tiny containers when there are a
559 // small number of updates and a large number of workers.
560 .flat_map(|w| std::iter::repeat_n(w, 1000))
561 .cycle();
562 let round_robin = Exchange::new(move |_| next_worker.next().unwrap());
563 let snapshot_updates = raw_data
564 .map::<Vec<_>, _, _>(Clone::clone)
565 .unary(round_robin, "PgCastSnapshotRows", |_, _| {
566 move |input, output| {
567 while let Some((time, data)) = input.next() {
568 let mut session = output.session(&time);
569 for ((oid, output_index, event), time, diff) in data.drain(..) {
570 let output = &table_info
571 .get(&oid)
572 .and_then(|outputs| outputs.get(&output_index))
573 .expect("table_info contains all outputs");
574
575 let event = event
576 .as_ref()
577 .map_err(|e: &DataflowError| e.clone())
578 .and_then(|bytes| {
579 decode_copy_row(bytes, output.casts.len(), &mut text_row)?;
580 let datums = datum_vec.borrow_with(&text_row);
581 super::cast_row(&output.casts, &datums, &mut final_row)?;
582 Ok(SourceMessage {
583 key: Row::default(),
584 value: final_row.clone(),
585 metadata: Row::default(),
586 })
587 });
588
589 session.give(((output_index, event), time, diff));
590 }
591 }
592 }
593 })
594 .as_collection();
595
596 let errors = definite_errors.concat(&transient_errors.map(ReplicationError::from));
597
598 (
599 snapshot_updates,
600 rewinds,
601 slot_ready,
602 stats_stream,
603 errors,
604 button.press_on_drop(),
605 )
606}
607
608/// Starts a read-only transaction on the SQL session of `client` at a consistent LSN point by
609/// creating a replication slot. Returns a snapshot identifier that can be imported in
610/// other SQL session and the LSN of the consistent point.
611async fn export_snapshot(
612 client: &Client,
613 slot: &str,
614 temporary: bool,
615) -> Result<(String, MzOffset), TransientError> {
616 match export_snapshot_inner(client, slot, temporary).await {
617 Ok(ok) => Ok(ok),
618 Err(err) => {
619 // We don't want to leave the client inside a failed tx
620 client.simple_query("ROLLBACK;").await?;
621 Err(err)
622 }
623 }
624}
625
626async fn export_snapshot_inner(
627 client: &Client,
628 slot: &str,
629 temporary: bool,
630) -> Result<(String, MzOffset), TransientError> {
631 client
632 .simple_query("BEGIN READ ONLY ISOLATION LEVEL REPEATABLE READ;")
633 .await?;
634
635 // Note: Using unchecked here is okay because we're using it in a SQL query.
636 let slot = Ident::new_unchecked(slot).to_ast_string_simple();
637 let temporary_str = if temporary { " TEMPORARY" } else { "" };
638 let query =
639 format!("CREATE_REPLICATION_SLOT {slot}{temporary_str} LOGICAL \"pgoutput\" USE_SNAPSHOT");
640 let row = match simple_query_opt(client, &query).await {
641 Ok(row) => Ok(row.unwrap()),
642 Err(PostgresError::Postgres(err)) if err.code() == Some(&SqlState::DUPLICATE_OBJECT) => {
643 return Err(TransientError::ReplicationSlotAlreadyExists);
644 }
645 Err(err) => Err(err),
646 }?;
647
648 // When creating a replication slot postgres returns the LSN of its consistent point, which is
649 // the LSN that must be passed to `START_REPLICATION` to cleanly transition from the snapshot
650 // phase to the replication phase. `START_REPLICATION` includes all transactions that commit at
651 // LSNs *greater than or equal* to the passed LSN. Therefore the snapshot phase must happen at
652 // the greatest LSN that is not beyond the consistent point. That LSN is `consistent_point - 1`
653 let consistent_point: PgLsn = row.get("consistent_point").unwrap().parse().unwrap();
654 let consistent_point = u64::from(consistent_point)
655 .checked_sub(1)
656 .expect("consistent point is always non-zero");
657
658 let row = simple_query_opt(client, "SELECT pg_export_snapshot();")
659 .await?
660 .unwrap();
661 let snapshot = row.get("pg_export_snapshot").unwrap().to_owned();
662
663 Ok((snapshot, MzOffset::from(consistent_point)))
664}
665
666/// Starts a read-only transaction on the SQL session of `client` at a the consistent LSN point of
667/// `snapshot`.
668async fn use_snapshot(client: &Client, snapshot: &str) -> Result<(), TransientError> {
669 client
670 .simple_query("BEGIN READ ONLY ISOLATION LEVEL REPEATABLE READ;")
671 .await?;
672 let query = format!("SET TRANSACTION SNAPSHOT '{snapshot}';");
673 client.simple_query(&query).await?;
674 Ok(())
675}
676
677async fn set_statement_timeout(client: &Client, timeout: Duration) -> Result<(), TransientError> {
678 // Value is known to accept milliseconds w/o units.
679 // https://www.postgresql.org/docs/current/runtime-config-client.html
680 client
681 .simple_query(&format!("SET statement_timeout = {}", timeout.as_millis()))
682 .await?;
683 Ok(())
684}
685
686/// Decodes a row of `col_len` columns obtained from a text encoded COPY query into `row`.
687fn decode_copy_row(data: &[u8], col_len: usize, row: &mut Row) -> Result<(), DefiniteError> {
688 let mut packer = row.packer();
689 let row_parser = mz_pgcopy::CopyTextFormatParser::new(data, b'\t', "\\N");
690 let mut column_iter = row_parser.iter_raw_truncating(col_len);
691 for _ in 0..col_len {
692 let value = match column_iter.next() {
693 Some(Ok(value)) => value,
694 Some(Err(_)) => return Err(DefiniteError::InvalidCopyInput),
695 None => return Err(DefiniteError::MissingColumn),
696 };
697 let datum = value.map(super::decode_utf8_text).transpose()?;
698 packer.push(datum.unwrap_or(Datum::Null));
699 }
700 Ok(())
701}
702
703/// Record the sizes of the tables being snapshotted in `PgSnapshotMetrics`.
704async fn fetch_snapshot_size(
705 client: &Client,
706 // The table names, oids, and number of outputs for this table owned by this worker.
707 tables: Vec<(String, Oid, usize)>,
708 metrics: PgSnapshotMetrics,
709 config: &RawSourceCreationConfig,
710) -> Result<u64, anyhow::Error> {
711 // TODO(guswynn): delete unused configs
712 let snapshot_config = config.config.parameters.pg_snapshot_config;
713
714 let mut total = 0;
715 for (table, oid, output_count) in tables {
716 let stats = collect_table_statistics(client, snapshot_config, &table, oid).await?;
717 metrics.record_table_count_latency(table, stats.count_latency);
718 total += stats.count * u64::cast_from(output_count);
719 }
720 Ok(total)
721}
722
723#[derive(Default)]
724struct TableStatistics {
725 count: u64,
726 count_latency: f64,
727}
728
729async fn collect_table_statistics(
730 client: &Client,
731 config: PgSourceSnapshotConfig,
732 table: &str,
733 oid: u32,
734) -> Result<TableStatistics, anyhow::Error> {
735 use mz_ore::metrics::MetricsFutureExt;
736 let mut stats = TableStatistics::default();
737
738 let estimate_row = simple_query_opt(
739 client,
740 &format!("SELECT reltuples::bigint AS estimate_count FROM pg_class WHERE oid = '{oid}'"),
741 )
742 .wall_time()
743 .set_at(&mut stats.count_latency)
744 .await?;
745 stats.count = match estimate_row {
746 Some(row) => row.get("estimate_count").unwrap().parse().unwrap_or(0),
747 None => bail!("failed to get estimate count for {table}"),
748 };
749
750 // If the estimate is low enough we can attempt to get an exact count. Note that not yet
751 // vacuumed tables will report zero rows here and there is a possibility that they are very
752 // large. We accept this risk and we offer the feature flag as an escape hatch if it becomes
753 // problematic.
754 if config.collect_strict_count && stats.count < 1_000_000 {
755 let count_row = simple_query_opt(client, &format!("SELECT count(*) as count from {table}"))
756 .wall_time()
757 .set_at(&mut stats.count_latency)
758 .await?;
759 stats.count = match count_row {
760 Some(row) => row.get("count").unwrap().parse().unwrap(),
761 None => bail!("failed to get count for {table}"),
762 }
763 }
764
765 Ok(stats)
766}