conjurer/conjurer_librarian/search_bot.py

"""
This module contains functions for searching for DOI (Digital Object Identifier)
in a list of live results.
It includes a producer-consumer pattern implementation using threads and queues.

Functions:
- producer: Reads lines from a file and puts them into an output queue.
- consumer: Consumes items from an input queue and checks if DOI exists in the live results list.
- search_for_doi: Searches for DOI in live results using the producer-consumer pattern.

Global Variables:
- MAXTHREADS: Maximum number of worker threads.
- DATABASE_PATH: Path to the database files.
- ENCODING: Encoding of the database files.
- CHUNK: File name pattern for the database files.
- _sentinel: Sentinel object used to signal termination.
- result_list: List to store the search results.
- WORK_Q_SIZE: Maximum size of the work queue.
"""

# TODO: Wpiemdolić to wszystko w klasę z loggerem przysłanym z góry
from queue import Empty, Queue
from threading import Thread

q = Queue()

# DATA FOR TEST ONLY
# MAXTHREADS = 5
# DATABASE_PATH = r'C:\\Database\\chunks_1\\'

# DEPLOYMENT DATA
MAXTHREADS = 41
DATABASE_PATH = r"C:\\Database\\chunks\\"

ENCODING = "utf-8"
CHUNK = "_chunk.txt"
_sentinel = object()
WORK_Q_SIZE = 35500000


def producer(out_q, control_q, filename, _logger):
    """
    Produces items from the output queue and puts them into the control queue.

    Args:
        out_q (Queue): Output queue.
        control_q (Queue): Control queue.
        filename (str): Name of the file.
        _logger: Logger object for logging.
    """
    with open(DATABASE_PATH + filename, "r+", encoding=ENCODING) as operated_file:
        _logger.info(f"Worker {filename} ")
        while True:
            line = operated_file.readline()
            if not line:
                break
            out_q.put(line)
            try:
                check = control_q.get(block=False)
            except Empty:
                check = False
            if check is _sentinel:
                _logger.info("TERM signal received")
                control_q.put(check)
                break
        _logger.info(f"Worker finished: {filename}")
        out_q.put(_sentinel)


# IMPORTANT!!! ONLY ONE CONSUMER THREAD AS WE ARE NOT PUTTING SENTINELS BACK
def consumer(in_q, control_q, doi, live_results, result_list, control_dict, no, _logger):
    """
    Consumes items from an input queue and checks if DOI exists in the live results list.

    Args:
        in_q (Queue): Input queue.
        control_q (Queue): Control queue.
        doi (list): List of DOI to search for.
        live_results (list): List to store the search results.
        _logger: Logger object for logging.
    """
    _logger.info("Consumer thread started:%s", no)
    for item in doi:
        result_list.append({"DOI": item[0], "exists": False, "data": item[1]})
    while True:
        done_check = True
        try:
            data = in_q.get(block=False)
            if data is _sentinel:
                _logger.info("Worker finished")
                control_dict["sentinels"] += 1
            else:
                for item in result_list:
                    if item["DOI"] in data and not item["exists"]:
                        _logger.info(data)
                        _logger.info("HIT")
                        item["exists"] = True
                        live_results.append(item)
                    done_check = done_check and item["exists"]
                if done_check:
                    control_q.put(_sentinel)
        except Empty:
            pass
        if control_dict["sentinels"] >= MAXTHREADS:
            _logger.info("All workers finished")
            break



def search_for_doi(doi, live_results, _logger):
    """
    Search for DOI in live_results using _logger for logging.

    Args:
        doi (list): List of DOI to search for.
        live_results (list): List to store the search results.
        _logger: Logger object for logging.
    """
    control_dict = {"sentinels":0}
    result_list = []
    threads = []
    work_q = Queue(maxsize=WORK_Q_SIZE)
    control_q = Queue()
    for i in range (0, (len(doi)//1000)+2):
        t_cons = Thread(
            target=consumer, args=(work_q, control_q, doi, live_results, result_list, control_dict, i, _logger)
        )
        _logger.info("Consumer thread created")
        threads.append(t_cons)
    for i in range(0, MAXTHREADS):
        # TEST DATA
        # filename = "test" + str(i) + CHUNk
        # DEPLOYMENT DATA
        filename = str(i) + CHUNK
        _logger.info(f"Creating worker thread no: {i}")
        threads.append(
            Thread(target=producer, args=(work_q, control_q, filename, _logger))
        )
    for worker in threads:
        worker.start()
    for worker in threads:
        worker.join()
    return result_list


if __name__ == "__main__":
    import logging

    logger = logging.getLogger()
    logger.setLevel("DEBUG")
    h1 = logging.StreamHandler()
    logger.addHandler(h1)
    logger.info("TEST RUN")
    live_result = []
    logger.info(
        search_for_doi(
            [
                ("10.1002/9781118786352.wbieg0998.pub2", "DATA"),
                ("10.1002/j.2050-0416.2002.tb00563.x", "DATA"),
                ("10.1111/j.1365-2958.1994.tb00448.x", "DATA"),
                ("10.2165/00128415-200309690-00017", "DATA"),
                ("10.5772/48313", "DATA"),
                ("10.15803/ijnc.7.2_419", "DATA"),
                ("10.1051/0004-6361/201321596e", "DATA"),
                ("10.2307/40835941", "DATA"),
            ],
            live_result,
            logger,
        )
    )
    logger.info(live_result)