pg_client/
pg_dump.rs

1use std::borrow::Cow;
2
3use core::fmt::{Display, Formatter};
4use core::str::FromStr;
5
6use crate::identifier::{QualifiedTable, Schema};
7
8/// Alphanumeric key for the `\restrict` / `\unrestrict` psql meta-commands
9/// emitted by `pg_dump` (CVE-2025-8714).
10///
11/// When set, `pg_dump` uses this fixed key instead of generating a random one,
12/// making plain-text dump output deterministic across invocations.
13///
14/// Constraints: non-empty, alphanumeric (`a-zA-Z0-9`), max 63 bytes
15/// (matching the auto-generated key length).
16///
17/// See: <https://www.postgresql.org/docs/current/app-pgdump.html>
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub struct RestrictKey(Cow<'static, str>);
20
21/// Maximum length of a restrict key in bytes.
22const MAX_LENGTH: usize = 63;
23
24/// Const-compatible validation that returns an optional error.
25const fn validate_restrict_key(input: &str) -> Option<RestrictKeyParseError> {
26    if input.is_empty() {
27        return Some(RestrictKeyParseError::Empty);
28    }
29
30    if input.len() > MAX_LENGTH {
31        return Some(RestrictKeyParseError::TooLong);
32    }
33
34    let bytes = input.as_bytes();
35    let mut index = 0;
36
37    while index < bytes.len() {
38        if !bytes[index].is_ascii_alphanumeric() {
39            return Some(RestrictKeyParseError::NotAlphanumeric);
40        }
41        index += 1;
42    }
43
44    None
45}
46
47impl RestrictKey {
48    /// Creates a new restrict key from a static string.
49    ///
50    /// # Panics
51    ///
52    /// Panics if the input is empty, exceeds 63 bytes, or contains non-alphanumeric bytes.
53    #[must_use]
54    pub const fn from_static_or_panic(input: &'static str) -> Self {
55        match validate_restrict_key(input) {
56            Some(error) => panic!("{}", error.message()),
57            None => Self(Cow::Borrowed(input)),
58        }
59    }
60
61    /// Returns the value as a string slice.
62    #[must_use]
63    pub fn as_str(&self) -> &str {
64        &self.0
65    }
66}
67
68impl From<&RestrictKey> for RestrictKey {
69    fn from(key: &RestrictKey) -> Self {
70        match &key.0 {
71            Cow::Borrowed(s) => Self(Cow::Borrowed(s)),
72            Cow::Owned(s) => Self(Cow::Owned(s.clone())),
73        }
74    }
75}
76
77impl AsRef<str> for RestrictKey {
78    fn as_ref(&self) -> &str {
79        &self.0
80    }
81}
82
83impl Display for RestrictKey {
84    fn fmt(&self, formatter: &mut Formatter<'_>) -> core::fmt::Result {
85        formatter.write_str(&self.0)
86    }
87}
88
89impl FromStr for RestrictKey {
90    type Err = RestrictKeyParseError;
91
92    fn from_str(input: &str) -> Result<Self, Self::Err> {
93        match validate_restrict_key(input) {
94            Some(error) => Err(error),
95            None => Ok(Self(Cow::Owned(input.to_owned()))),
96        }
97    }
98}
99
100/// Error parsing a restrict key.
101#[derive(Debug, Clone, Copy, PartialEq, Eq)]
102pub enum RestrictKeyParseError {
103    /// Key cannot be empty.
104    Empty,
105    /// Key exceeds maximum length.
106    TooLong,
107    /// Key contains non-alphanumeric bytes.
108    NotAlphanumeric,
109}
110
111impl RestrictKeyParseError {
112    /// Returns the error message.
113    #[must_use]
114    pub const fn message(&self) -> &'static str {
115        match self {
116            Self::Empty => "restrict key cannot be empty",
117            Self::TooLong => "restrict key exceeds 63 byte max length",
118            Self::NotAlphanumeric => "restrict key must be alphanumeric",
119        }
120    }
121}
122
123impl Display for RestrictKeyParseError {
124    fn fmt(&self, formatter: &mut Formatter<'_>) -> core::fmt::Result {
125        formatter.write_str(self.message())
126    }
127}
128
129impl std::error::Error for RestrictKeyParseError {}
130
131#[must_use]
132pub struct PgSchemaDump {
133    exclude_schemas: Vec<Schema>,
134    exclude_tables: Vec<QualifiedTable>,
135    no_comments: bool,
136    no_owner: bool,
137    no_privileges: bool,
138    no_tablespaces: bool,
139    restrict_key: Option<RestrictKey>,
140    schemas: Vec<Schema>,
141    tables: Vec<QualifiedTable>,
142    verbose: bool,
143}
144
145impl Default for PgSchemaDump {
146    fn default() -> Self {
147        Self::new()
148    }
149}
150
151impl PgSchemaDump {
152    pub fn new() -> Self {
153        Self {
154            exclude_schemas: Vec::new(),
155            exclude_tables: Vec::new(),
156            no_comments: false,
157            no_owner: false,
158            no_privileges: false,
159            no_tablespaces: false,
160            restrict_key: None,
161            schemas: Vec::new(),
162            tables: Vec::new(),
163            verbose: false,
164        }
165    }
166
167    pub fn exclude_schema(mut self, schema: Schema) -> Self {
168        self.exclude_schemas.push(schema);
169        self
170    }
171
172    pub fn exclude_table(mut self, table: QualifiedTable) -> Self {
173        self.exclude_tables.push(table);
174        self
175    }
176
177    pub fn no_comments(mut self) -> Self {
178        self.no_comments = true;
179        self
180    }
181
182    pub fn no_owner(mut self) -> Self {
183        self.no_owner = true;
184        self
185    }
186
187    pub fn no_privileges(mut self) -> Self {
188        self.no_privileges = true;
189        self
190    }
191
192    pub fn no_tablespaces(mut self) -> Self {
193        self.no_tablespaces = true;
194        self
195    }
196
197    pub fn restrict_key(mut self, restrict_key: &RestrictKey) -> Self {
198        self.restrict_key = Some(RestrictKey::from(restrict_key));
199        self
200    }
201
202    pub fn schema(mut self, schema: Schema) -> Self {
203        self.schemas.push(schema);
204        self
205    }
206
207    pub fn table(mut self, table: QualifiedTable) -> Self {
208        self.tables.push(table);
209        self
210    }
211
212    pub fn verbose(mut self) -> Self {
213        self.verbose = true;
214        self
215    }
216
217    #[must_use]
218    pub fn arguments(&self) -> Vec<String> {
219        let mut args = vec!["--schema-only".to_string()];
220
221        for schema in &self.exclude_schemas {
222            args.push("--exclude-schema".to_string());
223            args.push(schema.to_string());
224        }
225
226        for table in &self.exclude_tables {
227            args.push("--exclude-table".to_string());
228            args.push(table.to_string());
229        }
230
231        if self.no_comments {
232            args.push("--no-comments".to_string());
233        }
234
235        if self.no_owner {
236            args.push("--no-owner".to_string());
237        }
238
239        if self.no_privileges {
240            args.push("--no-privileges".to_string());
241        }
242
243        if self.no_tablespaces {
244            args.push("--no-tablespaces".to_string());
245        }
246
247        if let Some(restrict_key) = &self.restrict_key {
248            args.push(format!("--restrict-key={restrict_key}"));
249        }
250
251        for schema in &self.schemas {
252            args.push("--schema".to_string());
253            args.push(schema.to_string());
254        }
255
256        for table in &self.tables {
257            args.push("--table".to_string());
258            args.push(table.to_string());
259        }
260
261        if self.verbose {
262            args.push("--verbose".to_string());
263        }
264
265        args
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn test_default() {
275        assert_eq!(PgSchemaDump::new().arguments(), vec!["--schema-only"]);
276    }
277
278    #[test]
279    fn test_verbose() {
280        assert_eq!(
281            PgSchemaDump::new().verbose().arguments(),
282            vec!["--schema-only", "--verbose"],
283        );
284    }
285
286    #[test]
287    fn test_exclude_schema() {
288        assert_eq!(
289            PgSchemaDump::new()
290                .exclude_schema("internal".parse().unwrap())
291                .arguments(),
292            vec!["--schema-only", "--exclude-schema", "internal"],
293        );
294    }
295
296    #[test]
297    fn test_exclude_table() {
298        assert_eq!(
299            PgSchemaDump::new()
300                .exclude_table(QualifiedTable {
301                    schema: Schema::PUBLIC,
302                    table: "some_table".parse().unwrap(),
303                })
304                .arguments(),
305            vec!["--schema-only", "--exclude-table", "public.some_table"],
306        );
307    }
308
309    #[test]
310    fn test_no_comments() {
311        assert_eq!(
312            PgSchemaDump::new().no_comments().arguments(),
313            vec!["--schema-only", "--no-comments"],
314        );
315    }
316
317    #[test]
318    fn test_no_owner() {
319        assert_eq!(
320            PgSchemaDump::new().no_owner().arguments(),
321            vec!["--schema-only", "--no-owner"],
322        );
323    }
324
325    #[test]
326    fn test_no_privileges() {
327        assert_eq!(
328            PgSchemaDump::new().no_privileges().arguments(),
329            vec!["--schema-only", "--no-privileges"],
330        );
331    }
332
333    #[test]
334    fn test_no_tablespaces() {
335        assert_eq!(
336            PgSchemaDump::new().no_tablespaces().arguments(),
337            vec!["--schema-only", "--no-tablespaces"],
338        );
339    }
340
341    #[test]
342    fn test_schema() {
343        assert_eq!(
344            PgSchemaDump::new().schema(Schema::PUBLIC).arguments(),
345            vec!["--schema-only", "--schema", "public"],
346        );
347    }
348
349    #[test]
350    fn test_table() {
351        assert_eq!(
352            PgSchemaDump::new()
353                .table(QualifiedTable {
354                    schema: Schema::PUBLIC,
355                    table: "users".parse().unwrap(),
356                })
357                .arguments(),
358            vec!["--schema-only", "--table", "public.users"],
359        );
360    }
361
362    #[test]
363    fn test_restrict_key() {
364        assert_eq!(
365            PgSchemaDump::new()
366                .restrict_key(&"abc123".parse().unwrap())
367                .arguments(),
368            vec!["--schema-only", "--restrict-key=abc123"],
369        );
370    }
371
372    #[test]
373    fn test_restrict_key_empty() {
374        assert_eq!("".parse::<RestrictKey>(), Err(RestrictKeyParseError::Empty),);
375    }
376
377    #[test]
378    fn test_restrict_key_too_long() {
379        let key: String = std::iter::repeat_n('a', 64).collect();
380
381        assert_eq!(
382            key.parse::<RestrictKey>(),
383            Err(RestrictKeyParseError::TooLong),
384        );
385    }
386
387    #[test]
388    fn test_restrict_key_max_length() {
389        let key: String = std::iter::repeat_n('a', 63).collect();
390
391        assert!(key.parse::<RestrictKey>().is_ok());
392    }
393
394    #[test]
395    fn test_restrict_key_non_alphanumeric() {
396        assert_eq!(
397            "abc-123".parse::<RestrictKey>(),
398            Err(RestrictKeyParseError::NotAlphanumeric),
399        );
400    }
401
402    #[test]
403    fn test_restrict_key_const() {
404        const KEY: RestrictKey = RestrictKey::from_static_or_panic("abc123");
405        assert_eq!(KEY.as_str(), "abc123");
406    }
407
408    #[test]
409    fn test_combined() {
410        assert_eq!(
411            PgSchemaDump::new()
412                .exclude_schema("internal".parse().unwrap())
413                .exclude_table(QualifiedTable {
414                    schema: Schema::PUBLIC,
415                    table: "temp".parse().unwrap(),
416                })
417                .no_owner()
418                .no_privileges()
419                .verbose()
420                .arguments(),
421            vec![
422                "--schema-only",
423                "--exclude-schema",
424                "internal",
425                "--exclude-table",
426                "public.temp",
427                "--no-owner",
428                "--no-privileges",
429                "--verbose",
430            ],
431        );
432    }
433}
pg_client/pg_dump.rs

pg_client/
pg_dump.rs
