pgdumpx 0.2.0

Read-only, bounded inspection, extraction, and row scanning for PostgreSQL custom-format dumps
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

use crate::{
    DumpId, PgDumpError,
    custom::primitives::{ArchiveIntegerSize, read_archive_integer},
    error::into_io_error,
    io::archive_reader::ArchiveReader,
};
use std::io::{self, Read};

pub(crate) const BLK_DATA: u8 = 1;

pub(crate) struct CustomChunkReader<'a, R> {
    reader: ArchiveReader<&'a mut R>,
    integer_size: ArchiveIntegerSize,
    dump_id: DumpId,
    remaining: u64,
    done: bool,
}

impl<'a, R> CustomChunkReader<'a, R> {
    pub(crate) const fn new(
        reader: ArchiveReader<&'a mut R>,
        integer_size: ArchiveIntegerSize,
        dump_id: DumpId,
    ) -> Self {
        Self {
            reader,
            integer_size,
            dump_id,
            remaining: 0,
            done: false,
        }
    }

    fn begin_next_chunk(&mut self) -> io::Result<()>
    where
        R: Read,
    {
        let length_offset = self.reader.offset();
        let length = read_archive_integer(&mut self.reader, self.integer_size).map_err(
            |error| match error {
                PgDumpError::UnexpectedEof { .. } => {
                    into_io_error(PgDumpError::TruncatedDataChunkLength {
                        dump_id: self.dump_id.as_i32(),
                        offset: length_offset,
                    })
                }
                other => into_io_error(other),
            },
        )?;
        if length < 0 {
            return Err(into_io_error(PgDumpError::InvalidDataChunkLength {
                dump_id: self.dump_id.as_i32(),
                length,
                offset: length_offset,
            }));
        }
        if length == 0 {
            self.done = true;
            return Ok(());
        }

        self.remaining = u64::try_from(length).map_err(|_| {
            into_io_error(PgDumpError::InvalidDataChunkLength {
                dump_id: self.dump_id.as_i32(),
                length,
                offset: length_offset,
            })
        })?;
        Ok(())
    }
}

impl<R: Read> Read for CustomChunkReader<'_, R> {
    fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
        if output.is_empty() || self.done {
            return Ok(0);
        }

        while self.remaining == 0 && !self.done {
            self.begin_next_chunk()?;
        }
        if self.done {
            return Ok(0);
        }

        let available = usize::try_from(self.remaining).unwrap_or(usize::MAX);
        let requested = output.len().min(available);
        let read = self
            .reader
            .read_some(&mut output[..requested])
            .map_err(into_io_error)?;
        if read == 0 {
            return Err(into_io_error(PgDumpError::TruncatedDataChunk {
                dump_id: self.dump_id.as_i32(),
                remaining: self.remaining,
                offset: self.reader.offset(),
            }));
        }

        let read_u64 = u64::try_from(read).map_err(|_| {
            into_io_error(PgDumpError::ArithmeticOverflow {
                offset: self.reader.offset(),
            })
        })?;
        self.remaining = self.remaining.checked_sub(read_u64).ok_or_else(|| {
            into_io_error(PgDumpError::ArithmeticOverflow {
                offset: self.reader.offset(),
            })
        })?;
        Ok(read)
    }
}