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

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

/// A copy source for Postgres that works well single-threaded workloads.
pub struct SequentialSafePostgresInstanceCopySourceStorage<'a> {
    client: TokioPoolableClient,
    identifier_quoter: Arc<IdentifierQuoter>,
    _lifetime: std::marker::PhantomData<&'a ()>,
}

impl Clone for SequentialSafePostgresInstanceCopySourceStorage<'_> {
    fn clone(&self) -> Self {
        panic!("SequentialSafePostgresInstanceCopySourceStorage should not be cloned")
    }
}

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

        let mut client = wrapper.pool().get_client().await?;
        client
            .execute_non_query_simple(
                "begin transaction isolation level repeatable read read only;",
            )
            .await?;

        Ok(SequentialSafePostgresInstanceCopySourceStorage {
            client,
            identifier_quoter: storage.identifier_quoter.clone(),
            _lifetime: std::marker::PhantomData,
        })
    }
}

pub struct BorrowedPostgresCopyOutReader<'a> {
    reader: CopyReader<'a, TokioConnectionFactory>,
}

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

impl CopySource for SequentialSafePostgresInstanceCopySourceStorage<'_> {
    type DataReader<'a>
        = BorrowedPostgresCopyOutReader<'a>
    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);

        let copy_reader = self.client.copy_out(&*copy_command, &[]).await?;

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