Recreate pg_wal if required in barman-cloud-restore

barman/clients/cloud_restore.py

@@ -65,7 +65,11 @@ def main(args=None):
                 logging.error("Bucket %s does not exist", cloud_interface.bucket_name)
                 raise SystemExit(1)
 
-            downloader.download_backup(config.backup_id, config.recovery_dir)
+            downloader.download_backup(
+                config.backup_id,
+                config.recovery_dir,
+                tablespace_map(config.tablespace),
+            )
 
     except KeyboardInterrupt as exc:
         logging.error("Barman cloud restore was interrupted by the user")
@@ -127,6 +131,13 @@ def parse_arguments(args=None):
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--tablespace",
+        help="tablespace relocation rule",
+        metavar="NAME:LOCATION",
+        action="append",
+        default=[],
+    )
     parser.add_argument(
         "--cloud-provider",
         help="The cloud provider to use as a storage backend",
@@ -148,6 +159,26 @@ def parse_arguments(args=None):
     return parser.parse_args(args=args)
 
 
+def tablespace_map(rules):
+    """
+    Return a mapping from tablespace names to locations built from any
+    `--tablespace name:/loc/ation` rules specified.
+    """
+    tablespaces = {}
+    for rule in rules:
+        try:
+            tablespaces.update([rule.split(":", 1)])
+        except ValueError:
+            logging.error(
+                "Invalid tablespace relocation rule '%s'\n"
+                "HINT: The valid syntax for a relocation rule is "
+                "NAME:LOCATION",
+                rule,
+            )
+            raise SystemExit(1)
+    return tablespaces
+
+
 class CloudBackupDownloader(object):
     """
     Cloud storage download client
@@ -166,12 +197,12 @@ def __init__(self, cloud_interface, server_name):
         self.server_name = server_name
         self.catalog = CloudBackupCatalog(cloud_interface, server_name)
 
-    def download_backup(self, backup_id, destination_dir):
+    def download_backup(self, backup_id, destination_dir, tablespaces):
         """
         Download a backup from cloud storage
 
-        :param str wal_name: Name of the WAL file
-        :param str wal_dest: Full path of the destination WAL file
+        :param str backup_id: The backup id to restore
+        :param str destination_dir: Path to the destination directory
         """
 
         backup_info = self.catalog.get_backup_info(backup_id)
@@ -184,19 +215,36 @@ def download_backup(self, backup_id, destination_dir):
 
         backup_files = self.catalog.get_backup_files(backup_info)
 
-        # Check that everything is ok
+        # We must download and restore a bunch of .tar files that contain PGDATA
+        # and each tablespace. First, we determine a target directory to extract
+        # each tar file into and record these in copy_jobs. For each tablespace,
+        # the location may be overriden by `--tablespace name:/new/location` on
+        # the command-line; and we must also add an entry to link_jobs to create
+        # a symlink from $PGDATA/pg_tblspc/oid to the correct location after the
+        # downloads.
+
         copy_jobs = []
+        link_jobs = []
         for oid in backup_files:
             file_info = backup_files[oid]
             # PGDATA is restored where requested (destination_dir)
             if oid is None:
                 target_dir = destination_dir
             else:
-                # Tablespaces are restored in the original location
-                # TODO: implement tablespace remapping
                 for tblspc in backup_info.tablespaces:
                     if oid == tblspc.oid:
                         target_dir = tblspc.location
+                        if tblspc.name in tablespaces:
+                            target_dir = os.path.realpath(tablespaces[tblspc.name])
+                        logging.debug(
+                            "Tablespace %s (oid=%s) will be located at %s",
+                            tblspc.name,
+                            oid,
+                            target_dir,
+                        )
+                        link_jobs.append(
+                            ["%s/pg_tblspc/%s" % (destination_dir, oid), target_dir]
+                        )
                         break
                 else:
                     raise AssertionError(
@@ -227,6 +275,17 @@ def download_backup(self, backup_id, destination_dir):
             )
             self.cloud_interface.extract_tar(file_info.path, target_dir)
 
+        for link, target in link_jobs:
+            os.symlink(target, link)
+
+        # If we did not restore the pg_wal directory from one of the uploaded
+        # backup files, we must recreate it here. (If pg_wal was originally a
+        # symlink, it would not have been uploaded.)
+
+        wal_path = os.path.join(destination_dir, backup_info.wal_directory())
+        if not os.path.exists(wal_path):
+            os.mkdir(wal_path)
+
 
 if __name__ == "__main__":
     main()

barman/infofile.py

@@ -584,6 +584,13 @@ def pg_major_version(self):
         else:
             return str(major)
 
+    def wal_directory(self):
+        """
+        Returns "pg_wal" (v10 and above) or "pg_xlog" (v9.6 and below) based on
+        the Postgres version represented by this backup
+        """
+        return "pg_wal" if self.version >= 100000 else "pg_xlog"
+
 
 class LocalBackupInfo(BackupInfo):
     __slots__ = "server", "config", "backup_manager"

doc/barman-cloud-restore.1

@@ -44,6 +44,10 @@ show program\[cq]s version number and exit
 -t, \[en]test
 test connectivity to the cloud destination and exit
 .TP
+\[en]tablespace NAME:LOCATION
+extract the named tablespace to the given directory instead of its
+original location (you may repeat the option for multiple tablespaces)
+.TP
 -P, \[en]profile
 profile name (e.g.\ INI section in AWS credentials file)
 .TP

doc/barman-cloud-restore.1.md

@@ -49,6 +49,10 @@ RECOVERY_DIR
 -t, --test
 : test connectivity to the cloud destination and exit
 
+--tablespace NAME:LOCATION
+: extract the named tablespace to the given directory instead of its
+original location (you may repeat the option for multiple tablespaces)
+
 -P, --profile
 : profile name (e.g. INI section in AWS credentials file)
 

tests/test_infofile.py

@@ -594,6 +594,8 @@ def test_pg_version(self, tmpdir):
         b_info = LocalBackupInfo(server, info_file=infofile.strpath)
         # BASE_BACKUP_INFO has version 90400 so expect 9.4
         assert b_info.pg_major_version() == "9.4"
+        assert b_info.wal_directory() == "pg_xlog"
         # Set backup_info.version to 100600 so expect 10
         b_info.version = 100600
         assert b_info.pg_major_version() == "10"
+        assert b_info.wal_directory() == "pg_wal"