Implement tablespace remapping for barman-cloud-restore

Each tablespace has a name, an oid, and a location (directory), and
there is a symbolic link from $PGDATA/pg_tblspc/$oid to the location.
barman-cloud-backup uploads a $oid.tar file for each tablespace, and
records the properties of each tablespace in backup_info.tablespaces.

Until now, barman-cloud-restore knew only how to download $oid.tar and
extract its contents to the original location. Here we add a couple of
new features:

1. You may optionally specify `--tablespace name:/new/location` to
   extract a given tablespace to a different location (you can repeat
   this option multiple times)

2. After downloading and extracting $oid.tar to the desired location, we
   now recreate the symbolic link from $PGDATA/pg_tblspc/$oid to that
   directory.

Closes #343
Closes #330

barman/clients/cloud_restore.py

@@ -71,7 +71,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")
@@ -133,6 +137,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",
@@ -162,6 +173,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
@@ -180,7 +211,7 @@ 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
 
@@ -198,19 +229,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(
@@ -241,6 +289,9 @@ 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.)

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)