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