databaseconnect.py

import logging
import logger_config
import chatbot
import random
import config
import mysql.connector
from time import sleep

log = logging.getLogger(__name__)
log.info("Entered module: %s" % __name__)

@logger_config.logger
def connection_to_database() -> object:
    """
    Connects to mysql database

    Returns:
        Database connection object
    """
    max_tries = 5
    tries = 1
    conn = None
    while tries <= max_tries:
        try:
            conn = mysql.connector.connect(
                user=config.user,
                password=config.password,
                host=config.host,
                port=config.port,
                database=config.database,
            )
            if conn.is_connected():
                # logging.debug("Connected")
                logging.debug("MySQL connected")
                break

        except mysql.connector.Error as e:
            tries += 1
            logging.debug(e, "...Retrying")
            sleep(20)
    try:
        if conn.is_connected():
            logging.debug("MySQL connected")
            return conn
    except Exception as e:
        raise Exception("DATABASE NOT CONNECTED", str(e))


@logger_config.logger
# setup database
def setup_database() -> object:
    """
    Setup database with chat table, question table, statement table and directions table

    Returns:
        None
    """
    db = connection_to_database()
    cursor = db.cursor()
    # TABLES array containing all required tables to be setup
    TABLES = [
        "chat_table",
        "statement_table",
        "question_table",
        "directions_table"
    ]

    for table in TABLES[:3]:
        cursor.execute("CREATE TABLE IF NOT EXISTS {table}(id INTEGER PRIMARY KEY AUTO_INCREMENT, root_word VARCHAR(40), subject VARCHAR(40), verb VARCHAR(40), sentence VARCHAR(200))")
    
    # setup for directions table
    cursor.execute(
      "CREATE TABLE IF NOT EXISTS {TABLES[3]}(id INTEGER PRIMARY KEY AUTO_INCREMENT, origin_location VARCHAR(100), destination_location VARCHAR(100))"  # noqa: E501
    ) 
    cursor.close()
    db.close()
    

@logger_config.logger
# add classified sentences to database
def add_to_database(classification: str, subject: str, root: str, verb: str, sentence: str) -> None:
    """
    Adds the classified sentence to database, if not already exists.

    Args:
        classification (str): The classification of statement: a chat, question or statement
        subject (str): The subject of the sentence from user input. For ex: Origin, place, etc
        root (str): The root word extracted from the sentence.
        verb (str): Verb relation of sentence from user input.

    Returns:
        None

    """
    db = connection_to_database()
    cursor = db.cursor(prepared=True)
    # If classification of user input is chat ('C'), insert into chat table
    if classification == 'C':
        query = "INSERT INTO chat_table(root_word, verb, sentence) VALUES (%s, %s, %s)"
        # Below params can be used for data validation and sanitation before binding to sql query
        params = (str(root), str(verb), sentence,)
        cursor.execute(query, params)
        db.commit()

    # If classification of user input is question ('Q'), checks the question table if it's already present
    elif classification == "Q":
        cursor.execute("SELECT sentence FROM question_table")
        result = cursor.fetchall()
        exist = 0
        for r in result:
            if r[-1] == sentence:
                exist = 1
                # do not add if question already exists
                break
        if exist == 0:
            query = "INSERT INTO question_table(subject, root_word, verb, sentence) VALUES (%s, %s, %s, %s)"
            params = (str(subject), str(root),str(verb),sentence,)
            cursor.execute(query, params)
            db.commit()

    else:
        # Adds to statement table otherwise
        cursor.execute("SELECT sentence FROM statement_table")
        result = cursor.fetchall()
        exist = 0
        for r in result:
            if r[-1] == sentence:
                exist = 1
                # do not add if statement already exists
                break
        if exist == 0:
            query = "INSERT INTO statement_table(subject, root_word, verb, sentence) VALUES (%s, %s, %s, %s)"
            params = (str(subject), str(root), str(verb), sentence,)
            cursor.execute(query, params)   
            db.commit()
    cursor.close()
    db.close()


@logger_config.logger
# get a random chat response
def get_chat_response() -> str:
    """
    Gets a random chat response from the chat table.

    Returns:
        response(str): A chat response from the chat table.
    """

    db = connection_to_database()
    cursor = db.cursor(prepared=True)
    cursor.execute("SELECT COUNT(*) FROM chat_table")
    res = cursor.fetchone()
    total_chat_records = res[0]

    chat_id = random.randint(1, total_chat_records)
    query = "SELECT sentence FROM chat_table WHERE id = %s"
    params = (str(chat_id),)
    cursor.execute(query, params)
    result = cursor.fetchone()
    response = result[0]
    cursor.close()
    db.close()
    return response


