Merge pull request #2 from siemens/feat/restore/allow-destination

feat(restore): Allow providing restore destination

Reviewed-By : anupam.ghosh@siemens.com

src/snap_to_bucket/run.py

@@ -23,7 +23,7 @@ def main(args):
     snap_to_bucket = SnapToBucket(args.bucket, args.tag, args.verbose,
                                   args.type, args.storage_class, args.mount,
                                   args.delete, args.restore, args.key,
-                                  args.boot)
+                                  args.boot, args.restore_dir)
     snap_to_bucket.update_proxy(args.proxy, args.noproxy)
     snap_to_bucket.update_split_size(args.split)
     if args.gzip:
@@ -74,7 +74,7 @@ def entrypoint():
     parser = argparse.ArgumentParser(description='''
     snap_to_bucket is a simple tool based on boto3 to move snapshots to S3
     buckets
-    ''', epilog='use delete with caution')
+    ''')
     parser.add_argument("-v", "--verbose", help="increase output verbosity " +
                         "(-vvv for more verbosity)", action="count", default=0)
     parser.add_argument("-b", "--bucket", help="S3 bucket to push snaps in",
@@ -100,8 +100,8 @@ def entrypoint():
                         "(default: %(default)s)", required=False,
                         metavar="DIR", default="/mnt/snaps")
     parser.add_argument("-d", "--delete", help="delete snapshot after " +
-                        "transfer (default: %(default)s)", required=False,
-                        action="store_true", default=False)
+                        "transfer. Use with caution! (default: %(default)s)",
+                        required=False, action="store_true", default=False)
     parser.add_argument("-s", "--split", help="split tar in chunks no bigger " +
                         "than (allowed suffix b,k,m,g,t) (default: %(default)s)",
                         metavar="SIZE", required=False, default="5t",
@@ -115,6 +115,9 @@ def entrypoint():
                         required=False)
     parser.add_argument("--boot", help="was the snapshot a bootable volume?",
                         action="store_true", default=False, required=False)
+    parser.add_argument("--restore-dir", help="directory to store S3 objects " +
+                        "for restoring (default: %(default)s)",
+                        default="/tmp/snap-to-bucket", required=False)
     args = parser.parse_args()
 
     if os.geteuid() != 0:

src/snap_to_bucket/s3_handler.py

@@ -12,7 +12,6 @@
 import os
 import sys
 import base64
-import shutil
 import hashlib
 import threading
 from datetime import datetime
@@ -344,7 +343,7 @@ def get_object_count_and_size(self, key):
         return (self.__get_object_count(key),
                 self.restore_partition_size)
 
-    def download_key(self, key, partno):
+    def download_key(self, key, partno, restore_dir):
         """
         Download the key from S3
 
@@ -355,6 +354,8 @@ def download_key(self, key, partno):
         :param partno: Part number of the key to be downloaded (-1 if there is
             only one part)
         :type partno: integer
+        :param restore_dir: Location to store S3 object for restore
+        :type restore_dir: string
 
         :return: Location of downloaded file and size of restored partition (in
             bytes)
@@ -365,28 +366,27 @@ def download_key(self, key, partno):
         response = self.s3client.list_objects_v2(Bucket=self.bucket,
                                                  Prefix=key)
         keys = [o['Key'] for o in response['Contents']]
-        download_key = None
+        download_key_name = None
         if partno == -1:
-            download_key = keys[0]
+            download_key_name = keys[0]
         else:
             for key in keys:
                 if f"-part{partno}.tar" in key:
-                    download_key = key
+                    download_key_name = key
                     break
-        if download_key == None:
+        if download_key_name == None:
             raise Exception(f"Unable to part '{partno}' under key {key}")
-        self.temp_download = f"/tmp/snap-to-bucket/{download_key}"
+        self.temp_download = os.path.join(restore_dir, download_key_name)
         size = self.s3client.head_object(Bucket=self.bucket,
-                                         Key=download_key)['ContentLength']
-        os.makedirs(os.path.dirname(self.temp_download), exist_ok=True)
+                                         Key=download_key_name)['ContentLength']
         progress = ProgressPercentage(key, size)
         try:
-            self.s3client.download_file(self.bucket, download_key,
+            self.s3client.download_file(self.bucket, download_key_name,
                                         self.temp_download, Callback=progress)
             print()
         except Exception as e:
-            print(f"Failed while downloading s3://{self.bucket}/{download_key}",
+            print(f"Failed while downloading s3://{self.bucket}/{download_key_name}",
                   file=sys.stderr)
-            shutil.rmtree("/tmp/snap-to-bucket")
+            os.remove(self.temp_download)
             raise e
         return self.temp_download

src/snap_to_bucket/snap_handler.py

@@ -32,12 +32,14 @@ class SnapToBucket:
     :ivar restore_boot: Was the snapshot, being restored, bootable?
     :ivar split_size: Size in bytes to split tar at
     :ivar gzip: True to compress tar with gzip
+    :ivar restore_dir: Location to store S3 object for restore
     """
 
     def __init__(self, bucket, tag="snap-to-bucket", verbose=0,
                  volume_type="gp2", storage_class="STANDARD",
                  mount_point="/mnt/snaps", delete_snap=False, restore=False,
-                 restore_key="", restore_boot=False):
+                 restore_key="", restore_boot=False,
+                 restore_dir="/tmp/snap-to-bucket"):
         """
         Initializer for the class attributes.
 
@@ -61,6 +63,8 @@ def __init__(self, bucket, tag="snap-to-bucket", verbose=0,
         :type restore_key: string
         :param restore_boot: Was the snapshot, being restored, bootable?
         :type restore_boot: boolean
+        :param restore_dir: Location to store S3 object for restore
+        :type restore_dir: string
         """
         self.__bucket = bucket
         self.__tag = tag
@@ -80,6 +84,7 @@ def __init__(self, bucket, tag="snap-to-bucket", verbose=0,
         self.__restore = restore
         self.__restore_key = restore_key
         self.__restore_boot = restore_boot
+        self.__restore_dir = restore_dir
         self.__split_size = 5 * 1024.0 * 1024.0 * 1024.0 * 1024.0
         self.__gzip = False
 
@@ -133,6 +138,9 @@ def initiate_migration(self):
         if self.__restore == True:
             if self.__restore_key == None:
                 raise Exception("missing key argument for restore")
+            os.makedirs(self.__restore_dir, exist_ok=True)
+            if not os.access(self.__restore_dir, os.W_OK):
+                raise Exception(f"Directory {self.__restore_dir} is not writeable")
             self.__restore_snapshot()
             return
 
@@ -206,12 +214,12 @@ def __restore_snapshot(self):
             self.__fshandler.mount_volume(volume_id)
             if no_of_objects == 1:
                 temp_path = self.__s3handler.download_key(self.__restore_key,
-                                                          -1)
+                                                        -1, self.__restore_dir)
                 self.__fshandler.untar(temp_path)
             else:
                 for i in range(1, no_of_objects + 1):
                     temp_path = self.__s3handler.download_key(
-                        self.__restore_key, i)
+                        self.__restore_key, i, self.__restore_dir)
                     self.__fshandler.untar(temp_path)
             self.__fshandler.terminate_tar()
             if self.__restore_boot: