import pymysql.cursors
from data_model import MysqlConfig


class Database:
    def __init__(self, config: MysqlConfig):
        self.config = config
        self.connection = None

    def connect(self):
        if not self.connection or not self.connection.open:
            self.connection = pymysql.connect(
                host=self.config.host,
                user=self.config.user,
                password=self.config.password,
                database=self.config.db_name,
                port=self.config.port,
                cursorclass=pymysql.cursors.DictCursor
            )
        return self.connection

    def close(self):
        if self.connection and self.connection.open:
            self.connection.close()

    def execute(self, query, args=None):
        connection = self.connect()
        try:
            with connection.cursor() as cursor:
                cursor.execute(query, args)
            connection.commit()
        finally:
            self.close()

    def fetchone(self, query, args=None):
        connection = self.connect()
        try:
            with connection.cursor() as cursor:
                cursor.execute(query, args)
                result = cursor.fetchone()
            return result
        finally:
            self.close()

    def fetchall(self, query, args=None):
        connection = self.connect()
        try:
            with connection.cursor() as cursor:
                cursor.execute(query, args)
                results = cursor.fetchall()
            return results
        finally:
            self.close()