@logger_config.logger
def get_question_response(subject: str, root: str, verb: str) -> tuple[str, str]:
    """
    Responds to the user's input if subject or verb is present in the table. 
    If not, requests the user to train the chatbot.

    Args:
        subject (str): The subject of the sentence from user input. For ex: "Origin", "Place", etc.
        root (str): The root word extracted from the sentence.
        verb (str): Verb relation of sentence from user input.

    Returns:
        tuple[str, str] where:
            response(str): Response to the user either from the subject table if subject found
            or statement table
    """
    db = connection_to_database()
    cursor = db.cursor(prepared=True)
    # If subject is empty, check if the extracted verb is present in the statement table
    if str(subject) == "[]":
        cursor.execute("SELECT verb FROM statement_table")
        res = cursor.fetchall()
        found = 0
        for r in res:
            if r[-1] == str(verb):
                found = 1
                break
        if found == 1:
            query = "SELECT sentence FROM statement_table WHERE verb= %s"
            params = (str(verb),)
            cursor.execute(query, params)
            res = cursor.fetchone()
            response = res[0]
            return response, chatbot.LearnResponse.MESSAGE.name
        else:
            response = "Sorry I don't know the response to this. Please train me."
            return response, chatbot.LearnResponse.TRAIN_ME.name

     # If subject exists, check if the subject is present in the statement table
    else:
        cursor.execute("SELECT subject FROM statement_table")
        res = cursor.fetchall()
        found = 0
        for r in res:
            if r[-1] == str(subject[0]):
                found = 1
                break

        # If subject is present, check the correspinding verb for it    
        if found == 1:
            query = "SELECT verb FROM statement_table WHERE subject= %s"
            params = (str(subject[0]),)
            cursor.execute(query, params)
            res = cursor.fetchone()
            checkVerb = res[0]   

            # checkVerb is a string while verb is a list. checkVerb ['verb']
            if checkVerb == "[]" or checkVerb[2:-2] == verb[0]:
                query = "SELECT sentence FROM statement_table WHERE subject= %s"
                params = (str(subject[0]),)
                cursor.execute(query, params)
                res = cursor.fetchone()
                response = res[0]
                return response, chatbot.LearnResponse.MESSAGE.name
            else:
                # If subject exists and verb not found, responds with training request
                response = "Sorry I don't know the response to this. Please train me."
                return response, chatbot.LearnResponse.TRAIN_ME.name
            
        else:
            # If subject exists and verb not found, responds with training request
            response = "Sorry I don't know the response to this. Please train me."
            return response, chatbot.LearnResponse.TRAIN_ME.name
        

@logger_config.logger
def learn_question_response(sentence: str) -> tuple[str, str]:
    """
    Update the database with the sentence from the user input

    Args:
        sentence (str): The user input.

    Returns:
        tuple[str, str] where:
            response(str): Response to the user
            chatbot.LearnResponse.MESSAGE.name(str): Chatbot's response
    """
    db = connection_to_database()
    cursor = db.cursor(buffered=True)
    cursor.execute("SELECT id FROM statement_table ORDER BY id DESC")
    res = cursor.fetchone()
    last_id = res[0]
    query = "UPDATE statement_table SET sentence = %s WHERE id = %s"
    params = (sentence, str(last_id),)
    cursor.execute(query, params)
    db.commit()
    cursor.close()
    db.close()
    response = "Thank you! I have learnt this."
    return response, chatbot.LearnResponse.MESSAGE.name


@logger_config.logger
def clear_table(table_name) -> object:
    """
    Deletes the question table and/or statement table based on user agreement

    Args:
        table_name (str): The table name to be cleaned.

    Returns:
        object: Database object
    """
    db = connection_to_database()
    cursor = db.cursor()
    
    if table_name in ("question_table", "statement_table"):
        tables_to_be_cleaned = ("question_table", "statement_table")
        logging.debug("The following tables will be cleaned:\n")
        for table in tables_to_be_cleaned:
            describe_table(cursor, table)
        
        if input("Enter 'Y' to confirm cleaning of BOTH tables: ") in ("Y", "y",):
            for table in tables_to_be_cleaned:
                cursor.execute("DELETE FROM %s",(table,))  
            db.commit()
            logging.debug("Tables cleaned successfully")
        else:
            logging.debug("Table cleaning skipped.")

    else:
        logging.debug("The following table will be cleaned:\n")
        describe_table(cursor, table_name)

        if input("Enter 'Y' to confirm: ") in ("Y", "y"):
            cursor.execute("DELETE FROM %s",(table_name,))
            logging.debug("Table cleaned successfully")
            db.commit()
        else:
            logging.debug("Table cleaning skipped.")

    cursor.close()
    db.close()


@logger_config.logger
def describe_table(cursor, table_name):
    cursor.execute("DESC %s",(table_name,))
    res = cursor.fetchall()
    column_names = [col[0] for col in res]
    
    cursor.execute("SELECT COUNT(*) FROM %s",(table_name,))
    res = cursor.fetchall()
    records_no = res[0][0]

    logging.debug("Table Name:", table_name)
    logging.debug("Columns:", column_names)
    logging.debug("Number of existing records:", records_no)
    logging.debug()
    
    return records_no