def do_execute_redshift_query(query_name, supplied_params, queue_name): """A common entry point for the actual execution of queries on Lambda""" logger = logging.getLogger("hsreplaynet.lambdas.execute_redshift_query") logger.info("Query Name: %s" % query_name) logger.info("Query Params: %s" % supplied_params) query = get_redshift_catalogue().get_query(query_name) parameterized_query = query.build_full_params(supplied_params) try: wlm_queue = settings.REDSHIFT_QUERY_QUEUES[queue_name]["wlm_queue"] with get_concurrent_redshift_query_queue_semaphore(queue_name): _do_execute_query(parameterized_query, wlm_queue) logger.info("Query Execution Complete") return True except NotAvailable: logger.warn( "The Redshift query queue was already at max concurrency. Skipping query." ) metric_fields = {"count": 1} metric_fields.update(parameterized_query.supplied_non_filters_dict) influx_metric("redshift_query_lambda_execution_concurrency_exceeded", metric_fields, query_name=query_name, **parameterized_query.supplied_filters_dict) return False
def do_drain_redshift_query_queue_iteration(queue_name): logger = logging.getLogger( "hsreplaynet.lambdas.drain_redshift_query_queue") logger.info("Queue name: %s" % queue_name) semaphore = get_concurrent_redshift_query_queue_semaphore(queue_name) available_slots = semaphore.available_count logger.info("There are %i available slots" % available_slots) # We don't bother to pull more messages than open slots for running queries if available_slots: messages = get_messages(queue_name, available_slots) if messages: countdown_latch = CountDownLatch(len(messages)) def redshift_query_runner(message): query_name = None try: receipt = message["ReceiptHandle"] body = json.loads(message["Body"]) query_name = body["query_name"] params = body["supplied_parameters"] success = do_execute_redshift_query( query_name, params, queue_name) logger.info("Do Execute Result: %s" % str(success)) if success: # If we don't delete the message then it will be retried in a subsequent lambda SQS.delete_message( QueueUrl=get_or_create_queue(queue_name), ReceiptHandle=receipt) finally: logger.debug("Lambda completed for %s Decrementing latch.", str(query_name)) countdown_latch.count_down() for message in messages: lambda_invocation = Thread(target=redshift_query_runner, args=(message, )) lambda_invocation.start() # We will exit once all child redshift_query_runners have returned. countdown_latch. await () return True else: return False else: return False