Update single table restore method

config/config.yaml

@@ -21,8 +21,8 @@ project_settings:
   # (it usually is the same project on which BBQ runs)
   backup_project_id: 'BBQ-project-id'
 
-  # restoration_project_id - project into which data will be restored during restoration process
-  restoration_project_id: 'restoration-storage-project-id'
+  # default_restoration_project_id - project into which data will be restored by default during restoration process
+  default_restoration_project_id: 'default-restoration-storage-project-id'
   debug_mode: False
 
 tests:

src/commons/config/configuration.py

@@ -39,8 +39,8 @@ def backup_project_id(self):
         return self.__project_config['project_settings'].get('backup_project_id')
 
     @property
-    def restoration_project_id(self):
-        return self.__project_config['project_settings'].get('restoration_project_id')
+    def default_restoration_project_id(self):
+        return self.__project_config['project_settings'].get('default_restoration_project_id')
 
     @property
     def projects_to_skip(self):

src/commons/handlers/validators.py

@@ -0,0 +1,84 @@
+import re
+
+from src.commons.exceptions import ParameterValidationException
+
+
+class WrongDatasetNameException(Exception):
+    pass
+
+
+class WrongProjectNameException(Exception):
+    pass
+
+
+class WrongWriteDispositionException(Exception):
+    pass
+
+
+class WrongCreateDispositionException(Exception):
+    pass
+
+
+project_id_pattern = re.compile("^[a-zA-Z0-9-]+$")
+dataset_id_pattern = re.compile("^[a-zA-Z0-9_]+$")
+available_create_dispositions = ["CREATE_IF_NEEDED", "CREATE_NEVER"]
+available_write_dispositions = ["WRITE_TRUNCATE", "WRITE_APPEND", "WRITE_EMPTY"]
+
+
+def validate_restore_request_params(
+        source_project_id=None, source_dataset_id=None,
+        target_project_id=None, target_dataset_id=None,
+        create_disposition=None, write_disposition=None):
+    try:
+        if source_project_id:
+            validate_project_id(source_project_id)
+        if source_dataset_id:
+            validate_dataset_id(source_dataset_id)
+        if target_project_id:
+            validate_project_id(target_project_id)
+        if target_dataset_id:
+            validate_dataset_id(target_dataset_id)
+        if write_disposition:
+            validate_write_disposition(write_disposition)
+        if create_disposition:
+            validate_create_disposition(create_disposition)
+
+    except (WrongDatasetNameException,
+            WrongProjectNameException,
+            WrongWriteDispositionException,
+            WrongCreateDispositionException), e:
+
+        raise ParameterValidationException(e.message)
+
+
+def validate_project_id(project_id):
+    if not project_id or not project_id_pattern.match(project_id):
+        error_message = "Invalid project value: '{}'. Project IDs may " \
+                        "contain letters, numbers, and " \
+                        "dash".format(project_id)
+        raise WrongProjectNameException(error_message)
+
+
+def validate_dataset_id(dataset_id):
+    if not dataset_id or not dataset_id_pattern.match(dataset_id):
+        error_message = "Invalid dataset value: '{}'. Dataset IDs may " \
+                        "contain letters, numbers, and " \
+                        "underscores".format(dataset_id)
+        raise WrongDatasetNameException(error_message)
+
+
+def validate_write_disposition(write_disposition):
+    if write_disposition not in available_write_dispositions:
+        error_message = "Invalid write disposition: '{}'. " \
+                        "The following values are supported: {}." \
+            .format(write_disposition, ', '.join(available_write_dispositions))
+        raise WrongWriteDispositionException(error_message)
+
+
+def validate_create_disposition(create_disposition):
+    if create_disposition not in available_create_dispositions:
+        error_message = "Invalid create disposition: '{}'. " \
+                        "The following values are supported: {}." \
+            .format(create_disposition,
+                    ', '.join(available_create_dispositions))
+        raise WrongCreateDispositionException(error_message)

src/restore/dataset/dataset_restore_handler.py

@@ -3,12 +3,11 @@
 
 import webapp2
 
