Spaces:
Runtime error
Runtime error
import logging | |
import os | |
import sqlite3 | |
from contextlib import contextmanager | |
from typing import Generator | |
logger = logging.getLogger(__name__) | |
logger.setLevel(logging.DEBUG) | |
db_file = os.getenv("DB_FILE_NAME", "sqlite-data.db") | |
schema_entries = """ | |
CREATE TABLE entries | |
( | |
id TEXT PRIMARY KEY, | |
author TEXT NOT NULL, | |
source TEXT NOT NULL, | |
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP | |
) | |
""" | |
# create schema for 'summary' table, id is a uuid4 | |
schema_summary = """ | |
CREATE TABLE summaries | |
( | |
id INTEGER PRIMARY KEY AUTOINCREMENT, | |
entry_id TEXT NOT NULL, | |
summary TEXT NOT NULL, | |
summarizer_name TEXT NOT NULL, | |
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
FOREIGN KEY(entry_id) REFERENCES entries(id) | |
) | |
""" | |
schema_tag = """ | |
CREATE TABLE tags | |
( | |
id INTEGER PRIMARY KEY AUTOINCREMENT, | |
entry_id TEXT NOT NULL, | |
tag TEXT NOT NULL, | |
tagger_name TEXT NOT NULL, | |
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
FOREIGN KEY(entry_id) REFERENCES entries(id) | |
) | |
""" | |
schema_job = """ | |
CREATE TABLE jobs | |
( | |
id INTEGER PRIMARY KEY AUTOINCREMENT, | |
entry_id TEXT NOT NULL, | |
status TEXT NOT NULL DEFAULT 'pending', | |
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP, | |
FOREIGN KEY(entry_id) REFERENCES entries(id) | |
) | |
""" | |
TABLES = { | |
"entries": schema_entries, | |
"summaries": schema_summary, | |
"tags": schema_tag, | |
"jobs": schema_job, | |
} | |
TABLES_CREATED = False | |
def _get_db_connection() -> sqlite3.Connection: | |
global TABLES_CREATED | |
# sqlite cannot deal with concurrent access, so we set a big timeout | |
conn = sqlite3.connect(db_file, timeout=30) | |
if TABLES_CREATED: | |
return conn | |
cursor = conn.cursor() | |
# create tables if needed | |
for table_name, schema in TABLES.items(): | |
cursor.execute( | |
"SELECT name FROM sqlite_master WHERE type='table' AND name=?", | |
(table_name,), | |
) | |
table_exists = cursor.fetchone() is not None | |
if not table_exists: | |
logger.info(f"'{table_name}' table does not exist, creating it now...") | |
cursor.execute(schema) | |
conn.commit() | |
logger.info("done") | |
TABLES_CREATED = True | |
return conn | |
def get_db_cursor() -> Generator[sqlite3.Cursor, None, None]: | |
conn = _get_db_connection() | |
cursor = conn.cursor() | |
try: | |
yield cursor | |
finally: | |
conn.commit() | |
cursor.close() | |
conn.close() | |