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