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 AsRef<str> for RestrictKey {
69    fn as_ref(&self) -> &str {
70        &self.0
71    }
72}
73
74impl Display for RestrictKey {
75    fn fmt(&self, formatter: &mut Formatter<'_>) -> core::fmt::Result {
76        formatter.write_str(&self.0)
77    }
78}
79
80impl FromStr for RestrictKey {
81    type Err = RestrictKeyParseError;
82
83    fn from_str(input: &str) -> Result<Self, Self::Err> {
84        match validate_restrict_key(input) {
85            Some(error) => Err(error),
86            None => Ok(Self(Cow::Owned(input.to_owned()))),
87        }
88    }
89}
90
91/// Error parsing a restrict key.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum RestrictKeyParseError {
94    /// Key cannot be empty.
95    Empty,
96    /// Key exceeds maximum length.
97    TooLong,
98    /// Key contains non-alphanumeric bytes.
99    NotAlphanumeric,
100}
101
102impl RestrictKeyParseError {
103    /// Returns the error message.
104    #[must_use]
105    pub const fn message(&self) -> &'static str {
106        match self {
107            Self::Empty => "restrict key cannot be empty",
108            Self::TooLong => "restrict key exceeds 63 byte max length",
109            Self::NotAlphanumeric => "restrict key must be alphanumeric",
110        }
111    }
112}
113
114impl Display for RestrictKeyParseError {
115    fn fmt(&self, formatter: &mut Formatter<'_>) -> core::fmt::Result {
116        formatter.write_str(self.message())
117    }
118}
119
120impl std::error::Error for RestrictKeyParseError {}
121
122#[must_use]
123pub struct PgSchemaDump {
124    exclude_schemas: Vec<Schema>,
125    exclude_tables: Vec<QualifiedTable>,
126    no_comments: bool,
127    no_owner: bool,
128    no_privileges: bool,
129    no_tablespaces: bool,
130    restrict_key: Option<RestrictKey>,
131    schemas: Vec<Schema>,
132    tables: Vec<QualifiedTable>,
133    verbose: bool,
134}
135
136impl Default for PgSchemaDump {
137    fn default() -> Self {
138        Self::new()
139    }
140}
141
142impl PgSchemaDump {
143    pub fn new() -> Self {
144        Self {
145            exclude_schemas: Vec::new(),
146            exclude_tables: Vec::new(),
147            no_comments: false,
148            no_owner: false,
149            no_privileges: false,
150            no_tablespaces: false,
151            restrict_key: None,
152            schemas: Vec::new(),
153            tables: Vec::new(),
154            verbose: false,
155        }
156    }
157
158    pub fn exclude_schema(mut self, schema: Schema) -> Self {
159        self.exclude_schemas.push(schema);
160        self
161    }
162
163    pub fn exclude_table(mut self, table: QualifiedTable) -> Self {
164        self.exclude_tables.push(table);
165        self
166    }
167
168    pub fn no_comments(mut self) -> Self {
169        self.no_comments = true;
170        self
171    }
172
173    pub fn no_owner(mut self) -> Self {
174        self.no_owner = true;
175        self
176    }
177
178    pub fn no_privileges(mut self) -> Self {
179        self.no_privileges = true;
180        self
181    }
182
183    pub fn no_tablespaces(mut self) -> Self {
184        self.no_tablespaces = true;
185        self
186    }
187
188    pub fn restrict_key(mut self, restrict_key: RestrictKey) -> Self {
189        self.restrict_key = Some(restrict_key);
190        self
191    }
192
193    pub fn schema(mut self, schema: Schema) -> Self {
194        self.schemas.push(schema);
195        self
196    }
197
198    pub fn table(mut self, table: QualifiedTable) -> Self {
199        self.tables.push(table);
200        self
201    }
202
203    pub fn verbose(mut self) -> Self {
204        self.verbose = true;
205        self
206    }
207
208    #[must_use]
209    pub fn arguments(&self) -> Vec<String> {
210        let mut args = vec!["--schema-only".to_string()];
211
212        for schema in &self.exclude_schemas {
213            args.push("--exclude-schema".to_string());
214            args.push(schema.to_string());
215        }
216
217        for table in &self.exclude_tables {
218            args.push("--exclude-table".to_string());
219            args.push(table.to_string());
220        }
221
222        if self.no_comments {
223            args.push("--no-comments".to_string());
224        }
225
226        if self.no_owner {
227            args.push("--no-owner".to_string());
228        }
229
230        if self.no_privileges {
231            args.push("--no-privileges".to_string());
232        }
233
234        if self.no_tablespaces {
235            args.push("--no-tablespaces".to_string());
236        }
237
238        if let Some(restrict_key) = &self.restrict_key {
239            args.push(format!("--restrict-key={restrict_key}"));
240        }
241
242        for schema in &self.schemas {
243            args.push("--schema".to_string());
244            args.push(schema.to_string());
245        }
246
247        for table in &self.tables {
248            args.push("--table".to_string());
249            args.push(table.to_string());
250        }
251
252        if self.verbose {
253            args.push("--verbose".to_string());
254        }
255
256        args
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn test_default() {
266        assert_eq!(PgSchemaDump::new().arguments(), vec!["--schema-only"]);
267    }
268
269    #[test]
270    fn test_verbose() {
271        assert_eq!(
272            PgSchemaDump::new().verbose().arguments(),
273            vec!["--schema-only", "--verbose"],
274        );
275    }
276
277    #[test]
278    fn test_exclude_schema() {
279        assert_eq!(
280            PgSchemaDump::new()
281                .exclude_schema("internal".parse().unwrap())
282                .arguments(),
283            vec!["--schema-only", "--exclude-schema", "internal"],
284        );
285    }
286
287    #[test]
288    fn test_exclude_table() {
289        assert_eq!(
290            PgSchemaDump::new()
291                .exclude_table(QualifiedTable {
292                    schema: Schema::PUBLIC,
293                    table: "some_table".parse().unwrap(),
294                })
295                .arguments(),
296            vec!["--schema-only", "--exclude-table", "public.some_table"],
297        );
298    }
299
300    #[test]
301    fn test_no_comments() {
302        assert_eq!(
303            PgSchemaDump::new().no_comments().arguments(),
304            vec!["--schema-only", "--no-comments"],
305        );
306    }
307
308    #[test]
309    fn test_no_owner() {
310        assert_eq!(
311            PgSchemaDump::new().no_owner().arguments(),
312            vec!["--schema-only", "--no-owner"],
313        );
314    }
315
316    #[test]
317    fn test_no_privileges() {
318        assert_eq!(
319            PgSchemaDump::new().no_privileges().arguments(),
320            vec!["--schema-only", "--no-privileges"],
321        );
322    }
323
324    #[test]
325    fn test_no_tablespaces() {
326        assert_eq!(
327            PgSchemaDump::new().no_tablespaces().arguments(),
328            vec!["--schema-only", "--no-tablespaces"],
329        );
330    }
331
332    #[test]
333    fn test_schema() {
334        assert_eq!(
335            PgSchemaDump::new().schema(Schema::PUBLIC).arguments(),
336            vec!["--schema-only", "--schema", "public"],
337        );
338    }
339
340    #[test]
341    fn test_table() {
342        assert_eq!(
343            PgSchemaDump::new()
344                .table(QualifiedTable {
345                    schema: Schema::PUBLIC,
346                    table: "users".parse().unwrap(),
347                })
348                .arguments(),
349            vec!["--schema-only", "--table", "public.users"],
350        );
351    }
352
353    #[test]
354    fn test_restrict_key() {
355        assert_eq!(
356            PgSchemaDump::new()
357                .restrict_key("abc123".parse().unwrap())
358                .arguments(),
359            vec!["--schema-only", "--restrict-key=abc123"],
360        );
361    }
362
363    #[test]
364    fn test_restrict_key_empty() {
365        assert_eq!("".parse::<RestrictKey>(), Err(RestrictKeyParseError::Empty),);
366    }
367
368    #[test]
369    fn test_restrict_key_too_long() {
370        let key: String = std::iter::repeat_n('a', 64).collect();
371
372        assert_eq!(
373            key.parse::<RestrictKey>(),
374            Err(RestrictKeyParseError::TooLong),
375        );
376    }
377
378    #[test]
379    fn test_restrict_key_max_length() {
380        let key: String = std::iter::repeat_n('a', 63).collect();
381
382        assert!(key.parse::<RestrictKey>().is_ok());
383    }
384
385    #[test]
386    fn test_restrict_key_non_alphanumeric() {
387        assert_eq!(
388            "abc-123".parse::<RestrictKey>(),
389            Err(RestrictKeyParseError::NotAlphanumeric),
390        );
391    }
392
393    #[test]
394    fn test_restrict_key_const() {
395        const KEY: RestrictKey = RestrictKey::from_static_or_panic("abc123");
396        assert_eq!(KEY.as_str(), "abc123");
397    }
398
399    #[test]
400    fn test_combined() {
401        assert_eq!(
402            PgSchemaDump::new()
403                .exclude_schema("internal".parse().unwrap())
404                .exclude_table(QualifiedTable {
405                    schema: Schema::PUBLIC,
406                    table: "temp".parse().unwrap(),
407                })
408                .no_owner()
409                .no_privileges()
410                .verbose()
411                .arguments(),
412            vec![
413                "--schema-only",
414                "--exclude-schema",
415                "internal",
416                "--exclude-table",
417                "public.temp",
418                "--no-owner",
419                "--no-privileges",
420                "--verbose",
421            ],
422        );
423    }
424}
pg_client/pg_dump.rs

pg_client/
pg_dump.rs
