43 lines
1.3 KiB
Python
43 lines
1.3 KiB
Python
|
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
|
||
|
|
);""")
|
||
|
|
|
||
|
|
def add_link(self, doc_filepath: str, link: str, total_cost: int):
|
||
|
|
with self.get_cursor() as cur:
|
||
|
|
cur.execute("INSERT INTO ltl (doc_filepath, link, total_cost) VALUES (?, ?, ?)",
|
||
|
|
(doc_filepath, link, total_cost))
|
||
|
|
self.con.commit()
|
||
|
|
|
||
|
|
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()
|