From 348bc3c3435f69a825e8345d0e3e426ddca97eb7 Mon Sep 17 00:00:00 2001
From: alisa <alisa@alisaqaq.moe>
Date: Fri, 8 Mar 2024 09:33:26 +0800
Subject: [PATCH] feat: support --role options in pg_dump

---
 example/database/postgresql.yaml   | 26 +++++++++++++++++++-------
 src/configs/client/pg_dump.rs      |  1 +
 src/configs/database/postgresql.rs |  2 ++
 src/provider/pg_dump.rs            | 13 +++++++++++++
 4 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/example/database/postgresql.yaml b/example/database/postgresql.yaml
index a576940..eb36243 100644
--- a/example/database/postgresql.yaml
+++ b/example/database/postgresql.yaml
@@ -30,6 +30,12 @@ postgresql:
       - "Data"
       - "PostData"
 
+    # Set role before dumping
+    #   Can be omitted
+    #   Can be overridden by database specific settings
+    # --role={} option in pg_dump
+    ## role: "my_role"
+
     # Whether to create the database before restoring
     #   Can be omitted
     #   Can be overridden by database specific settings
@@ -76,7 +82,7 @@ postgresql:
 
       # Format of the backup (Plain, Custom, Tar)
       #   Can be omitted
-      #   Will override the pg_dump.format option
+      #   Will override the client.format option
       # Note that the directory type is not supported
       # --format={} option in pg_dump
       ## format: "Plain"
@@ -103,28 +109,34 @@ postgresql:
 
       # Section to include in the backup (PreData, Data, PostData)
       #   Can be omitted
-      #   Will override the pg_dump.sections option
+      #   Will override the client.sections option
       # --section={} option in pg_dump
       ## sections:
       ##   - "PreData"
       ##   - "Data"
       ##   - "PostData"
 
+      # Set role before dumping
+      #   Can be omitted
+      #   Will override the client.role option
+      # --role={} option in pg_dump
+      ## role: "my_role"
+
       # Whether to create the database before restoring
       #   Can be omitted
-      #   Will override the pg_dump.create option
+      #   Will override the client.create option
       # --create option in pg_dump
       ## create: false
 
       # Whether to drop the database before restoring
       #   Can be omitted
-      #   Will override the pg_dump.clean option
+      #   Will override the client.clean option
       # --clean option in pg_dump
       ## clean: false
 
       # Do not save the corresponding data
       #   Can be omitted
-      #   Will override the pg_dump.do_not_save option
+      #   Will override the client.do_not_save option
       # Available values:
       #   LargeObjects, Owner, Privileges, Comments,
       #   Publications, SecurityLabels, Subscriptions,
@@ -135,7 +147,7 @@ postgresql:
 
       # Disable something
       #   Can be omitted
-      #   Will override the pg_dump.disable option
+      #   Will override the client.disable option
       # Available values:
       #   DollarQuoting, Triggers
       # --disable={} option in pg_dump
@@ -143,5 +155,5 @@ postgresql:
 
       # Extra arguments for pg_dump
       #   Can be omitted
-      #   Will override the pg_dump.extra_args option
+      #   Will override the client.extra_args option
       ## extra_args: []
diff --git a/src/configs/client/pg_dump.rs b/src/configs/client/pg_dump.rs
index 1c2772e..4f0ad3b 100644
--- a/src/configs/client/pg_dump.rs
+++ b/src/configs/client/pg_dump.rs
@@ -11,6 +11,7 @@ pub struct PgDump {
     pub binary_path: String,
     pub format: PgFormat,
     pub sections: Vec<PgSection>,
+    pub role: Option<String>,
     pub clean: Option<bool>,
     pub create: Option<bool>,
     pub do_not_save: Option<Vec<PgDoNotSave>>,
diff --git a/src/configs/database/postgresql.rs b/src/configs/database/postgresql.rs
index 6918265..a8d6de5 100644
--- a/src/configs/database/postgresql.rs
+++ b/src/configs/database/postgresql.rs
@@ -21,6 +21,7 @@ pub struct PostgreSQLDatabase {
     pub exclude_schemata: Option<Vec<String>>,
     pub include_tables: Option<Vec<String>>,
     pub exclude_tables: Option<Vec<String>>,
+    pub role: Option<String>,
     pub clean: Option<bool>,
     pub create: Option<bool>,
     pub sections: Option<Vec<PgSection>>,
@@ -38,6 +39,7 @@ impl From<&String> for PostgreSQLDatabase {
             exclude_schemata: None,
             include_tables: None,
             exclude_tables: None,
+            role: None,
             clean: None,
             create: None,
             sections: None,
diff --git a/src/provider/pg_dump.rs b/src/provider/pg_dump.rs
index 6c73d2b..f3f072b 100644
--- a/src/provider/pg_dump.rs
+++ b/src/provider/pg_dump.rs
@@ -141,6 +141,19 @@ impl Dump<PgDump, PostgreSQLConnection> for PostgreSQLDatabase {
         apply_selective_args_some(cmd, "schema", &self.include_schemata);
         apply_selective_args_some(cmd, "exclude-schema", &self.exclude_schemata);
 
+        // role
+        match &self.role {
+            Some(val) => {
+                cmd.arg(format!("--role={role}", role = val));
+            }
+            None => match &pg_dump.role {
+                Some(val) => {
+                    cmd.arg(format!("--role={role}", role = val));
+                }
+                None => (),
+            },
+        }
+
         // extra_args
         match &self.extra_args {
             Some(val) => {