Merge 8ca86b2 into 42f83d4

src/commons/handlers/validators.py

@@ -25,26 +25,7 @@ class WrongCreateDispositionException(Exception):
 available_write_dispositions = ["WRITE_TRUNCATE", "WRITE_APPEND", "WRITE_EMPTY"]
 
 
-def validate_dataset_restore_params(source_project_id, source_dataset_id,
-                                    target_project_id, target_dataset_id,
-                                    create_disposition, write_disposition):
-    __validate_restore_request_params(source_project_id=source_project_id,
-                                      source_dataset_id=source_dataset_id,
-                                      target_project_id=target_project_id,
-                                      target_dataset_id=target_dataset_id,
-                                      create_disposition=create_disposition,
-                                      write_disposition=write_disposition)
-
-
-def validate_list_restore_params(target_project_id, target_dataset_id,
-                                 create_disposition, write_disposition):
-    __validate_restore_request_params(target_project_id=target_project_id,
-                                      target_dataset_id=target_dataset_id,
-                                      create_disposition=create_disposition,
-                                      write_disposition=write_disposition)
-
-
-def __validate_restore_request_params(
+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):

src/restore/dataset/dataset_restore_handler.py

@@ -23,7 +23,7 @@ def post(self, project_id, dataset_id):
         write_disposition = self.request.get('writeDisposition', None)
         max_partition_days = self.__get_max_partition_days()
 
-        validators.validate_dataset_restore_params(
+        validators.validate_restore_request_params(
             source_project_id=project_id,
             source_dataset_id=dataset_id,
             target_project_id=target_project_id,

src/restore/list/backup_list_restore_handler.py

@@ -27,9 +27,11 @@ def post(self):
         write_disposition = self.request.get('writeDisposition', None)
         create_disposition = self.request.get('createDisposition', None)
 
-        validators.validate_list_restore_params(
-            target_project_id, target_dataset_id,
-            create_disposition, write_disposition
+        validators.validate_restore_request_params(
+            target_project_id=target_project_id,
+            target_dataset_id=target_dataset_id,
+            create_disposition=create_disposition,
+            write_disposition=write_disposition
         )
 
         body_json = self.__parse_body_json()

src/restore/list/backup_list_restore_service.py

@@ -5,7 +5,6 @@
 
 from src.backup.datastore.Table import Table
 from src.commons.collections import paginated
-from src.commons.config.configuration import configuration
 from src.commons.decorators.log_time import log_time
 from src.commons.decorators.retry import retry
 from src.commons.exceptions import ParameterValidationException
@@ -36,7 +35,8 @@ def __str__(self):
 
 class BackupListRestoreRequest(object):
     def __init__(self, backup_items, target_project_id, target_dataset_id,
-                 create_disposition, write_disposition):
+                 create_disposition='CREATE_IF_NEEDED',
+                 write_disposition='WRITE_EMPTY'):
         self.backup_items = backup_items
         self.target_project_id = target_project_id
         self.target_dataset_id = target_dataset_id
@@ -110,7 +110,7 @@ def __get_backup_entities(backup_items):
     def __create_restore_item(self, restore_request, backup_entity, backup_item):
         source_entity = self.__get_source_table_entity(backup_entity)
 
-        source_table_reference = RestoreTableReference\
+        source_table_reference = RestoreTableReference \
             .backup_table_reference(source_entity, backup_entity)
         target_table_reference = self.__create_target_table_reference(
             restore_request, source_entity)

src/restore/table/table_restore_handler.py

@@ -1,32 +1,43 @@
-from datetime import datetime
-
 import webapp2
+from datetime import datetime
 
+from src.commons.config.configuration import configuration
 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.handlers import validators
-from src.commons.config.configuration import configuration
-from src.restore.table.table_restore_service import TableRestoreService
+from src.commons.handlers.bbq_authenticated_handler import \
+    BbqAuthenticatedHandler
+from src.commons.handlers.json_handler import JsonHandler
 from src.commons.table_reference import TableReference
+from src.restore.table.table_restore_service import TableRestoreService
 
 
 class TableRestoreHandler(JsonHandler):
 
     def get(self, project_id, dataset_id, table_id):
         partition_id = self.request.get('partitionId', None)
+        target_project_id = self.request.get('targetProjectId', None)
         target_dataset_id = self.request.get('targetDatasetId', None)
-        if target_dataset_id:
-            validators.validate_dataset_id(target_dataset_id)
+        create_disposition = self.request.get('createDisposition', None)
+        write_disposition = self.request.get('writeDisposition', None)
+
+        validators.validate_restore_request_params(
+            target_project_id=target_project_id,
+            target_dataset_id=target_dataset_id,
+            create_disposition=create_disposition,
+            write_disposition=write_disposition
+        )
 
         restoration_datetime = self.__get_restoration_datetime()
 
         table_reference = TableReference(project_id, dataset_id,
                                          table_id, partition_id)
 
-        restore_data = TableRestoreService.restore(
-            table_reference, target_dataset_id, restoration_datetime
-        )
+        restore_data = TableRestoreService.restore(table_reference,
+                                                   target_project_id,
+                                                   target_dataset_id,
+                                                   create_disposition,
+                                                   write_disposition,
+                                                   restoration_datetime)
         self._finish_with_success(restore_data)
 
     def __get_restoration_datetime(self):

src/restore/table/table_restore_service.py

@@ -8,12 +8,18 @@
 
 
 class TableRestoreService(object):
+
     @classmethod
-    def restore(cls, table_reference, target_dataset_id, restoration_datetime):
+    def restore(cls, table_reference, target_project_id, target_dataset_id,
+                create_disposition, write_disposition, restoration_datetime):
+
         backup = BackupFinder.for_table(table_reference, restoration_datetime)
 
-        restore_request = BackupListRestoreRequest(
-            [BackupItem(backup.key)], None, target_dataset_id, None, None)
+        restore_request = BackupListRestoreRequest([BackupItem(backup.key)],
+                                                   target_project_id,
+                                                   target_dataset_id,
+                                                   create_disposition,
+                                                   write_disposition)
 
         restoration_job_id = BackupListRestoreService().restore(restore_request)
         logging.info("Scheduled restoration job: %s", restoration_job_id)