Skip to main content

postgresql_commands/
pg_restore.rs

1use crate::Settings;
2use crate::traits::CommandBuilder;
3use std::convert::AsRef;
4use std::ffi::{OsStr, OsString};
5use std::path::PathBuf;
6
7/// `pg_restore` restores a `PostgreSQL` database from an archive created by `pg_dump`.
8#[derive(Clone, Debug, Default)]
9pub struct PgRestoreBuilder {
10    program_dir: Option<PathBuf>,
11    envs: Vec<(OsString, OsString)>,
12    dbname: Option<OsString>,
13    file: Option<OsString>,
14    format: Option<OsString>,
15    list: bool,
16    verbose: bool,
17    version: bool,
18    help: bool,
19    data_only: bool,
20    clean: bool,
21    create: bool,
22    exit_on_error: bool,
23    index: Option<OsString>,
24    jobs: Option<OsString>,
25    use_list: Option<OsString>,
26    schema: Option<OsString>,
27    exclude_schema: Option<OsString>,
28    no_owner: bool,
29    function: Option<OsString>,
30    schema_only: bool,
31    superuser: Option<OsString>,
32    table: Option<OsString>,
33    trigger: Option<OsString>,
34    no_privileges: bool,
35    single_transaction: bool,
36    disable_triggers: bool,
37    enable_row_security: bool,
38    if_exists: bool,
39    no_comments: bool,
40    no_data_for_failed_tables: bool,
41    no_publications: bool,
42    no_security_labels: bool,
43    no_subscriptions: bool,
44    no_table_access_method: bool,
45    no_tablespaces: bool,
46    section: Option<OsString>,
47    strict_names: bool,
48    use_set_session_authorization: bool,
49    host: Option<OsString>,
50    port: Option<u16>,
51    username: Option<OsString>,
52    no_password: bool,
53    password: bool,
54    pg_password: Option<OsString>,
55    role: Option<OsString>,
56}
57
58impl PgRestoreBuilder {
59    /// Create a new [`PgRestoreBuilder`]
60    #[must_use]
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Create a new [`PgRestoreBuilder`] from [Settings]
66    pub fn from(settings: &dyn Settings) -> Self {
67        let mut builder = Self::new()
68            .program_dir(settings.get_binary_dir())
69            .host(settings.get_host())
70            .port(settings.get_port())
71            .username(settings.get_username())
72            .pg_password(settings.get_password());
73        if let Some(socket_dir) = settings.get_socket_dir() {
74            builder = builder.host(socket_dir.to_string_lossy().to_string());
75        }
76        builder
77    }
78
79    /// Location of the program binary
80    #[must_use]
81    pub fn program_dir<P: Into<PathBuf>>(mut self, path: P) -> Self {
82        self.program_dir = Some(path.into());
83        self
84    }
85
86    /// connect to database name
87    #[must_use]
88    pub fn dbname<S: AsRef<OsStr>>(mut self, name: S) -> Self {
89        self.dbname = Some(name.as_ref().to_os_string());
90        self
91    }
92
93    /// output file name (- for stdout)
94    #[must_use]
95    pub fn file<S: AsRef<OsStr>>(mut self, filename: S) -> Self {
96        self.file = Some(filename.as_ref().to_os_string());
97        self
98    }
99
100    /// backup file format (should be automatic)
101    #[must_use]
102    pub fn format<S: AsRef<OsStr>>(mut self, format: S) -> Self {
103        self.format = Some(format.as_ref().to_os_string());
104        self
105    }
106
107    /// print summarized TOC of the archive
108    #[must_use]
109    pub fn list(mut self) -> Self {
110        self.list = true;
111        self
112    }
113
114    /// verbose mode
115    #[must_use]
116    pub fn verbose(mut self) -> Self {
117        self.verbose = true;
118        self
119    }
120
121    /// output version information, then exit
122    #[must_use]
123    pub fn version(mut self) -> Self {
124        self.version = true;
125        self
126    }
127
128    /// show help, then exit
129    #[must_use]
130    pub fn help(mut self) -> Self {
131        self.help = true;
132        self
133    }
134
135    /// restore only the data, no schema
136    #[must_use]
137    pub fn data_only(mut self) -> Self {
138        self.data_only = true;
139        self
140    }
141
142    /// clean (drop) database objects before recreating
143    #[must_use]
144    pub fn clean(mut self) -> Self {
145        self.clean = true;
146        self
147    }
148
149    /// create the target database
150    #[must_use]
151    pub fn create(mut self) -> Self {
152        self.create = true;
153        self
154    }
155
156    /// exit on error, default is to continue
157    #[must_use]
158    pub fn exit_on_error(mut self) -> Self {
159        self.exit_on_error = true;
160        self
161    }
162
163    /// restore named index
164    #[must_use]
165    pub fn index<S: AsRef<OsStr>>(mut self, name: S) -> Self {
166        self.index = Some(name.as_ref().to_os_string());
167        self
168    }
169
170    /// use this many parallel jobs to restore
171    #[must_use]
172    pub fn jobs<S: AsRef<OsStr>>(mut self, num: S) -> Self {
173        self.jobs = Some(num.as_ref().to_os_string());
174        self
175    }
176
177    /// use table of contents from this file for selecting/ordering output
178    #[must_use]
179    pub fn use_list<S: AsRef<OsStr>>(mut self, filename: S) -> Self {
180        self.use_list = Some(filename.as_ref().to_os_string());
181        self
182    }
183
184    /// restore only objects in this schema
185    #[must_use]
186    pub fn schema<S: AsRef<OsStr>>(mut self, name: S) -> Self {
187        self.schema = Some(name.as_ref().to_os_string());
188        self
189    }
190
191    /// do not restore objects in this schema
192    #[must_use]
193    pub fn exclude_schema<S: AsRef<OsStr>>(mut self, name: S) -> Self {
194        self.exclude_schema = Some(name.as_ref().to_os_string());
195        self
196    }
197
198    /// skip restoration of object ownership
199    #[must_use]
200    pub fn no_owner(mut self) -> Self {
201        self.no_owner = true;
202        self
203    }
204
205    /// restore named function
206    #[must_use]
207    pub fn function<S: AsRef<OsStr>>(mut self, name: S) -> Self {
208        self.function = Some(name.as_ref().to_os_string());
209        self
210    }
211
212    /// restore only the schema, no data
213    #[must_use]
214    pub fn schema_only(mut self) -> Self {
215        self.schema_only = true;
216        self
217    }
218
219    /// superuser user name to use for disabling triggers
220    #[must_use]
221    pub fn superuser<S: AsRef<OsStr>>(mut self, name: S) -> Self {
222        self.superuser = Some(name.as_ref().to_os_string());
223        self
224    }
225
226    /// restore named relation (table, view, etc.)
227    #[must_use]
228    pub fn table<S: AsRef<OsStr>>(mut self, name: S) -> Self {
229        self.table = Some(name.as_ref().to_os_string());
230        self
231    }
232
233    /// restore named trigger
234    #[must_use]
235    pub fn trigger<S: AsRef<OsStr>>(mut self, name: S) -> Self {
236        self.trigger = Some(name.as_ref().to_os_string());
237        self
238    }
239
240    /// skip restoration of access privileges (grant/revoke)
241    #[must_use]
242    pub fn no_privileges(mut self) -> Self {
243        self.no_privileges = true;
244        self
245    }
246
247    /// restore as a single transaction
248    #[must_use]
249    pub fn single_transaction(mut self) -> Self {
250        self.single_transaction = true;
251        self
252    }
253
254    /// disable triggers during data-only restore
255    #[must_use]
256    pub fn disable_triggers(mut self) -> Self {
257        self.disable_triggers = true;
258        self
259    }
260
261    /// enable row security
262    #[must_use]
263    pub fn enable_row_security(mut self) -> Self {
264        self.enable_row_security = true;
265        self
266    }
267
268    /// use IF EXISTS when dropping objects
269    #[must_use]
270    pub fn if_exists(mut self) -> Self {
271        self.if_exists = true;
272        self
273    }
274
275    /// do not restore comments
276    #[must_use]
277    pub fn no_comments(mut self) -> Self {
278        self.no_comments = true;
279        self
280    }
281
282    /// do not restore data of tables that could not be created
283    #[must_use]
284    pub fn no_data_for_failed_tables(mut self) -> Self {
285        self.no_data_for_failed_tables = true;
286        self
287    }
288
289    /// do not restore publications
290    #[must_use]
291    pub fn no_publications(mut self) -> Self {
292        self.no_publications = true;
293        self
294    }
295
296    /// do not restore security labels
297    #[must_use]
298    pub fn no_security_labels(mut self) -> Self {
299        self.no_security_labels = true;
300        self
301    }
302
303    /// do not restore subscriptions
304    #[must_use]
305    pub fn no_subscriptions(mut self) -> Self {
306        self.no_subscriptions = true;
307        self
308    }
309
310    /// do not restore table access methods
311    #[must_use]
312    pub fn no_table_access_method(mut self) -> Self {
313        self.no_table_access_method = true;
314        self
315    }
316
317    /// do not restore tablespace assignments
318    #[must_use]
319    pub fn no_tablespaces(mut self) -> Self {
320        self.no_tablespaces = true;
321        self
322    }
323
324    /// restore named section (pre-data, data, or post-data)
325    #[must_use]
326    pub fn section<S: AsRef<OsStr>>(mut self, section: S) -> Self {
327        self.section = Some(section.as_ref().to_os_string());
328        self
329    }
330
331    /// require table and/or schema include patterns to match at least one entity each
332    #[must_use]
333    pub fn strict_names(mut self) -> Self {
334        self.strict_names = true;
335        self
336    }
337
338    /// use SET SESSION AUTHORIZATION commands instead of ALTER OWNER commands to set ownership
339    #[must_use]
340    pub fn use_set_session_authorization(mut self) -> Self {
341        self.use_set_session_authorization = true;
342        self
343    }
344
345    /// database server host or socket directory
346    #[must_use]
347    pub fn host<S: AsRef<OsStr>>(mut self, hostname: S) -> Self {
348        self.host = Some(hostname.as_ref().to_os_string());
349        self
350    }
351
352    /// database server port number
353    #[must_use]
354    pub fn port(mut self, port: u16) -> Self {
355        self.port = Some(port);
356        self
357    }
358
359    /// connect as specified database user
360    #[must_use]
361    pub fn username<S: AsRef<OsStr>>(mut self, name: S) -> Self {
362        self.username = Some(name.as_ref().to_os_string());
363        self
364    }
365
366    /// never prompt for password
367    #[must_use]
368    pub fn no_password(mut self) -> Self {
369        self.no_password = true;
370        self
371    }
372
373    /// force password prompt (should happen automatically)
374    #[must_use]
375    pub fn password(mut self) -> Self {
376        self.password = true;
377        self
378    }
379
380    /// user password
381    #[must_use]
382    pub fn pg_password<S: AsRef<OsStr>>(mut self, pg_password: S) -> Self {
383        self.pg_password = Some(pg_password.as_ref().to_os_string());
384        self
385    }
386
387    /// do SET ROLE before restore
388    #[must_use]
389    pub fn role<S: AsRef<OsStr>>(mut self, rolename: S) -> Self {
390        self.role = Some(rolename.as_ref().to_os_string());
391        self
392    }
393}
394
395impl CommandBuilder for PgRestoreBuilder {
396    /// Get the program name
397    fn get_program(&self) -> &'static OsStr {
398        "pg_restore".as_ref()
399    }
400
401    /// Location of the program binary
402    fn get_program_dir(&self) -> &Option<PathBuf> {
403        &self.program_dir
404    }
405
406    /// Get the arguments for the command
407    #[expect(clippy::too_many_lines)]
408    fn get_args(&self) -> Vec<OsString> {
409        let mut args: Vec<OsString> = Vec::new();
410
411        if let Some(name) = &self.dbname {
412            args.push("--dbname".into());
413            args.push(name.into());
414        }
415
416        if let Some(filename) = &self.file {
417            args.push("--file".into());
418            args.push(filename.into());
419        }
420
421        if let Some(format) = &self.format {
422            args.push("--format".into());
423            args.push(format.into());
424        }
425
426        if self.list {
427            args.push("--list".into());
428        }
429
430        if self.verbose {
431            args.push("--verbose".into());
432        }
433
434        if self.version {
435            args.push("--version".into());
436        }
437
438        if self.help {
439            args.push("--help".into());
440        }
441
442        if self.data_only {
443            args.push("--data-only".into());
444        }
445
446        if self.clean {
447            args.push("--clean".into());
448        }
449
450        if self.create {
451            args.push("--create".into());
452        }
453
454        if self.exit_on_error {
455            args.push("--exit-on-error".into());
456        }
457
458        if let Some(name) = &self.index {
459            args.push("--index".into());
460            args.push(name.into());
461        }
462
463        if let Some(num) = &self.jobs {
464            args.push("--jobs".into());
465            args.push(num.into());
466        }
467
468        if let Some(filename) = &self.use_list {
469            args.push("--use-list".into());
470            args.push(filename.into());
471        }
472
473        if let Some(name) = &self.schema {
474            args.push("--schema".into());
475            args.push(name.into());
476        }
477
478        if let Some(name) = &self.exclude_schema {
479            args.push("--exclude-schema".into());
480            args.push(name.into());
481        }
482
483        if self.no_owner {
484            args.push("--no-owner".into());
485        }
486
487        if let Some(name) = &self.function {
488            args.push("--function".into());
489            args.push(name.into());
490        }
491
492        if self.schema_only {
493            args.push("--schema-only".into());
494        }
495
496        if let Some(name) = &self.superuser {
497            args.push("--superuser".into());
498            args.push(name.into());
499        }
500
501        if let Some(name) = &self.table {
502            args.push("--table".into());
503            args.push(name.into());
504        }
505
506        if let Some(name) = &self.trigger {
507            args.push("--trigger".into());
508            args.push(name.into());
509        }
510
511        if self.no_privileges {
512            args.push("--no-privileges".into());
513        }
514
515        if self.single_transaction {
516            args.push("--single-transaction".into());
517        }
518
519        if self.disable_triggers {
520            args.push("--disable-triggers".into());
521        }
522
523        if self.enable_row_security {
524            args.push("--enable-row-security".into());
525        }
526
527        if self.if_exists {
528            args.push("--if-exists".into());
529        }
530
531        if self.no_comments {
532            args.push("--no-comments".into());
533        }
534
535        if self.no_data_for_failed_tables {
536            args.push("--no-data-for-failed-tables".into());
537        }
538
539        if self.no_publications {
540            args.push("--no-publications".into());
541        }
542
543        if self.no_security_labels {
544            args.push("--no-security-labels".into());
545        }
546
547        if self.no_subscriptions {
548            args.push("--no-subscriptions".into());
549        }
550
551        if self.no_table_access_method {
552            args.push("--no-table-access-method".into());
553        }
554
555        if self.no_tablespaces {
556            args.push("--no-tablespaces".into());
557        }
558
559        if let Some(section) = &self.section {
560            args.push("--section".into());
561            args.push(section.into());
562        }
563
564        if self.strict_names {
565            args.push("--strict-names".into());
566        }
567
568        if self.use_set_session_authorization {
569            args.push("--use-set-session-authorization".into());
570        }
571
572        if let Some(hostname) = &self.host {
573            args.push("--host".into());
574            args.push(hostname.into());
575        }
576
577        if let Some(port) = &self.port {
578            args.push("--port".into());
579            args.push(port.to_string().into());
580        }
581
582        if let Some(name) = &self.username {
583            args.push("--username".into());
584            args.push(name.into());
585        }
586
587        if self.no_password {
588            args.push("--no-password".into());
589        }
590
591        if self.password {
592            args.push("--password".into());
593        }
594
595        if let Some(role) = &self.role {
596            args.push("--role".into());
597            args.push(role.into());
598        }
599
600        args
601    }
602
603    /// Get the environment variables for the command
604    fn get_envs(&self) -> Vec<(OsString, OsString)> {
605        let mut envs: Vec<(OsString, OsString)> = self.envs.clone();
606
607        if let Some(password) = &self.pg_password {
608            envs.push(("PGPASSWORD".into(), password.into()));
609        }
610
611        envs
612    }
613
614    /// Set an environment variable for the command
615    fn env<S: AsRef<OsStr>>(mut self, key: S, value: S) -> Self {
616        self.envs
617            .push((key.as_ref().to_os_string(), value.as_ref().to_os_string()));
618        self
619    }
620}
621
622#[cfg(test)]
623mod tests {
624    use super::*;
625    use crate::TestSettings;
626    use crate::TestSocketSettings;
627    use crate::traits::CommandToString;
628    use test_log::test;
629
630    #[test]
631    fn test_builder_new() {
632        let command = PgRestoreBuilder::new().program_dir(".").build();
633        assert_eq!(
634            PathBuf::from(".").join("pg_restore"),
635            PathBuf::from(command.to_command_string().replace('"', ""))
636        );
637    }
638
639    #[test]
640    fn test_builder_from() {
641        let command = PgRestoreBuilder::from(&TestSettings).build();
642        #[cfg(not(target_os = "windows"))]
643        let command_prefix = r#"PGPASSWORD="password" "./pg_restore" "#;
644        #[cfg(target_os = "windows")]
645        let command_prefix = r#"".\\pg_restore" "#;
646
647        assert_eq!(
648            format!(
649                r#"{command_prefix}"--host" "localhost" "--port" "5432" "--username" "postgres""#
650            ),
651            command.to_command_string()
652        );
653    }
654
655    #[test]
656    fn test_builder_from_socket() {
657        let command = PgRestoreBuilder::from(&TestSocketSettings).build();
658        #[cfg(not(target_os = "windows"))]
659        let command_prefix = r#"PGPASSWORD="password" "./pg_restore" "#;
660        #[cfg(target_os = "windows")]
661        let command_prefix = r#"".\\pg_restore" "#;
662        assert_eq!(
663            format!(
664                r#"{command_prefix}"--host" "/tmp/pg_socket" "--port" "5432" "--username" "postgres""#
665            ),
666            command.to_command_string()
667        );
668    }
669    #[test]
670    fn test_builder() {
671        let command = PgRestoreBuilder::new()
672            .env("PGDATABASE", "database")
673            .dbname("dbname")
674            .file("file")
675            .format("format")
676            .list()
677            .verbose()
678            .version()
679            .help()
680            .data_only()
681            .clean()
682            .create()
683            .exit_on_error()
684            .index("index")
685            .jobs("jobs")
686            .use_list("use_list")
687            .schema("schema")
688            .exclude_schema("exclude_schema")
689            .no_owner()
690            .function("function")
691            .schema_only()
692            .superuser("superuser")
693            .table("table")
694            .trigger("trigger")
695            .no_privileges()
696            .single_transaction()
697            .disable_triggers()
698            .enable_row_security()
699            .if_exists()
700            .no_comments()
701            .no_data_for_failed_tables()
702            .no_publications()
703            .no_security_labels()
704            .no_subscriptions()
705            .no_table_access_method()
706            .no_tablespaces()
707            .section("section")
708            .strict_names()
709            .use_set_session_authorization()
710            .host("localhost")
711            .port(5432)
712            .username("username")
713            .no_password()
714            .password()
715            .pg_password("password")
716            .role("role")
717            .build();
718        #[cfg(not(target_os = "windows"))]
719        let command_prefix = r#"PGDATABASE="database" PGPASSWORD="password" "#;
720        #[cfg(target_os = "windows")]
721        let command_prefix = String::new();
722
723        assert_eq!(
724            format!(
725                r#"{command_prefix}"pg_restore" "--dbname" "dbname" "--file" "file" "--format" "format" "--list" "--verbose" "--version" "--help" "--data-only" "--clean" "--create" "--exit-on-error" "--index" "index" "--jobs" "jobs" "--use-list" "use_list" "--schema" "schema" "--exclude-schema" "exclude_schema" "--no-owner" "--function" "function" "--schema-only" "--superuser" "superuser" "--table" "table" "--trigger" "trigger" "--no-privileges" "--single-transaction" "--disable-triggers" "--enable-row-security" "--if-exists" "--no-comments" "--no-data-for-failed-tables" "--no-publications" "--no-security-labels" "--no-subscriptions" "--no-table-access-method" "--no-tablespaces" "--section" "section" "--strict-names" "--use-set-session-authorization" "--host" "localhost" "--port" "5432" "--username" "username" "--no-password" "--password" "--role" "role""#
726            ),
727            command.to_command_string()
728        );
729    }
730}