b2bcenter-parser/storage.py

import sqlite3
from contextlib import contextmanager
from typing import ContextManager


class Storage:
    def __init__(self):
        self.con = sqlite3.connect("database.db")

    @contextmanager
    def get_cursor(self) -> ContextManager[sqlite3.Cursor]:
        cur = self.con.cursor()
        try:
            yield cur
        finally:
            cur.close()

    def create_tables(self):
        with self.get_cursor() as cur:
            cur.execute("""
CREATE TABLE IF NOT EXISTS ltl (
                id INTEGER PRIMARY KEY,
                doc_filepath TEXT,
                link TEXT,
                total_cost INTEGER,
                query TEXT,
                answer TEXT
            );

""")
            cur.execute("""
            create unique index if not exists fp_link_index on ltl(doc_filepath, link);
            """)

            cur.execute("""
           create trigger if not exists query_answer_trig
            after insert
                on ltl
            begin
                update ltl set total_cost = new.total_cost,
                query = new.query,
                answer = new.answer
                where link = new.link;
            end; 
            """)

            cur.execute("""
            create table if not exists users (
            id INTEGER PRIMARY KEY,
            tg_user_id TEXT
);
            """)
            self.con.commit()

    def add_link(self, doc_filepath: str, link: str, total_cost: int, query: str, answer: str):
        with self.get_cursor() as cur:
            cur.execute("INSERT INTO ltl (doc_filepath, link, total_cost, query, answer) VALUES (?, ?, ?, ?, ?)",
                        (doc_filepath, link, total_cost, query, answer))
            self.con.commit()

    def get_links(self):
        with self.get_cursor() as cur:
            res = cur.execute("SELECT DISTINCT link FROM ltl")
            return [row[0] for row in res.fetchall()]

    def is_link_exists(self, link: str) -> bool:
        with self.get_cursor() as cur:
            res = cur.execute("SELECT * FROM ltl WHERE link = ?", (link,))
            return res.fetchone()

    def is_doc_exists(self, doc_filepath: str) -> bool:
        with self.get_cursor() as cur:
            res = cur.execute("SELECT * FROM ltl WHERE doc_filepath = ?", (doc_filepath,))
            return res.fetchone()

    def get_users(self):
        with self.get_cursor() as cur:
            res = cur.execute("SELECT DISTINCT tg_user_id FROM users")
            return [row[0] for row in res.fetchall()]


if __name__ == "__main__":
    s = Storage()
    print(s.get_users())