Merge pull request #18 from GoogleCloudPlatform/tomerlf-snapshot-fr

Tomerlf snapshot fr

gce_rescue/bin/rescue.py

@@ -68,11 +68,12 @@ def main():
 
     print('Restoring VM...')
     action = 'reset_rescue_mode'
-    msg = messages.tip_restore_disk(vm)
+    has_snapshot = vm.snapshot
+    msg = messages.tip_restore_disk(vm, snapshot=has_snapshot)
 
   call_tasks(vm=vm, action=action)
   print(msg)
 
 
 if __name__ == '__main__':
-  main()
+  main()

gce_rescue/config.py

@@ -24,6 +24,7 @@
 config = {
   'version': VERSION,
   'debug': False,
+  'skip-snapshot': False,
   'startup-script-file': os.path.join(dirname, 'startup-script.txt'),
   'source_guests': {
     'x86_64':[
@@ -56,9 +57,11 @@ def process_args():
                       help='Print to the log file in debug leve')
   parser.add_argument('-f', '--force', action='store_true',
                       help='Don\'t ask for confirmation.')
-
+  parser.add_argument('--skip-snapshot', action='store_true',
+                      help='Skip backing up the disk using a snapshot.')
   return parser
 
 
 def set_configs(user_args):
   config['debug'] = getattr(user_args, 'debug')
+  config['skip-snapshot'] = getattr(user_args, 'skip_snapshot')

gce_rescue/gce.py

@@ -20,7 +20,7 @@
 from typing import Dict, List, Union
 from time import time
 from gce_rescue.tasks.backup import backup_metadata_items
-from gce_rescue.tasks.disks import list_disk
+from gce_rescue.tasks.disks import list_disk, list_snapshot
 from gce_rescue.tasks.pre_validations import Validations
 from gce_rescue.config import get_config
 
@@ -70,7 +70,7 @@ def validate_instance_mode(data: Dict) -> Dict:
       'rescue-mode': False,
       'ts': generate_ts()
   }
-  if 'metadata' in data and  'items' in data['metadata']:
+  if 'metadata' in data and 'items' in data['metadata']:
     metadata = data['metadata']
     for item in metadata['items']:
       if item['key'] == 'rescue-mode':
@@ -216,3 +216,10 @@ def backup_items(self, v: List[str]) -> None:
   @property
   def disks(self) -> List[str]:
     return self._disks
+
+  @property
+  def snapshot(self) -> str:
+    if not self.rescue_mode_status['rescue-mode']:
+      return f"{self.disks['disk_name']}-{self.ts}"
+    return list_snapshot(self)
+

gce_rescue/messages.py

@@ -26,9 +26,14 @@ def tip_connect_ssh(vm: Instance) -> str:
     f'{vm.zone}/instances/{vm.name}?authuser=0&hl=en_US&useAdminProxy=true&'
     f'troubleshoot4005Enabled=true\n')
 
-def tip_restore_disk(vm: Instance) -> str:
-  return (f'└── The instance {vm.name} was restored! Use the snapshot below '
+def tip_restore_disk(vm: Instance, snapshot=False) -> str:
+  if not snapshot:
+    snapshot_restore_msg = ''
+  else:
+    snapshot_restore_msg = (f' Use the snapshot below '
     f'if you need to restore the modification made while the instance was '
     f'in rescue mode.\n Snapshot name: {vm.disks["disk_name"]}-{vm.ts}\n'
     f' More information: '
     f'https://cloud.google.com/compute/docs/disks/restore-snapshot\n')
+
+  return f'└── The instance {vm.name} was restored!' + snapshot_restore_msg

gce_rescue/tasks/actions.py

@@ -19,7 +19,8 @@
 
 from gce_rescue.gce import Instance
 from gce_rescue.tasks.disks import (
-  config_rescue_disks,
+  take_snapshot,
+  create_rescue_disk,
   restore_original_disk,
   attach_disk
 )
@@ -32,7 +33,7 @@
   restore_metadata_items
 )
 from gce_rescue.utils import Tracker
