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()