From 7e5475549c76a3dd8c49448ef2daf8514f37ae29 Mon Sep 17 00:00:00 2001
From: theandrew168 <andrew@shallowbrooksoftware.com>
Date: Mon, 26 Sep 2022 23:15:03 -0500
Subject: [PATCH] Add conf option to restore specific schemas

---
 README.md         |  3 ++-
 pg2s3.conf        |  1 +
 pg2s3.go          | 49 +++++++++++++++++++++++++++++++++--------------
 script/pg2s3.conf |  7 +++++--
 4 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index ac06dfc..7a9fa6e 100644
--- a/README.md
+++ b/README.md
@@ -36,9 +36,10 @@ The following settings are required to run pg2s3:
 | `s3_access_key_id`     | Yes       | S3-compatible storage access key ID |
 | `s3_secret_access_key` | Yes       | S3-compatible storage secret access key |
 | `s3_bucket_name`       | Yes       | S3-compatible storage bucket name |
-| `backup_prefix`        | Yes       | Prefix attached to the name of each backup |
+| `backup_prefix`        | No        | Prefix attached to the name of each backup (default "pg2s3") |
 | `backup_retention`     | No        | Number of backups to retain after pruning |
 | `backup_schedule`      | No        | Backup schedule as a standard cron expression (UTC) |
+| `restore_schemas`      | No        | List of schemas to restore (default "public") |
 
 ## Encryption
 Backups managed by pg2s3 can be optionally encrypted using [age](https://github.com/FiloSottile/age).
diff --git a/pg2s3.conf b/pg2s3.conf
index eb3dcf4..2ab60dc 100644
--- a/pg2s3.conf
+++ b/pg2s3.conf
@@ -6,4 +6,5 @@ s3_bucket_name = "pg2s3"
 backup_prefix = "pg2s3"
 backup_retention = 30
 backup_schedule = "0 9 * * *"
+restore_schemas = ["public"]
 age_public_key = "age156hm5jvxfvf8xf0zjs52gc5hhq64rt23gw3fehqj2vu77sk07a5qvplj52"
diff --git a/pg2s3.go b/pg2s3.go
index 2a78990..2538f3e 100644
--- a/pg2s3.go
+++ b/pg2s3.go
@@ -21,24 +21,42 @@ import (
 )
 
 type Config struct {
-	PGConnectionURI   string `toml:"pg_connection_uri"`
-	S3Endpoint        string `toml:"s3_endpoint"`
-	S3AccessKeyID     string `toml:"s3_access_key_id"`
-	S3SecretAccessKey string `toml:"s3_secret_access_key"`
-	S3BucketName      string `toml:"s3_bucket_name"`
-	BackupPrefix      string `toml:"backup_prefix"`
-	BackupRetention   int    `toml:"backup_retention"`
-	BackupSchedule    string `toml:"backup_schedule"`
-	AgePublicKey      string `toml:"age_public_key"`
+	PGConnectionURI   string   `toml:"pg_connection_uri"`
+	S3Endpoint        string   `toml:"s3_endpoint"`
+	S3AccessKeyID     string   `toml:"s3_access_key_id"`
+	S3SecretAccessKey string   `toml:"s3_secret_access_key"`
+	S3BucketName      string   `toml:"s3_bucket_name"`
+	BackupPrefix      string   `toml:"backup_prefix"`
+	BackupRetention   int      `toml:"backup_retention"`
+	BackupSchedule    string   `toml:"backup_schedule"`
+	RestoreSchemas    []string `toml:"restore_schemas"`
+	AgePublicKey      string   `toml:"age_public_key"`
 }
 
 func ReadConfig(path string) (Config, error) {
-	var cfg Config
+	// init Config struct with default values
+	cfg := Config{
+		BackupPrefix:   "pg2s3",
+		RestoreSchemas: []string{"public"},
+	}
 	meta, err := toml.DecodeFile(path, &cfg)
 	if err != nil {
 		return Config{}, err
 	}
 
+	// gather extra values
+	var extra []string
+	for _, keys := range meta.Undecoded() {
+		key := keys[0]
+		extra = append(extra, key)
+	}
+
+	// error upon extra values
+	if len(extra) > 0 {
+		msg := strings.Join(extra, ", ")
+		return Config{}, fmt.Errorf("extra config values: %s", msg)
+	}
+
 	// build set of present config keys
 	present := make(map[string]bool)
 	for _, keys := range meta.Keys() {
@@ -52,7 +70,6 @@ func ReadConfig(path string) (Config, error) {
 		"s3_access_key_id",
 		"s3_secret_access_key",
 		"s3_bucket_name",
-		"backup_prefix",
 	}
 
 	// ensure required keys are present
@@ -147,7 +164,7 @@ func New(cfg Config) (*Client, error) {
 
 func (c *Client) CreateBackup() (io.Reader, error) {
 	args := []string{
-		"-Fc",
+		"-Fc", // custom output format (compressed and flexible)
 		c.cfg.PGConnectionURI,
 	}
 	cmd := exec.Command("pg_dump", args...)
@@ -169,10 +186,14 @@ func (c *Client) CreateBackup() (io.Reader, error) {
 
 func (c *Client) RestoreBackup(backup io.Reader) error {
 	args := []string{
-		"-c",
-		"-d",
+		"-c", // clean DB object before recreating them
+		"-d", // database to be restored
 		c.cfg.PGConnectionURI,
 	}
+	for _, schema := range c.cfg.RestoreSchemas {
+		// specify which schemas should be restored
+		args = append(args, "-n", schema)
+	}
 	cmd := exec.Command("pg_restore", args...)
 
 	cmd.Stdin = backup
diff --git a/script/pg2s3.conf b/script/pg2s3.conf
index 4c38363..c040496 100644
--- a/script/pg2s3.conf
+++ b/script/pg2s3.conf
@@ -13,8 +13,8 @@
 # REQUIRED - S3-compatible storage bucket name
 #s3_bucket_name = ""
 
-# REQUIRED - Prefix attached to the name of each backup
-#backup_prefix = ""
+# OPTIONAL - Prefix attached to the name of each backup (default "pg2s3")
+#backup_prefix = "pg2s3"
 
 # OPTIONAL - Number of backups to retain after pruning
 #backup_retention = 0
@@ -22,5 +22,8 @@
 # OPTIONAL - Backup schedule as a standard cron expression (UTC)
 #backup_schedule = ""
 
+# OPTIONAL - List of schemas to restore (default ["public"])
+#restore_schemas = ["public"]
+
 # OPTIONAL - Public key for backup encryption
 #age_public_key = ""