-
+from gce_rescue.config import get_config
 _logger = logging.getLogger(__name__)
 
 def _list_tasks(vm: Instance, action: str) -> List:
@@ -50,7 +51,13 @@ def _list_tasks(vm: Instance, action: str) -> List:
         }]
       },
       {
-        'name': config_rescue_disks,
+        'name': take_snapshot,
+        'args': [{
+          'vm': vm
+        }]
+      },
+      {
+        'name': create_rescue_disk,
         'args': [{
           'vm': vm
         }]
@@ -120,6 +127,9 @@ def _list_tasks(vm: Instance, action: str) -> List:
 def call_tasks(vm: Instance, action: str) -> None:
   """ Loop tasks dict and execute """
   tasks = _list_tasks(vm = vm, action = action)
+  if get_config('skip-snapshot'):
+    _logger.info(f'Skipping snapshot backup.')
+    tasks = [task for task in tasks if task['name'].__name__ != 'take_snapshot']
   total_tasks = len(tasks)
 
   tracker = Tracker(total_tasks)

gce_rescue/tasks/backup.py

@@ -30,7 +30,7 @@ def backup_metadata_items(data: Dict) -> List:
     return data['metadata']['items']
   return []
 
-def _create_snapshot(vm) -> Dict:
+def create_snapshot(vm) -> Dict:
   """
   Create a snaphost of the instance boot disk, adding self._ts to the disk name.
   https://cloud.google.com/compute/docs/reference/rest/v1/disks/createSnapshot
@@ -51,8 +51,3 @@ def _create_snapshot(vm) -> Dict:
   result = wait_for_operation(vm, oper=operation)
   return result
 
-def backup(vm) -> None:
-  """
-  List of methods to backup data and information from the orignal instance
-  """
-  _create_snapshot(vm)

gce_rescue/tasks/backup_test.py

@@ -39,7 +39,7 @@ def test_backup_metadata_items(self):
 
   def test_backup(self):
     """Test backup task."""
-    backup.backup(self.vm)
+    backup.create_snapshot(self.vm)
 
 
 if __name__ == '__main__':

gce_rescue/tasks/disks.py

@@ -20,8 +20,9 @@
 import googleapiclient.errors
 
 from gce_rescue.tasks.keeper import wait_for_operation
-from gce_rescue.tasks.backup import  backup
+from gce_rescue.tasks.backup import create_snapshot
 from gce_rescue.utils import ThreadHandler as Handler
+from googleapiclient.errors import HttpError
 
 _logger = logging.getLogger(__name__)
 
@@ -177,9 +178,12 @@ def _detach_disk(vm, disk: str) -> Dict:
   return result
 
 
-def config_rescue_disks(vm) -> None:
+def take_snapshot(vm) -> None:
+  create_snapshot(vm)
+
+
+def create_rescue_disk(vm) -> None:
   device_name = vm.disks['device_name']
-  backup(vm)
   # task1 = multitasks.Handler(
   #     target = backup,
   #     kwargs={'vm' : vm}
@@ -199,6 +203,18 @@ def config_rescue_disks(vm) -> None:
     boot=True
   )
 
+def list_snapshot(vm) -> str:
+  snapshot_name = f"{vm.disks['disk_name']}-{vm.ts}"
+  try:
+    result = vm.compute.snapshots().get(
+      snapshot=snapshot_name,
+      project=vm.project
+    ).execute()
+  except HttpError:
+    _logger.info('Snapshot was not found for VM in active rescue mode')
+    return ''
+  return snapshot_name
+
 def restore_original_disk(vm) -> None:
   """ Restore tasks to the original disk """
   device_name = vm.disks['device_name']

gce_rescue/tasks/disks_test.py

@@ -73,7 +73,7 @@ def test_config_rescue_disks(self):
       'operations',
       'disks',
     ])
-    disks.config_rescue_disks(self.vm)
+    disks.create_rescue_disk(self.vm)
 
 
   def test_restore_original_disk(self):