-from src.commons.exceptions import ParameterValidationException
-from src.commons.handlers.json_handler import JsonHandler
-from src.commons.handlers.bbq_authenticated_handler import BbqAuthenticatedHandler
-from src.commons.big_query import validators
-from src.commons.big_query.validators import WrongDatasetNameException
+from src.commons.handlers import validators
 from src.commons.config.configuration import configuration
+from src.commons.handlers.bbq_authenticated_handler import \
+    BbqAuthenticatedHandler
+from src.commons.handlers.json_handler import JsonHandler
 from src.restore.dataset.dataset_restore_service import \
     DatasetRestoreService
 from src.restore.status.restoration_job_status_service import \
@@ -18,11 +17,20 @@
 class DatasetRestoreHandler(JsonHandler):
 
     def post(self, project_id, dataset_id):
-
+        target_project_id = self.request.get('targetProjectId', None)
         target_dataset_id = self.request.get('targetDatasetId', None)
+        create_disposition = self.request.get('createDisposition', None)
+        write_disposition = self.request.get('writeDisposition', None)
         max_partition_days = self.__get_max_partition_days()
 
-        self.__validate_params(dataset_id, target_dataset_id)
+        validators.validate_restore_request_params(
+            source_project_id=project_id,
+            source_dataset_id=dataset_id,
+            target_project_id=target_project_id,
+            target_dataset_id=target_dataset_id,
+            create_disposition=create_disposition,
+            write_disposition=write_disposition
+        )
 
         restoration_job_id = str(uuid.uuid4())
         logging.info("Created restoration_job_id: %s", restoration_job_id)
@@ -31,34 +39,26 @@ def post(self, project_id, dataset_id):
             restoration_job_id=restoration_job_id,
             project_id=project_id,
             dataset_id=dataset_id,
+            target_project_id=target_project_id,
             target_dataset_id=target_dataset_id,
+            create_disposition=create_disposition,
+            write_disposition=write_disposition,
             max_partition_days=max_partition_days
         )
 
         restore_data = {
             'restorationJobId': restoration_job_id,
             'projectId': project_id,
             'datasetId': dataset_id,
-            'restorationStatusEndpoint': RestorationJobStatusService.get_status_endpoint(
-                restoration_job_id),
-            'restorationWarningsOnlyStatusEndpoint': RestorationJobStatusService.get_warnings_only_status_endpoint(
-                restoration_job_id)
+            'restorationStatusEndpoint': RestorationJobStatusService.get_status_endpoint(restoration_job_id),
+            'restorationWarningsOnlyStatusEndpoint': RestorationJobStatusService.get_warnings_only_status_endpoint(restoration_job_id)
         }
         self._finish_with_success(restore_data)
 
     def __get_max_partition_days(self):
         max_partition_days = self.request.get('maxPartitionDays', None)
         return int(max_partition_days) if max_partition_days else None
 
-    @staticmethod
-    def __validate_params(dataset_id, target_dataset_id):
-        try:
-            validators.validate_dataset_id(dataset_id)
-            if target_dataset_id:
-                validators.validate_dataset_id(target_dataset_id)
-        except WrongDatasetNameException, e:
-            raise ParameterValidationException(e.message)
-
 
 class DatasetRestoreAuthenticatedHandler(DatasetRestoreHandler,
                                          BbqAuthenticatedHandler):
@@ -69,8 +69,10 @@ def __init__(self, request=None, response=None):
 
 
 app = webapp2.WSGIApplication([
-    webapp2.Route('/restore/project/<project_id:.*>/dataset/<dataset_id:.*>',
-     DatasetRestoreHandler),
-    webapp2.Route('/schedule/restore/project/<project_id:.*>/dataset/<dataset_id:.*>',
-     DatasetRestoreAuthenticatedHandler)
+    webapp2.Route(
+        '/restore/project/<project_id:.*>/dataset/<dataset_id:.*>',
+        DatasetRestoreHandler),
+    webapp2.Route(
+        '/schedule/restore/project/<project_id:.*>/dataset/<dataset_id:.*>',
+        DatasetRestoreAuthenticatedHandler)
 ], debug=configuration.debug_mode)

