class TestSQSHandler(MatrixTestCaseUsingMockAWS): def setUp(self): super(TestSQSHandler, self).setUp() self.sqs_handler = SQSHandler() self.sqs.meta.client.purge_queue(QueueUrl="test_query_job_q_name") def test_add_message_to_queue(self): payload = {'test_key': "test_value"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) messages = self.sqs.meta.client.receive_message( QueueUrl="test_query_job_q_name") message_body = json.loads(messages['Messages'][0]['Body']) self.assertEqual(message_body['test_key'], "test_value") def test_receive_messages_from_queue__returns_None_when_no_messages_found( self): message = self.sqs_handler.receive_messages_from_queue( "test_query_job_q_name", 1) self.assertEqual(message, None) def test_retrieve_messages_from_queue__returns_message_when_message_is_found( self): payload = {'test_key': "test_value"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) messages = self.sqs_handler.receive_messages_from_queue( queue_url="test_query_job_q_name") message_body = json.loads(messages[0]['Body']) self.assertEqual(len(messages), 1) self.assertEqual(message_body['test_key'], "test_value") def test_delete_message_from_queue(self): payload = {'test_key': "test_value"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) messages = self.sqs_handler.receive_messages_from_queue( queue_url="test_query_job_q_name") receipt_handle = messages[0]['ReceiptHandle'] self.sqs_handler.delete_message_from_queue("test_query_job_q_name", receipt_handle) message = self.sqs_handler.receive_messages_from_queue( "test_query_job_q_name", 1) self.assertEqual(message, None)
class TestQueryRunner(MatrixTestCaseUsingMockAWS): def setUp(self): super(TestQueryRunner, self).setUp() self.query_runner = QueryRunner() self.matrix_infra_config.set(self.__class__.TEST_CONFIG) self.query_runner.matrix_infra_config = self.matrix_infra_config self.sqs_handler = SQSHandler() self.sqs.meta.client.purge_queue(QueueUrl="test_query_job_q_name") self.sqs.meta.client.purge_queue( QueueUrl="test_deadletter_query_job_q_name") @mock.patch( "matrix.common.aws.s3_handler.S3Handler.load_content_from_obj_key") @mock.patch( "matrix.common.aws.sqs_handler.SQSHandler.receive_messages_from_queue") def test_run__with_no_messages_in_queue(self, mock_receive_messages, mock_load_obj): mock_receive_messages.return_value = None self.query_runner.run(max_loops=1) mock_receive_messages.assert_called_once_with( self.query_runner.query_job_q_url) mock_load_obj.assert_not_called() @mock.patch( "matrix.common.aws.batch_handler.BatchHandler.schedule_matrix_conversion" ) @mock.patch( "matrix.common.request.request_tracker.RequestTracker.is_request_ready_for_conversion" ) @mock.patch( "matrix.common.request.request_tracker.RequestTracker.complete_subtask_execution" ) @mock.patch( "matrix.common.aws.redshift_handler.RedshiftHandler.transaction") @mock.patch( "matrix.common.aws.s3_handler.S3Handler.load_content_from_obj_key") def test_run__with_one_message_in_queue_and_not_ready_for_conversion( self, mock_load_obj, mock_transaction, mock_complete_subtask, mock_is_ready_for_conversion, mock_schedule_conversion): request_id = str(uuid.uuid4()) payload = {'request_id': request_id, 's3_obj_key': "test_s3_obj_key"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) mock_is_ready_for_conversion.return_value = False self.query_runner.run(max_loops=1) mock_load_obj.assert_called_once_with("test_s3_obj_key") mock_transaction.assert_called() mock_complete_subtask.assert_called_once_with(Subtask.QUERY) mock_schedule_conversion.assert_not_called() @mock.patch( "matrix.common.request.request_tracker.RequestTracker.write_batch_job_id_to_db" ) @mock.patch("matrix.common.request.request_tracker.RequestTracker.format") @mock.patch( "matrix.common.aws.batch_handler.BatchHandler.schedule_matrix_conversion" ) @mock.patch( "matrix.common.request.request_tracker.RequestTracker.is_request_ready_for_conversion" ) @mock.patch( "matrix.common.request.request_tracker.RequestTracker.complete_subtask_execution" ) @mock.patch( "matrix.common.aws.redshift_handler.RedshiftHandler.transaction") @mock.patch( "matrix.common.aws.s3_handler.S3Handler.load_content_from_obj_key") def test_run__with_one_message_in_queue_and_ready_for_conversion( self, mock_load_obj, mock_transaction, mock_complete_subtask, mock_is_ready_for_conversion, mock_schedule_conversion, mock_request_format, mock_write_batch_job_id_to_db): request_id = str(uuid.uuid4()) payload = {'request_id': request_id, 's3_obj_key': "test_s3_obj_key"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) mock_is_ready_for_conversion.return_value = True mock_schedule_conversion.return_value = "123-123" self.query_runner.run(max_loops=1) mock_schedule_conversion.assert_called_once_with(request_id, mock.ANY) mock_write_batch_job_id_to_db.assert_called_once_with("123-123") @mock.patch( "matrix.common.request.request_tracker.RequestTracker.log_error") @mock.patch("matrix.common.request.request_tracker.RequestTracker.format") @mock.patch( "matrix.common.request.request_tracker.RequestTracker.complete_subtask_execution" ) @mock.patch( "matrix.common.aws.redshift_handler.RedshiftHandler.transaction") @mock.patch( "matrix.common.aws.s3_handler.S3Handler.load_content_from_obj_key") def test_run__with_one_message_in_queue_and_fails(self, mock_load_obj, mock_transaction, mock_complete_subtask, mock_request_format, mock_log_error): request_id = str(uuid.uuid4()) payload = {'request_id': request_id, 's3_obj_key': "test_s3_obj_key"} self.sqs_handler.add_message_to_queue("test_query_job_q_name", payload) mock_complete_subtask.side_effect = MatrixException( status=requests.codes.not_found, title=f"Unable to find") self.query_runner.run(max_loops=1) mock_log_error.assert_called_once() query_queue_messages = self.sqs_handler.receive_messages_from_queue( "test_query_job_q_name", 1) self.assertEqual(query_queue_messages, None) deadletter_queue_messages = self.sqs_handler.receive_messages_from_queue( "test_deadletter_query_job_q_name", 1) self.assertEqual(len(deadletter_queue_messages), 1) message_body = json.loads(deadletter_queue_messages[0]['Body']) self.assertEqual(message_body['request_id'], request_id) self.assertEqual(message_body['s3_obj_key'], "test_s3_obj_key")
class QueryRunner: def __init__(self): self.sqs_handler = SQSHandler() self.s3_handler = S3Handler(os.environ["MATRIX_QUERY_BUCKET"]) self.batch_handler = BatchHandler() self.redshift_handler = RedshiftHandler() self.matrix_infra_config = MatrixInfraConfig() @property def query_job_q_url(self): return self.matrix_infra_config.query_job_q_url @property def query_job_deadletter_q_url(self): return self.matrix_infra_config.query_job_deadletter_q_url def run(self, max_loops=None): loops = 0 while max_loops is None or loops < max_loops: loops += 1 messages = self.sqs_handler.receive_messages_from_queue( self.query_job_q_url) if messages: message = messages[0] logger.info(f"Received {message} from {self.query_job_q_url}") payload = json.loads(message['Body']) request_id = payload['request_id'] request_tracker = RequestTracker(request_id) Logging.set_correlation_id(logger, value=request_id) obj_key = payload['s3_obj_key'] receipt_handle = message['ReceiptHandle'] try: logger.info(f"Fetching query from {obj_key}") query = self.s3_handler.load_content_from_obj_key(obj_key) logger.info(f"Running query from {obj_key}") self.redshift_handler.transaction([query], read_only=True) logger.info(f"Finished running query from {obj_key}") logger.info( f"Deleting {message} from {self.query_job_q_url}") self.sqs_handler.delete_message_from_queue( self.query_job_q_url, receipt_handle) logger.info( "Incrementing completed queries in state table") request_tracker.complete_subtask_execution(Subtask.QUERY) if request_tracker.is_request_ready_for_conversion(): logger.info("Scheduling batch conversion job") batch_job_id = self.batch_handler.schedule_matrix_conversion( request_id, request_tracker.format) request_tracker.write_batch_job_id_to_db(batch_job_id) except Exception as e: logger.info( f"QueryRunner failed on {message} with error {e}") request_tracker.log_error(str(e)) logger.info( f"Adding {message} to {self.query_job_deadletter_q_url}" ) self.sqs_handler.add_message_to_queue( self.query_job_deadletter_q_url, payload) logger.info( f"Deleting {message} from {self.query_job_q_url}") self.sqs_handler.delete_message_from_queue( self.query_job_q_url, receipt_handle) else: logger.info(f"No messages to read from {self.query_job_q_url}")