pgsync_postgres/
snapshot.rs

1//! Streaming table copy under an exported snapshot.
2
3use crate::client::PgClient;
4use crate::errors::from_pg;
5use bytes::{Bytes, BytesMut};
6use futures_util::{SinkExt, StreamExt};
7use pgsync_core::config::{TableName, quote_ident, quote_literal};
8use pgsync_core::metrics::Metrics;
9use pgsync_core::{Error, Result};
10use std::sync::Arc;
11use tracing::debug;
12
13/// Start a read-only REPEATABLE READ transaction on `src` that sees `snapshot_name`.
14/// Must be the first statements on the connection's transaction.
15pub async fn import_snapshot(src: &PgClient, snapshot_name: &str) -> Result<()> {
16    if snapshot_name.is_empty()
17        || !snapshot_name
18            .chars()
19            .all(|c| c.is_ascii_alphanumeric() || c == '-')
20    {
21        return Err(Error::internal(format!(
22            "suspicious snapshot name {snapshot_name:?}"
23        )));
24    }
25    src.batch_execute(&format!(
26        "BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY; SET TRANSACTION SNAPSHOT {};",
27        quote_literal(snapshot_name)
28    ))
29    .await
30    .map_err(|e| from_pg("import snapshot", e))
31}
32
33#[derive(Debug, Clone, Copy, Default)]
34pub struct CopyStats {
35    pub rows: u64,
36    pub bytes: u64,
37}
38
39/// Stream `table` from `src` (inside its snapshot transaction) into `dst` (inside the
40/// caller's target transaction) with `COPY ... TO STDOUT` / `COPY ... FROM STDIN`.
41/// Memory is bounded by `buffer_size` plus one server chunk.
42pub async fn copy_table(
43    src: &PgClient,
44    dst: &PgClient,
45    table: &TableName,
46    columns: &[String],
47    buffer_size: usize,
48    metrics: &Arc<Metrics>,
49) -> Result<CopyStats> {
50    let cols: Vec<String> = columns.iter().map(|c| quote_ident(c)).collect();
51    let col_list = cols.join(", ");
52    let out_sql = format!("COPY {} ({col_list}) TO STDOUT", table.quoted());
53    let in_sql = format!("COPY {} ({col_list}) FROM STDIN", table.quoted());
54
55    let out = src
56        .copy_out(&out_sql)
57        .await
58        .map_err(|e| from_pg(&format!("COPY OUT {table}"), e))?;
59    let sink = dst
60        .copy_in::<_, Bytes>(&in_sql)
61        .await
62        .map_err(|e| from_pg(&format!("COPY IN {table}"), e))?;
63    let mut out = std::pin::pin!(out);
64    let mut sink = std::pin::pin!(sink);
65    let mut buf = BytesMut::with_capacity(buffer_size);
66    let mut stats = CopyStats::default();
67
68    while let Some(chunk) = out.next().await {
69        let chunk = chunk.map_err(|e| from_pg(&format!("COPY OUT {table}"), e))?;
70        stats.bytes += chunk.len() as u64;
71        metrics.add(&metrics.bytes_copied, chunk.len() as u64);
72        buf.extend_from_slice(&chunk);
73        if buf.len() >= buffer_size {
74            sink.send(buf.split().freeze())
75                .await
76                .map_err(|e| from_pg(&format!("COPY IN {table}"), e))?;
77        }
78    }
79    if !buf.is_empty() {
80        sink.send(buf.split().freeze())
81            .await
82            .map_err(|e| from_pg(&format!("COPY IN {table}"), e))?;
83    }
84    stats.rows = sink
85        .as_mut()
86        .finish()
87        .await
88        .map_err(|e| from_pg(&format!("COPY IN {table}"), e))?;
89    metrics.add(&metrics.rows_copied, stats.rows);
90    debug!(%table, rows = stats.rows, bytes = stats.bytes, "table copied");
91    Ok(stats)
92}
pgsync_postgres/snapshot.rs

pgsync_postgres/
snapshot.rs
