def test_bug_regression_job_already_exists_after_internal_error(
self, _, insert_job, _create_random_job_id,
create_copy_job_result_check, table_metadata):
# given
post_copy_action_request = \
PostCopyActionRequest(url='/my/url', data={'key1': 'value1'})
table_metadata._BigQueryTableMetadata__get_table_or_partition.return_value.get_location.return_value = 'EU'
# when
CopyJobService().run_copy_job_request(
CopyJobRequest(task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
source_big_query_table=self.example_source_bq_table,
target_big_query_table=self.example_target_bq_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request))
# then
self.assertEqual(insert_job.call_count, 2)
create_copy_job_result_check.assert_called_once_with(
ResultCheckRequest(
task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
job_reference=BigQueryJobReference(
project_id='target_project_id_1',
job_id='random_job_123',
location='EU'),
retry_count=0,
post_copy_action_request=post_copy_action_request))
def test_copy_job_result_check_creation(self):
# given
result_check_request = ResultCheckRequest(
task_name_suffix='task-name-suffix',
copy_job_type_id='backups',
job_reference=BigQueryJobReference(project_id="project_abc",
job_id="job123",
location='EU'),
retry_count=2,
post_copy_action_request=PostCopyActionRequest(
url="/my/url",
data={"key1": "value1"})
)
TaskCreator.create_copy_job_result_check(result_check_request)
# then
expected_queue_name = 'backups-result-check'
executed_tasks = self.taskqueue_stub.get_filtered_tasks(
queue_names=expected_queue_name
)
self.assertEqual(len(executed_tasks), 1,
"Should create one task in queue")
executed_task = executed_tasks[0]
self.assertEqual(json.dumps(result_check_request, cls=RequestEncoder),
executed_task.extract_params()['resultCheckRequest'])
self.assertEqual('POST', executed_task.method)
self.assertEqual(executed_task.url,
'/tasks/copy_job_async/result_check')
def test_happy_path(self, result_check_mock):
# given
project_id = "target_project_id"
job_id = "job_id"
location = "EU"
retry_count = "1"
post_copy_action_request = \
PostCopyActionRequest(url="/my/url", data={"key1": "value1"})
result_check_request = self.create_example_result_check_request(
project_id, job_id, location, retry_count,
post_copy_action_request)
# when
self.under_test.post(url='/tasks/copy_job_async/result_check',
params={
"resultCheckRequest":
json.dumps(result_check_request,
cls=RequestEncoder)
})
# then
result_check_mock.assert_called_with(
self.create_example_result_check_request(project_id, job_id,
location, retry_count,
post_copy_action_request))
def test_that_should_re_trigger_copy_job_task_with_proper_create_and_write_dispositions_if_retry_error_occurs(
self, create_copy_job, _):
# given
retry_count = 0
post_copy_action_request = \
PostCopyActionRequest(url="/my/url", data={"key1": "value1"})
create_disposition = "CREATE_NEVER"
write_disposition = "WRITE_TRUNCATE"
# when
ResultCheck().check(
ResultCheckRequest(
task_name_suffix="task_name_suffix",
copy_job_type_id="backups",
job_reference=BigQueryJobReference(
project_id="target_project_id",
job_id="job_id",
location='EU'),
retry_count=retry_count,
post_copy_action_request=post_copy_action_request))
# then
copy_job_result = CopyJobResult(
JobResultExample.DONE_WITH_RETRY_ERRORS)
copy_job_request = CopyJobRequest(
task_name_suffix=None,
copy_job_type_id="backups",
source_big_query_table=copy_job_result.source_bq_table,
target_big_query_table=copy_job_result.target_bq_table,
create_disposition=create_disposition,
write_disposition=write_disposition,
retry_count=retry_count + 1,
post_copy_action_request=post_copy_action_request)
create_copy_job.assert_called_once_with(copy_job_request)
def __run_copy_job_for_each(self, restore_items, restoration_job):
logging.info("Scheduling %s", len(restore_items))
for restore_item in restore_items:
source_table_reference = restore_item.source_table_reference
target_table_reference = restore_item.target_table_reference
try:
self.restore_workspace_creator.create_workspace(
source_table_reference, target_table_reference)
CopyJobServiceAsync(
copy_job_type_id='restore',
task_name_suffix=restoration_job.key.id()
).with_post_action(
PostCopyActionRequest(
url='/callback/restore-finished/',
data={'restoreItemKey': restore_item.key.urlsafe()})
).with_create_disposition(
restoration_job.create_disposition).with_write_disposition(
restoration_job.write_disposition).copy_table(
source_table_reference.create_big_query_table(),
target_table_reference.create_big_query_table())
except Exception as ex:
logging.error(
"Error during creating copy job. Marking restore "
"item as FAILED, Error message: %s", ex.message)
restore_item.update_with_failed(restore_item.key, ex.message)
def test_that_copy_table_should_create_correct_post_copy_action_if_404_http_error_thrown_on_copy_job_creation(
self, create_post_copy_action, insert_job):
# given
error = HttpError(Mock(status=404), 'not found')
error._get_reason = Mock(return_value='not found')
insert_job.side_effect = error
post_copy_action_request = PostCopyActionRequest(
url='/my/url', data={'key1': 'value1'})
request = CopyJobRequest(
task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
source_big_query_table=self.example_source_bq_table,
target_big_query_table=self.example_target_bq_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request)
# when
CopyJobService().run_copy_job_request(request)
# then
create_post_copy_action.assert_called_once_with(
copy_job_type_id='test-process',
post_copy_action_request=post_copy_action_request,
job_json={
'status': {
'state':
'DONE',
'errors': [{
'reason':
'Invalid',
'message':
("404 while creating Copy Job from {} to {}".format(
self.example_source_bq_table,
self.example_target_bq_table))
}]
},
'configuration': {
'copy': {
'sourceTable': {
'projectId':
self.example_source_bq_table.get_project_id(),
'tableId':
self.example_source_bq_table.get_table_id(),
'datasetId':
self.example_source_bq_table.get_dataset_id()
},
'destinationTable': {
'projectId':
self.example_target_bq_table.get_project_id(),
'tableId':
self.example_target_bq_table.get_table_id(),
'datasetId':
self.example_target_bq_table.get_dataset_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_copy_job_creation(self, _):
# given
copy_job_request = CopyJobRequest(
task_name_suffix='task-name-suffix',
copy_job_type_id="backups",
source_big_query_table=BigQueryTable('source_project',
'source_dataset',
'source_table'),
target_big_query_table=BigQueryTable('target_project',
'target_dataset',
'target_table'),
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
post_copy_action_request=PostCopyActionRequest(url="/my/url", data={
"key1": "value1"})
)
# when
TaskCreator.create_copy_job(
copy_job_request=copy_job_request
)
# then
expected_queue_name = 'backups-copy-job'
executed_tasks = self.taskqueue_stub.get_filtered_tasks(
queue_names=expected_queue_name
)
self.assertEqual(len(executed_tasks), 1,
"Should create one task in queue")
executed_task = executed_tasks[0]
self.assertEqual(json.dumps(copy_job_request, cls=RequestEncoder),
executed_task.extract_params()['copyJobRequest'])
self.assertEqual('POST', executed_task.method)
self.assertEqual('task_name', executed_task.name)
self.assertEqual(executed_task.url, '/tasks/copy_job_async/copy_job')
def test_that_post_copy_action_request_is_passed(
self, create_copy_job_result_check, _):
# given
post_copy_action_request = \
PostCopyActionRequest(url='/my/url', data={'key1': 'value1'})
# when
CopyJobService().run_copy_job_request(
CopyJobRequest(task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
source_big_query_table=self.example_source_bq_table,
target_big_query_table=self.example_target_bq_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request))
# then
create_copy_job_result_check.assert_called_once_with(
ResultCheckRequest(
task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
job_reference=BigQueryJobReference(project_id='test_project',
job_id='job123',
location='EU'),
retry_count=0,
post_copy_action_request=post_copy_action_request))
def test_that_after_successful_job_post_copy_action_request_is_created(
self, create_post_copy_action, _):
# given
# when
ResultCheck().check(self.create_example_result_check_request())
# then
create_post_copy_action.assert_called_once_with(
copy_job_type_id="backups",
post_copy_action_request=PostCopyActionRequest(
url="/my/url", data={"key1": "value1"}),
job_json=JobResultExample.DONE)
def test_that_should_execute_post_copy_action_request_if_not_repetitive_error_occurs(
self, create_post_copy_action, create_copy_job, _):
# given
# when
ResultCheck().check(self.create_example_result_check_request())
# then
create_post_copy_action.assert_called_once_with(
copy_job_type_id="backups",
post_copy_action_request=PostCopyActionRequest(
url="/my/url", data={"key1": "value1"}),
job_json=JobResultExample.DONE_WITH_NOT_REPETITIVE_ERRORS)
create_copy_job.assert_not_called()
def test_create_post_copy_action_throws_error_on_unknown_queue(self):
# when
with self.assertRaises(UnknownQueueError) as error:
TaskCreator.create_post_copy_action(
copy_job_type_id="unknown-copying",
post_copy_action_request=PostCopyActionRequest(
'/my/post/copy/url', {'mypayload': 'mypayload_value'}),
job_json={"state": "DONE"})
self.assertEqual(
error.exception.message, "There is no queue "
"'unknown-copying-post-copy-action'. "
"Please add it to your queue.yaml "
"definition.")
def __copy_table_async(source_bq_table, destination_bq_table):
CopyJobServiceAsync(
copy_job_type_id='backups',
task_name_suffix=request_correlation_id.get()).with_post_action(
PostCopyActionRequest(
url='/callback/backup-created/{}/{}/{}'.format(
source_bq_table.project_id, source_bq_table.dataset_id,
source_bq_table.table_id),
data={
"sourceBqTable": source_bq_table,
"targetBqTable": destination_bq_table
})).copy_table(source_bq_table, destination_bq_table)
def create_example_result_check_request(self):
retry_count = 0
post_copy_action_request = \
PostCopyActionRequest(url="/my/url", data={"key1": "value1"})
# when
result_check_request = ResultCheckRequest(
task_name_suffix="task_name_suffix",
copy_job_type_id="backups",
job_reference=BigQueryJobReference(project_id="target_project_id",
job_id="job_id",
location='EU'),
retry_count=retry_count,
post_copy_action_request=post_copy_action_request)
return result_check_request
def from_json(cls, json):
from src.commons.big_query.big_query_job_reference import \
BigQueryJobReference
job_reference = BigQueryJobReference.from_json(json["job_reference"])
from src.commons.big_query.copy_job_async.post_copy_action_request import \
PostCopyActionRequest
post_copy_action_request = PostCopyActionRequest.from_json(
json["post_copy_action_request"])
return ResultCheckRequest(
task_name_suffix=json["task_name_suffix"],
copy_job_type_id=json["copy_job_type_id"],
job_reference=job_reference,
retry_count=json["retry_count"],
post_copy_action_request=post_copy_action_request)
def test_create_copy_job_result_check_throws_error_on_unknown_queue(self):
# when
with self.assertRaises(UnknownQueueError) as error:
TaskCreator.create_copy_job_result_check(ResultCheckRequest(
task_name_suffix=None,
copy_job_type_id="unknown-copying",
job_reference=BigQueryJobReference(project_id="project_abc",
job_id="job123",
location='EU'),
retry_count=0,
post_copy_action_request=PostCopyActionRequest(
'/my/post/copy/url', {'mypayload': 'mypayload_value'}))
)
self.assertEqual(error.exception.message,
"There is no queue 'unknown-copying-result-check'. "
"Please add it to your queue.yaml definition.")
def from_json(cls, json):
source_big_query_table = BigQueryTable.from_json(
json["source_big_query_table"])
target_big_query_table = BigQueryTable.from_json(
json["target_big_query_table"])
post_copy_action_request = PostCopyActionRequest.from_json(
json["post_copy_action_request"])
return CopyJobRequest(
task_name_suffix=json["task_name_suffix"],
copy_job_type_id=json["copy_job_type_id"],
source_big_query_table=source_big_query_table,
target_big_query_table=target_big_query_table,
create_disposition=json["create_disposition"],
write_disposition=json["write_disposition"],
retry_count=json["retry_count"],
post_copy_action_request=post_copy_action_request)
def test_happy_path(self, copy_table_mock):
# given
source_big_query_table = BigQueryTable("source_project_id",
"source_dataset_id",
"source_table_id")
target_big_query_table = BigQueryTable("target_project_id",
"target_dataset_id",
"target_table_id")
post_copy_action_request = PostCopyActionRequest(
url="/my/url", data={"key1": "value1"})
url = '/tasks/copy_job_async/copy_job'
# when
self.under_test.post(
url=url,
params={
"copyJobRequest":
json.dumps(CopyJobRequest(
task_name_suffix=None,
copy_job_type_id=None,
source_big_query_table=source_big_query_table,
target_big_query_table=target_big_query_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request),
cls=RequestEncoder)
})
# then
copy_table_mock.assert_called_with(
CopyJobRequest(task_name_suffix=None,
copy_job_type_id=None,
source_big_query_table=source_big_query_table,
target_big_query_table=target_big_query_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request))
def test_that_post_copy_action_request_is_passed(self, create_copy_job):
# given
post_copy_action_request = \
PostCopyActionRequest(url="/my/url", data={"key1": "value1"})
# when
CopyJobServiceAsync(copy_job_type_id="test-process",
task_name_suffix="example_sufix").with_post_action(
post_copy_action_request).copy_table(
self.create_example_source_bq_table(),
self.create_example_target_bq_table())
# then
create_copy_job.assert_called_once_with(
CopyJobRequest(
task_name_suffix="example_sufix",
copy_job_type_id="test-process",
source_big_query_table=(self.create_example_source_bq_table()),
target_big_query_table=(self.create_example_target_bq_table()),
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request))
def test_create_post_copy_action(self):
# when
TaskCreator.create_post_copy_action(
copy_job_type_id="backups",
post_copy_action_request=PostCopyActionRequest('/my/post/copy/url',
{
'mypayload': 'mypayload_value'}),
job_json={"state": "DONE"}
)
# then
expected_queue_name = 'backups-post-copy-action'
executed_tasks = self.taskqueue_stub.get_filtered_tasks(
queue_names=expected_queue_name
)
self.assertEqual(len(executed_tasks), 1,
"Should create one task in queue")
executed_task = executed_tasks[0]
self.assertEqual('POST', executed_task.method)
self.assertEqual('/my/post/copy/url', executed_task.url)
self.assertEqual(
'{"jobJson": {"state": "DONE"}, "data": {"mypayload": "mypayload_value"}}',
executed_task.payload)
def test_that_copy_table_should_create_correct_post_copy_action_if_access_denied_http_error_thrown_on_copy_job_creation(
self, create_post_copy_action, insert_job):
# given
http_error_content = "{\"error\": " \
" {\"errors\": [" \
" {\"reason\": \"Access Denied\"," \
" \"message\": \"Access Denied\"," \
" \"location\": \"US\"" \
" }]," \
" \"code\": 403," \
" \"message\": \"Access Denied\"}}"
insert_job.side_effect = HttpError(Mock(status=403),
http_error_content)
post_copy_action_request = PostCopyActionRequest(
url='/my/url', data={'key1': 'value1'})
request = CopyJobRequest(
task_name_suffix='task_name_suffix',
copy_job_type_id='test-process',
source_big_query_table=self.example_source_bq_table,
target_big_query_table=self.example_target_bq_table,
create_disposition="CREATE_IF_NEEDED",
write_disposition="WRITE_EMPTY",
retry_count=0,
post_copy_action_request=post_copy_action_request)
# when
CopyJobService().run_copy_job_request(request)
# then
create_post_copy_action.assert_called_once_with(
copy_job_type_id='test-process',
post_copy_action_request=post_copy_action_request,
job_json={
'status': {
'state':
'DONE',
'errors': [{
'reason':
'Invalid',
'message':
("Access Denied while creating Copy Job from {} to {}".
format(self.example_source_bq_table,
self.example_target_bq_table))
}]
},
'configuration': {
'copy': {
'sourceTable': {
'projectId':
self.example_source_bq_table.get_project_id(),
'tableId':
self.example_source_bq_table.get_table_id(),
'datasetId':
self.example_source_bq_table.get_dataset_id()
},
'destinationTable': {
'projectId':
self.example_target_bq_table.get_project_id(),
'tableId':
self.example_target_bq_table.get_table_id(),
'datasetId':
self.example_target_bq_table.get_dataset_id()
}
}
}
})