def test_that_proper_entities_were_stored_in_datastore(self): # given restoration_job_key = RestorationJob.create( HARDCODED_UUID, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") restore_item_tuples = self.__create_restore_items(count=3) # when AsyncBatchRestoreService().restore( restoration_job_key, [[restore_item_tuples[0][0], restore_item_tuples[1][0]], [restore_item_tuples[2][0]]]) # then restoration_job = RestorationJob.get_by_id(HARDCODED_UUID) self.assertEqual(restoration_job.items_count, 3) restore_items = list(RestoreItem.query().filter( RestoreItem.restoration_job_key == restoration_job.key)) self.assertEqual(restore_items[0].status, RestoreItem.STATUS_IN_PROGRESS) self.assertEqual(restore_items[0].completed, None) self.assertEqual(restore_items[0].source_table_reference, restore_item_tuples[0][1]) self.assertEqual(restore_items[0].target_table_reference, restore_item_tuples[0][2]) self.assertEqual(restore_items[1].status, RestoreItem.STATUS_IN_PROGRESS) self.assertEqual(restore_items[1].completed, None) self.assertEqual(restore_items[1].source_table_reference, restore_item_tuples[1][1]) self.assertEqual(restore_items[1].target_table_reference, restore_item_tuples[1][2])
def test_restoration_job_saved_in_datastore(self): # when RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # then self.assertIsNotNone(RestorationJob.get_by_id(TEST_RESTORATION_JOB_ID))
def test_restoration_job_has_initial_count_0(self): # when RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # then self.assertEqual( RestorationJob.get_by_id(TEST_RESTORATION_JOB_ID).items_count, 0)
def test_restoration_job_has_now_as_initial_created_time(self): # when RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # then self.assertEqual( RestorationJob.get_by_id(TEST_RESTORATION_JOB_ID).created, datetime(2012, 1, 14, 3, 21, 34))
def test_restoration_job_has_passed_write_disposition(self): # when RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # then self.assertEqual( RestorationJob.get_by_id( TEST_RESTORATION_JOB_ID).write_disposition, WRITE_DISPOSITION)
def test_create_if_not_exist_should_throw_exception_if_already_exist(self): # given RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # when then with self.assertRaises(DuplicatedRestorationJobIdException): RestorationJob.create(restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION)
def restore(self, restoration_job_id, project_id, dataset_id, target_project_id, target_dataset_id, create_disposition, write_disposition, max_partition_days): logging.info( "Executing restoration job '%s' of dataset '%s:%s' " "(target_dataset_id: '%s', max_partition_days:'%s')", restoration_job_id, project_id, dataset_id, target_dataset_id, max_partition_days) try: restoration_job_key = RestorationJob.create( restoration_job_id=restoration_job_id, create_disposition=create_disposition, write_disposition=write_disposition) except DuplicatedRestorationJobIdException: logging.warning( "Trying to create RestorationJob with already existed restoration job id" ) return restore_items = DatasetRestoreItemsGenerator.generate_restore_items( project_id=project_id, dataset_id=dataset_id, target_project_id=target_project_id, target_dataset_id=target_dataset_id, max_partition_days=max_partition_days) AsyncBatchRestoreService().restore(restoration_job_key, restore_items)
def __get_restoration_job_entity(restoration_job_id): entity = RestorationJob.get_by_id(restoration_job_id) if entity: return entity else: raise NotFoundException("Restoration job with id: '{}' doesn't " "exists".format(restoration_job_id))
def test_for_single_item_should_create_post_copy_action(self): # given restoration_job_key = RestorationJob.create( HARDCODED_UUID, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") restore_items_tuple = self.__create_restore_items(count=1) restore_item = restore_items_tuple[0][0] source_bq = restore_items_tuple[0][1].create_big_query_table() target_bq = restore_items_tuple[0][2].create_big_query_table() # when AsyncBatchRestoreService().restore(restoration_job_key, [[restore_item]]) # then self.copy_service.assert_has_calls([ call(copy_job_type_id='restore', task_name_suffix='123'), call().with_post_action( PostCopyActionRequest( url='/callback/restore-finished/', data={'restoreItemKey': (restore_item.key.urlsafe())})), call().with_post_action().with_create_disposition( 'CREATE_IF_NEEDED'), call().with_post_action().with_create_disposition( ).with_write_disposition('WRITE_EMPTY'), call().with_post_action().with_create_disposition(). with_write_disposition().copy_table(source_bq, target_bq) ])
def test_restoration_job_return_entity_key(self): # when entity_key = RestorationJob.create( restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION) # then self.assertIsNotNone(entity_key.get())
def restore(self, restore_request): restoration_job_key = RestorationJob.create( restoration_job_id=restore_request.restoration_job_id, create_disposition=restore_request.create_disposition, write_disposition=restore_request.write_disposition) restore_items = self.__generate_restore_items(restore_request) AsyncBatchRestoreService().restore(restoration_job_key, [restore_items]) return restore_request.restoration_job_id
def __create_restoration_job_with_one_item(restoration_job_id): restoration_job_key = RestorationJob.create( restoration_job_id=restoration_job_id, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") restoration_job_key.get().increment_count_by(1) restore_item = TestRestoreItem.__create_restore_item_example( restoration_job_key) restore_item_key = restore_item.put() return restoration_job_key, restore_item_key, restore_item
def test_failing_creating_dataset_should_update_restore_item_status( self, _): # given restoration_job_key = RestorationJob.create( HARDCODED_UUID, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") restore_items_tuple = self.__create_restore_items(count=1) restore_item = restore_items_tuple[0][0] # when AsyncBatchRestoreService().restore(restoration_job_key, [[restore_item]]) # then restore_items = list(RestoreItem.query().filter( RestoreItem.restoration_job_key == restoration_job_key)) self.assertEqual(restore_items[0].status, RestoreItem.STATUS_FAILED)
def test_that_enforcer_is_called_for_each_restore_item(self): # given restoration_job_key = RestorationJob.create( HARDCODED_UUID, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") restore_items_tuple = self.__create_restore_items(count=2) # when AsyncBatchRestoreService().restore( restoration_job_key, [[restore_items_tuple[0][0], restore_items_tuple[1][0]]]) # then expected_calls = [ call(restore_items_tuple[0][1], restore_items_tuple[0][2]), call(restore_items_tuple[1][1], restore_items_tuple[1][2]) ] self.enforcer.assert_has_calls(expected_calls)
def create_restoration_job_with_count(self, item_count): job_key = RestorationJob.create(TEST_RESTORATION_JOB_ID, create_disposition="CREATE_IF_NEEDED", write_disposition="WRITE_EMPTY") job_key.get().increment_count_by(item_count) return job_key
def __create_restoration_job(self): return RestorationJob.create( restoration_job_id=TEST_RESTORATION_JOB_ID, create_disposition=CREATE_DISPOSITION, write_disposition=WRITE_DISPOSITION).get()