src/restore/dataset/dataset_restore_items_generator.py

@@ -1,8 +1,9 @@
 import logging
 
+from src.backup.datastore.Table import Table
 from src.commons.collections import paginated
 from src.commons.decorators.log_time import log_time
-from src.backup.datastore.Table import Table
+from src.commons.table_reference import TableReference
 from src.restore.datastore.restore_item import RestoreItem
 from src.restore.restoration_table_reference import RestoreTableReference
 
@@ -11,8 +12,8 @@ class DatasetRestoreItemsGenerator(object):
 
     @classmethod
     @log_time
-    def generate_restore_items(cls, project_id, dataset_id, target_dataset_id,
-                               max_partition_days):
+    def generate_restore_items(cls, project_id, dataset_id, target_project_id,
+                               target_dataset_id, max_partition_days):
         if max_partition_days:
             table_entities = Table \
                 .get_tables_with_max_partition_days(project_id, dataset_id,
@@ -29,13 +30,16 @@ def generate_restore_items(cls, project_id, dataset_id, target_dataset_id,
                         RestoreTableReference.backup_table_reference(
                             table_entity, backup_entity)
 
-                    target_table_reference = \
-                        RestoreTableReference.target_table_reference(
-                            table_entity, target_dataset_id)
+                    target_table_reference = TableReference(
+                        target_project_id,
+                        target_dataset_id,
+                        table_entity.table_id,
+                        table_entity.partition_id
+                    )
 
-                    restore_items.append(
-                        RestoreItem.create(source_table_reference,
-                                           target_table_reference))
+                    restore_item = RestoreItem.create(source_table_reference,
+                                                      target_table_reference)
+                    restore_items.append(restore_item)
 
             logging.info("Restore items generator yields %s restore items",
                          len(restore_items))

src/restore/dataset/dataset_restore_parameters_validator.py

@@ -1,23 +1,26 @@
 import logging
 
-from src.commons.exceptions import ParameterValidationException
 from src.backup.datastore.Table import Table
 from src.commons.big_query.big_query import BigQuery
 from src.commons.config.configuration import configuration
+from src.commons.exceptions import ParameterValidationException
 
 
 class DatasetRestoreParametersValidator(object):
 
     def __init__(self):
         self.BQ = BigQuery()
 
-    def validate_parameters(self, project_id, dataset_id, target_dataset_id,
-                            max_partition_days):
+    def validate_parameters(self, project_id, dataset_id, target_project_id,
+            target_dataset_id, max_partition_days):
+
+        if target_project_id is None:
+            raise ParameterValidationException("Required target project id parameter is None")
 
         any_backup = self.__get_backup(project_id, dataset_id,
                                        max_partition_days)
 
-        self.__validate_locations(any_backup, target_dataset_id)
+        self.__validate_locations(any_backup, target_project_id, target_dataset_id)
 
     def __get_backup(self, project_id, dataset_id, max_partition_days):
         logging.info(
@@ -53,10 +56,10 @@ def __get_tables(self, project_id, dataset_id, max_partition_days):
                                                             max_partition_days,
                                                             page_size=20)
 
-    def __get_target_dataset_location(self, target_dataset_id):
+    def __get_target_dataset_location(self, target_project_id, target_dataset_id):
         target_dataset = self.BQ.get_dataset(
-            configuration.restoration_project_id,
-            target_dataset_id)
+            project_id=target_project_id,
+            dataset_id=target_dataset_id)
         if target_dataset is None:
             return None
         return self.__extract_location(target_dataset)
@@ -68,8 +71,8 @@ def __get_backup_dataset_location(self, any_backup):
             return None
         return self.__extract_location(backup_dataset)
 
-    def __validate_locations(self, any_backup, target_dataset_id):
-        target_location = self.__get_target_dataset_location(target_dataset_id)
+    def __validate_locations(self, any_backup, target_project_id, target_dataset_id):
+        target_location = self.__get_target_dataset_location(target_project_id,target_dataset_id)
 
         if target_location is None:
             return