check_subs/sqliter.py

import sqlite3

from utils import User, Ads, ScheduledMessage


# from loader import logger as log


class Database:
    def __init__(self, path_to_db="data/main.db"):
        self.path_to_db = path_to_db

    @property
    def connection(self):
        return sqlite3.connect(self.path_to_db)

    def execute(self, sql: str, parameters: tuple = None, fetchone=False, fetchall=False, commit=False):
        if not parameters:
            parameters = tuple()

        connection = self.connection
        connection.set_trace_callback(logger)
        cursor = connection.cursor()
        data = None
        cursor.execute(sql, parameters)
        if commit:
            connection.commit()
        if fetchone:
            data = cursor.fetchone()
        if fetchall:
            data = cursor.fetchall()
        connection.close()

        return data

    def create_table_(self):
        sql = ["""CREATE TABLE Users (
                id int NOT NULL,
                username varchar(255) NULL,
                Name varchar(255) NOT NULL,
                PRIMARY KEY(id)
                ); """,
               """CREATE TABLE "Chats" (
                   "client_id"	INTEGER NOT NULL,
                   "chat_id"	INTEGER NOT NULL,
                   "require_channel_id"	INTEGER NOT NULL
               )""",
               """CREATE TABLE "stop_words" (
                "client_id"	INTEGER,
                "word"	TEXT)""",
               """CREATE TABLE "ads" (
                "id"	INTEGER NOT NULL,
                "text"	TEXT,
                "photo"	TEXT DEFAULT NULL,
                PRIMARY KEY(id)
                )""",
               """CREATE TABLE "jobs" (
                "message_id"	INTEGER NOT NULL,
                "chat_id"	TEXT DEFAULT NULL,
                "text"	TEXT
                );""", ]
        for t in sql:
            try:
                self.execute(t, commit=True)
            except:
                pass

    def add_client(self, client_id: int, fullname: str, username):
        sql = "INSERT INTO Users VALUES(?,?,?);"
        parameters = (client_id, fullname, username,)
        try:
            self.execute(sql, parameters, commit=True)
        except sqlite3.IntegrityError:
            pass

    def get_all_clients(self):
        """
        return a list of all clients with id, chats, stop_words, require_channel_id
        """
        users = [c_id[0] for c_id in self.execute("SELECT DISTINCT(id) FROM Users", fetchall=True)]
        out = []
        for user in users:
            stop_words = [i[0] for i in
                          self.execute("SELECT word from stop_words WHERE client_id=?", (user,), fetchall=True)]
            chats = self.execute("SELECT chat_id, require_channel_id from Chats WHERE client_id=?", (user,),
                                 fetchall=True)
            require_channel_id = {}
            for chat in chats:
                require_channel_id[chat[0]] = []
            for chat in chats:
                if chat[1] != 0:
                    require_channel_id[chat[0]].append(chat[1])
            usr = User(client_id=user, chats_id=[c[0] for c in chats], stop_words=stop_words,
                       require_channel_id=require_channel_id)
            out.append(usr)
        return out

    def add_word(self, client_id: int, word: str):
        sql = "INSERT INTO stop_words VALUES(?,?);"
        parameters = (client_id, word,)
        self.execute(sql, parameters, commit=True)

    def select_all_word(self):
        clients = self.execute('SELECT DISTINCT(client_id) FROM stop_words', fetchall=True)
        words = {}
        for client_id in clients:
            sql = "SELECT word FROM stop_words where client_id=?"
            word = self.execute(sql, (client_id[0],), fetchall=True)
            words[client_id[0]] = [w[0] for w in word]

        return words

    def select_default_word(self):
        sql = "SELECT word FROM stop_words WHERE client_id=0"
        return [i[0] for i in self.execute(sql, fetchall=True)]

    def add_default_words(self, client_id):
        words = self.select_default_word()
        for word in words:
            self.add_word(client_id, word)

    def select_all_word_in_client(self, client_id):
        sql = "SELECT * FROM stop_words WHERE client_id=?"
        parameters = (client_id,)
        return self.execute(sql, parameters, fetchall=True)

    def delete_word(self, client_id: int, word: str):
        sql = 'DELETE FROM stop_words WHERE client_id=? AND word=?'
        return self.execute(sql, parameters=(client_id, word,), commit=True)

    @staticmethod
    def format_args(sql, parameters: dict):
        sql += " AND".join([
            f" {item} = ? " for item in parameters
        ])
        return sql, tuple(parameters.values())

    # def add_channel(self, user_id: int, chat_id: int, channel_id: int):
    def add_channel(self, client_id: int, chat_id: int, channel_id: int):
        sql = "INSERT INTO Chats VALUES(?,?,?);"
        parameters = (client_id, chat_id, channel_id,)
        self.execute(sql, parameters, commit=True)

    def delete_channel(self, client_id, chat_id, channel_id):
        sql = 'DELETE FROM Chats WHERE client_id=? AND chat_id = ? AND require_channel_id = ?'
        # sql, parameters = self.format_args(sql, kwargs)
        # return self.execute(sql, parameters, fetchone=True)
        params = (client_id, chat_id, channel_id,)
        return self.execute(sql, params, commit=True)

    def select_channels_by_user(self, client_id):
        sql = f'SELECT DISTINCT(chat_id) FROM Chats WHERE client_id=? AND require_channel_id=0;'
        params = (client_id,)
        return [i[0] for i in self.execute(sql, params, fetchall=True)]

    def select_channels_by_chat(self, client_id, chat_id=None, action='list'):
        if action == 'add':
            # TODO get chats except for those already assigned
            sql = f'SELECT DISTINCT(chat_id) FROM Chats WHERE client_id=? AND require_channel_id=0 OR require_channel_id=2;'
            params = (client_id,)
        elif action == 'delete':
            sql = f'SELECT DISTINCT(require_channel_id) FROM Chats WHERE client_id=? AND chat_id=? AND require_channel_id!=0;'
            params = (client_id, chat_id)
        elif action == 'list':
            sql = f'SELECT DISTINCT(require_channel_id) FROM Chats WHERE client_id=? AND chat_id=? ' \
                  f'AND require_channel_id!=0 AND require_channel_id!=2;'
            params = (client_id, chat_id,)

        return [i[0] for i in self.execute(sql, params, fetchall=True)]

    def select_all_chats(self):
        chats = [i[0] for i in self.execute('SELECT DISTINCT chat_id FROM Chats;', fetchall=True)]
        out = {}
        return chats

    def add_chat(self, client_id, chat_id, chat_type):
        '''
        type: group chat = 0
        group chat with req chat: 1
        channels: 2
        '''
        sql = "INSERT INTO Chats VALUES (?, ?, ?)"
        return self.execute(sql, parameters=(client_id, chat_id, chat_type), commit=True)

    def add_require_channel(self, client_id, chat_id, req_channel_id):
        # sql = "UPDATE Chats SET require_channel_id=? WHERE chat_id=?"
        sql = "INSERT INTO Chats VALUES (?,?,?)"
        return self.execute(sql, parameters=(client_id, chat_id, req_channel_id,), commit=True)

    def delete_require_channel(self, client_id, chat_id, req_channel_id):
        # sql = "UPDATE Chats SET require_channel_id=? WHERE chat_id=?"
        sql = "DELETE FROM Chats WHERE client_id=? AND chat_id=? AND require_channel_id=?"
        return self.execute(sql, parameters=(client_id, chat_id, req_channel_id,), commit=True)

    def select_all_ads(self):
        chats = [Ads(id=i[0], text=i[1], path_photo=i[2]) for i in self.execute('SELECT * FROM ads;', fetchall=True)]
        return chats

    def add_ads(self, text, photo=None):
        sql = "INSERT INTO ads(text, photo) VALUES (?,?)"
        return self.execute(sql, parameters=(text, photo), commit=True)

    def delete_ads(self, id: int):
        sql = "DELETE FROM ads WHERE id=?"
        return self.execute(sql, parameters=(id,), commit=True)

    def select_channels_for_ads(self):
        sql = f'SELECT DISTINCT(chat_id) FROM Chats WHERE require_channel_id!=2'
        # sql = f'SELECT DISTINCT(chat_id) FROM Chats WHERE require_channel_id!=0 AND require_channel_id!=2;'
        return [i[0] for i in self.execute(sql, fetchall=True)]

    def get_ads(self, id: int):
        c = self.execute('SELECT * FROM ads WHERE id=?;', (id,), fetchone=True)
        chats = Ads(id=c[0], text=c[1], path_photo=c[2])
        return chats

    def jobs_add(self, message_id, chat_id, text):
        sql = "INSERT INTO jobs VALUES (?,?,?)"
        return self.execute(sql, parameters=(message_id, chat_id, text), commit=True)

    def jobs_delete(self, message_id, chat_id, text):
        message_id = 0 if message_id is None else message_id
        sql = "DELETE FROM jobs WHERE message_id=? AND chat_id=? AND text=?"
        return self.execute(sql, parameters=(message_id, chat_id, text), commit=True)

    def get_jobs(self):
        jobs = []
        offset = 0
        while True:
            jobs_list = [ScheduledMessage(
                message_id=None if int(i[0]) == 0 else i[0],
                chat_id=i[1], text=i[2]
            ) for i in
                self.execute(f'SELECT * FROM jobs LIMIT 5 OFFSET {offset};', fetchall=True)]
            jobs.append(jobs_list)
            j = len(jobs_list)
            if j < 5 or len(jobs) > 10:
                break
            offset += 5

        return jobs[:10]


def logger(statement):
    # log.success(statement)
    print(f"""{'=' * 40}\nEXECUTING: {statement}\n{'=' * 40}""")