elefant-tools 0.0.8

A library for doing things like pg_dump and pg_restore, with extra features, and probably more bugs.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use crate::storage::postgres::postgres_instance_storage::PostgresInstanceStorage;
use crate::{
    CopySource, DataFormat, IdentifierQuoter, PostgresClientWrapper, PostgresSchema, PostgresTable,
    TableData, TableDataReader,
};
use elefant_client::tokio_connection::{TokioConnectionFactory, TokioPoolableClient};
use elefant_client::OwnedCopyReader;
use std::sync::Arc;
use tracing::instrument;

/// A copy source for Postgres that works well with parallelism.
///
/// This uses repeatable read isolation level and a snapshot to ensure that the data is consistent
/// across the entire dump.
pub struct ParallelSafePostgresInstanceCopySourceStorage<'a> {
    wrapper: &'a PostgresClientWrapper,
    transaction_id: String,
    identifier_quoter: Arc<IdentifierQuoter>,
    /// Holds the client with the snapshot-exporting transaction open.
    /// Must stay alive for the entire duration of the copy operation.
    _snapshot_client: Option<TokioPoolableClient>,
}

impl Clone for ParallelSafePostgresInstanceCopySourceStorage<'_> {
    fn clone(&self) -> Self {
        Self {
            wrapper: self.wrapper,
            transaction_id: self.transaction_id.clone(),
            identifier_quoter: self.identifier_quoter.clone(),
            _snapshot_client: None,
        }
    }
}

impl<'a> ParallelSafePostgresInstanceCopySourceStorage<'a> {
    #[instrument(skip_all)]
    pub async fn new(storage: &PostgresInstanceStorage<'a>) -> crate::Result<Self> {
        let wrapper = storage.connection;

        // Start a repeatable read transaction on a dedicated client and export snapshot
        let mut client = wrapper.pool().get_client().await?;
        client
            .execute_non_query_simple(
                "begin transaction isolation level repeatable read read only;",
            )
            .await?;

        let transaction_id: String = client
            .query_simple("select pg_export_snapshot();")
            .await?
            .collect_single_column_to_vec::<String>()
            .await?
            .into_iter()
            .next()
            .ok_or_else(|| {
                crate::ElefantToolsError::PostgresError(
                    elefant_client::ElefantClientError::PostgresError(
                        "pg_export_snapshot returned no rows".to_string(),
                    ),
                )
            })?;

        Ok(ParallelSafePostgresInstanceCopySourceStorage {
            wrapper,
            transaction_id,
            identifier_quoter: storage.identifier_quoter.clone(),
            _snapshot_client: Some(client),
        })
    }
}

pub struct PostgresCopyOutReader {
    reader: OwnedCopyReader<TokioConnectionFactory>,
}

impl TableDataReader for PostgresCopyOutReader {
    async fn read_chunk(&mut self) -> crate::Result<Option<&[u8]>> {
        let data = self.reader.read().await?;
        Ok(data.map(|cd| cd.data))
    }
}

impl PostgresCopyOutReader {
    pub async fn end(self) -> crate::Result<()> {
        self.reader.end().await?;
        Ok(())
    }
}

impl CopySource for ParallelSafePostgresInstanceCopySourceStorage<'_> {
    type DataReader<'a>
        = PostgresCopyOutReader
    where
        Self: 'a;
    type Cleanup = ();

    #[instrument(skip_all)]
    async fn get_data<'a>(
        &'a mut self,
        schema: &'a PostgresSchema,
        table: &'a PostgresTable,
        data_format: &'a DataFormat,
    ) -> crate::Result<TableData<Self::DataReader<'a>, Self::Cleanup>> {
        let copy_command = table.get_copy_out_command(schema, data_format, &self.identifier_quoter);

        // Get a new client from the pool and set the snapshot
        let mut client = self.wrapper.pool().get_client().await?;
        client
            .execute_non_query_simple(&format!(
                "begin transaction isolation level repeatable read read only; set transaction snapshot '{}';",
                self.transaction_id
            ))
            .await?;

        let reader = OwnedCopyReader::new(client, &*copy_command, &[]).await?;

        Ok(TableData {
            data_format: data_format.clone(),
            data: PostgresCopyOutReader { reader },
            cleanup: (),
        })
    }
}