m2-traitement-benchmark/drivers/psql.py

from psycopg2._psycopg import connection, cursor
from models.data_thread import Data_Thread, Threads
from interfaces.db_testing_interface import DB_Testing


class PSQL_Thread(Data_Thread):
    def __init__(self, thread: Data_Thread) -> None:
        super().__init__()
        for x in dir(thread):
            if not x.startswith("_"):
                self.__setattr__(x, thread.__getattribute__(x))
                
    def insert_into_db(self, cursor: cursor):        
        insert_query = '''
        INSERT INTO thread (
            sender_name, recipient_name, subject, body, sent_date, received_date, attachment_count, is_read, is_spam, importance_level, reply_count, forward_count, cc_recipients, bcc_recipients, folder
        ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s);
        '''
        
        data = (
            self.sender_name, self.recipient_name, self.subject, self.body,
            self.sent_date, self.received_date, self.attachment_count, self.is_read,
            self.is_spam, self.importance_level, self.reply_count, self.forward_count,
            self.cc_recipients, self.bcc_recipients, self.folder
        )
        
        cursor.execute(insert_query, data)

class PSQL_Testing(DB_Testing):
    driver_name = "PSQL"
    
    
    
    def __init__(self) -> None:
        super().__init__()
        self.data = [PSQL_Thread(x) for x in Threads]

    def create_table(self, driver: connection):
        cursor = driver.cursor()
        createTableQuery = '''
        CREATE TABLE thread (
            thread_id SERIAL PRIMARY KEY,
            sender_name VARCHAR(255) NOT NULL,
            recipient_name VARCHAR(255) NOT NULL,
            subject VARCHAR(255),
            body TEXT,
            sent_date TIMESTAMP,
            received_date TIMESTAMP,
            attachment_count INTEGER,
            is_read BOOLEAN,
            is_spam BOOLEAN,
            importance_level VARCHAR(50),
            reply_count INTEGER,
            forward_count INTEGER,
            cc_recipients TEXT,
            bcc_recipients TEXT,
            folder VARCHAR(100)
        );
        '''
        cursor.execute(createTableQuery)
        driver.commit()
    
    def delete_table(self, driver: connection):
        cursor = driver.cursor()
        dropQuery = 'DROP TABLE IF EXISTS thread;'
        cursor.execute(dropQuery)
        driver.commit()

    def add_singles(self, driver: connection):
        cursor = driver.cursor()
        for t in self.data:
            t.insert_into_db(cursor)
            driver.commit()
            
    def add_bulk(self, driver: connection):
        cursor = driver.cursor()
        for t in self.data:
            t.insert_into_db(cursor)    
        driver.commit()
        
    def delete_data(self, driver: connection):
        cursor = driver.cursor()
        dropQuery = 'DELETE FROM thread;'
        cursor.execute(dropQuery)
        driver.commit()
        
    def attach5_mr_mrs(self, driver: connection):
        cursor = driver.cursor()
        searchQuery = "SELECT COUNT(*) FROM thread WHERE attachment_count > 5 AND (subject = 'Mr' or subject  = 'Mrs');"
        cursor.execute(searchQuery)
        driver.commit()
        return f"Got {cursor.fetchone()[0]}"

    def search_mails_ends_microsoftcom(self, driver: connection):
        cursor = driver.cursor()
        searchQuery = "SELECT COUNT(*) FROM thread WHERE cc_recipients LIKE '%@microsoft.com';"
        cursor.execute(searchQuery)
        driver.commit()
        return f"Got {cursor.fetchone()[0]}"
        
    def search_lorem(self, driver: connection):
        cursor = driver.cursor()
        searchQuery = "select COUNT(*) from thread where body like '%Nullam sit amet turpis elementum ligula vehicula consequat. Morbi a ipsum. Integer a nibh.%';"
        cursor.execute(searchQuery)
        driver.commit()
        return f"Got {cursor.fetchone()[0]}"
        
    def get_sum_attachment_less_5(self, driver: connection):
        cursor = driver.cursor()
        searchQuery = "select SUM(attachment_count) from thread where attachment_count < 5;"
        cursor.execute(searchQuery)
        driver.commit()
        return f"Got {cursor.fetchone()[0]}"
            
    def update_add_replies_per_last_name(self, driver: connection) -> str:
        cursor = driver.cursor()
        alterQuery = """
        ALTER TABLE thread ADD COLUMN last_name_replies INT;
        """
        updateQuery = """
        WITH last_name_sums AS (
            SELECT 
                recipient_name,
                SUM(reply_count) OVER (PARTITION BY SPLIT_PART(recipient_name, ' ', -1)) AS total_replies
            FROM 
                thread
        )
        UPDATE thread AS t
        SET last_name_replies = ln_sums.total_replies
        FROM last_name_sums AS ln_sums
        WHERE t.recipient_name = ln_sums.recipient_name;
        """
        unalterQuery = """
        ALTER TABLE thread DROP COLUMN last_name_replies;
        """ 
        cursor.execute(alterQuery)
        cursor.execute(updateQuery)
        cursor.execute(unalterQuery)
        driver.commit()
        
    def update_add_replies_per_subject(self, driver: connection) -> str:
        cursor = driver.cursor()
        alterQuery = """
        ALTER TABLE thread ADD COLUMN subject_replies INT;
        """
        updateQuery = """
        WITH subject_sums AS (
            SELECT 
                recipient_name,
                SUM(reply_count) OVER (PARTITION BY SPLIT_PART(recipient_name, ' ', -1)) AS total_replies
            FROM 
                thread
        )
        UPDATE thread AS t
        SET subject_replies = ln_sums.total_replies
        FROM subject_sums AS ln_sums
        WHERE t.recipient_name = ln_sums.recipient_name;
        """
        unalterQuery = """
        ALTER TABLE thread DROP COLUMN subject_replies;
        """ 
        cursor.execute(alterQuery)
        cursor.execute(updateQuery)
        cursor.execute(unalterQuery)
        driver.commit()