HFCNotificator/db_access.py

import logging
import os
import time

from dotenv import load_dotenv
from mysql import connector as mysql

load_dotenv()
DB_USERNAME = os.getenv('DB_USERNAME')
DB_PASSWORD = os.getenv('DB_PASSWORD')


class Area:
    def __init__(self, id: int, name: str):
        self.id = id
        self.name = name


class District:
    def __init__(self, id: int, name: str, area_id: int, migun_time: int):
        self.id = id
        self.name = name
        self.area_id = area_id
        self.migun_time = migun_time


class Channel:
    def __init__(self, id: int, server_id: int | None, channel_lang: str):
        self.id = id
        self.server_id = server_id
        self.channel_lang = channel_lang


class Server:
    def __init__(self, id: int, lang: str):
        self.id = id
        self.lang = lang


class ChannelIterator:
    def __init__(self, cursor: mysql.connection.MySQLCursor):
        self.cursor = cursor

    def __iter__(self):
        return self

    def __next__(self) -> Channel:
        res = self.cursor.fetchone()
        if res is None:
            self.cursor.close()
            raise StopIteration
        return Channel(res[0], res[1], res[2])


class DBAccess:
    def __init__(self, handler: logging.Handler = None):

        log = logging.Logger('DBAccess')

        if handler is not None:
            log.addHandler(handler)
        else:
            log.addHandler(logging.StreamHandler())

        self.connection = None

        for i in range(12):
            try:
                self.connection = mysql.connect(
                    host='localhost',
                    user=DB_USERNAME,
                    password=DB_PASSWORD,
                    database='hfc_db'
                )
                break
            except mysql.Error as e:
                self.connection = None
                log.warning(f"Couldn't connect to db. This is attempt #{i}")
                time.sleep(5)
                
        if self.connection is None:
            self.connection = mysql.connect(
                host='localhost',
                user=DB_USERNAME,
                password=DB_PASSWORD,
                database='hfc_db'
            )

    def add_area(self, area_id: int, area_name: str):
        with self.connection.cursor() as crsr:
            crsr.execute(f'REPLACE INTO areas (area_id, area_name) VALUES (%s, %s)', (area_id, area_name))
            self.connection.commit()

    def add_district(self, district_id: int, district_name: str, area_id: int, area_name: str, migun_time: int):
        with self.connection.cursor() as crsr:
            crsr.execute(f'SELECT * FROM areas WHERE area_id=%s', (area_id,))
            crsr.fetchall()

            if crsr.rowcount == 0:
                self.add_area(area_id, area_name)

            crsr.execute(f'REPLACE INTO districts (district_id, district_name, area_id, migun_time) VALUES (%s, %s, %s, %s)', (district_id, district_name, area_id, migun_time))
            self.connection.commit()

    def add_server(self, server_id: int, server_lang: str):
        with self.connection.cursor() as crsr:
            crsr.execute(f'INSERT IGNORE INTO servers (server_id, server_lang) VALUES (%s, %s)', (server_id, server_lang))
        self.connection.commit()

    def add_channel(self, channel_id: int, server_id: int | None, channel_lang: str | None):
        with self.connection.cursor() as crsr:
            if server_id is not None:
                self.add_server(server_id, channel_lang)
            crsr.execute(f'REPLACE INTO channels (channel_id, server_id, channel_lang) VALUES (%s, %s, %s)', (channel_id, server_id, channel_lang))
            self.connection.commit()

    def get_area(self, id: int) -> Area | None:
        with self.connection.cursor() as crsr:
            crsr.execute('SELECT * FROM areas WHERE area_id=%s', (id,))
            res = crsr.fetchone()
            crsr.fetchall()

        if res is not None:
            return Area(res[0], res[1])
        else:
            return None

    def get_district(self, id: int) -> District | None:
        with self.connection.cursor() as crsr:
            crsr.execute('SELECT * FROM districts WHERE district_id=%s', (id,))
            res = crsr.fetchone()
            crsr.fetchall()

        if res is not None:
            return District(res[0], res[1], res[2], res[3])
        else:
            return None

    def get_district_area(self, district: District) -> Area | None:
        return self.get_area(district.area_id)

    def get_server(self, id: int) -> Server | None:
        with self.connection.cursor() as crsr:
            crsr.execute('SELECT * FROM servers WHERE server_id=%s', (id,))
            res = crsr.fetchone()
            crsr.fetchall()

        if res is not None:
            return Server(res[0], res[1])
        else:
            return None

    def get_channel(self, id: int) -> Channel | None:
        with self.connection.cursor() as crsr:
            crsr.execute('SELECT * FROM channels WHERE channel_id=%s', (id,))
            res = crsr.fetchone()
            crsr.fetchall()

        if res is not None:
            return Channel(res[0], res[1], res[2])
        else:
            return None

    def get_channel_server(self, channel: Channel) -> Server:
        return self.get_server(channel.server_id)

    def channel_iterator(self):
        crsr = self.connection.cursor()

        crsr.execute('SELECT * FROM channels')

        return ChannelIterator(crsr)

    def remove_channel(self, id: int):
        print(id)
        with self.connection.cursor() as crsr:
            crsr.execute('DELETE FROM channels WHERE channel_id=%s', (id,))
        self.connection.commit()

    def remove_server(self, id: int):
        with self.connection.cursor() as crsr:
            crsr.execute('DELETE FROM channels WHERE server_id=%s', (id,))
            crsr.execute('DELETE FROM servers WHERE server_id=%s', (id,))
        self.connection.commit()

    def remove_district(self, id: int):
        with self.connection.cursor() as crsr:
            crsr.execute('DELETE FROM districts WHERE district_id=%s', (id,))
        self.connection.commit()

    def remove_area(self, id: int):
        with self.connection.cursor() as crsr:
            crsr.execute('DELETE FROM districts WHERE area_id=%s', (id,))
            crsr.execute('DELETE FROM areas WHERE area_id=%s', (id,))
        self.connection.commit()

    def get_district_by_name(self, name: str) -> District | None:
        with self.connection.cursor() as crsr:
            crsr.execute('SELECT * FROM districts WHERE district_name=%s', (name,))
            res = crsr.fetchone()
            crsr.fetchall()

        if res is None:
            return None

        return District(res[0], res[1], res[2], res[3])