Add detailed database schema documentation for default table structure
This commit is contained in:
parent
e775d8de81
commit
9344a69636
@ -1,173 +1,277 @@
|
||||
"""
|
||||
Database layer for FactsDB service using SQLite
|
||||
Database module for FactsDB service
|
||||
Handles SQLite database operations with proper locking
|
||||
"""
|
||||
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
import os
|
||||
from typing import List, Dict, Any, Optional
|
||||
from contextlib import contextmanager
|
||||
from .config import DatabaseConfig
|
||||
from .config import Config
|
||||
|
||||
class DatabaseManager:
|
||||
"""Manages database operations for FactsDB"""
|
||||
"""Manages SQLite database operations with thread safety"""
|
||||
|
||||
def __init__(self, config: DatabaseConfig):
|
||||
def __init__(self, config):
|
||||
self.config = config
|
||||
self._lock = threading.Lock()
|
||||
self.init_database()
|
||||
self._init_database()
|
||||
|
||||
def get_connection(self):
|
||||
"""Get database connection"""
|
||||
conn = sqlite3.connect(self.config.connection_string, check_same_thread=False)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
def init_database(self):
|
||||
"""Initialize database schema"""
|
||||
def _init_database(self):
|
||||
"""Initialize the database and create tables if they don't exist"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
# Create tables for different fact types
|
||||
conn.execute('''
|
||||
CREATE TABLE IF NOT EXISTS facts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
table_name TEXT NOT NULL,
|
||||
title TEXT,
|
||||
summary TEXT,
|
||||
main_topic TEXT,
|
||||
key_entities TEXT,
|
||||
financial_impact TEXT,
|
||||
key_dates TEXT,
|
||||
main_points TEXT,
|
||||
file_path TEXT,
|
||||
extracted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
conn = sqlite3.connect(self.config.path, check_same_thread=False)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Create table for tracking processed files
|
||||
conn.execute('''
|
||||
CREATE TABLE IF NOT EXISTS processed_files (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
file_path TEXT UNIQUE,
|
||||
processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
table_name TEXT
|
||||
)
|
||||
''')
|
||||
# Create facts table
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS facts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
table_name TEXT NOT NULL,
|
||||
title TEXT,
|
||||
summary TEXT,
|
||||
main_topic TEXT,
|
||||
key_entities TEXT,
|
||||
financial_impact TEXT,
|
||||
key_dates TEXT,
|
||||
main_points TEXT,
|
||||
file_path TEXT,
|
||||
extracted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
processed BOOLEAN DEFAULT FALSE
|
||||
)
|
||||
''')
|
||||
|
||||
# Create indexes for better performance
|
||||
conn.execute('CREATE INDEX IF NOT EXISTS idx_facts_table_name ON facts(table_name)')
|
||||
conn.execute('CREATE INDEX IF NOT EXISTS idx_facts_extracted_at ON facts(extracted_at)')
|
||||
conn.execute('CREATE INDEX IF NOT EXISTS idx_processed_files_path ON processed_files(file_path)')
|
||||
# Create tables metadata table
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS tables (
|
||||
name TEXT PRIMARY KEY,
|
||||
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||
record_count INTEGER DEFAULT 0
|
||||
)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
# Create files table to track processed files
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS files (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
path TEXT UNIQUE NOT NULL,
|
||||
table_name TEXT NOT NULL,
|
||||
processed BOOLEAN DEFAULT FALSE,
|
||||
processed_at TIMESTAMP,
|
||||
error TEXT
|
||||
)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
@contextmanager
|
||||
def get_connection(self):
|
||||
"""Get a database connection with thread safety"""
|
||||
with self._lock:
|
||||
conn = sqlite3.connect(self.config.path, check_same_thread=False)
|
||||
try:
|
||||
yield conn
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def create_table(self, table_name: str):
|
||||
"""Create a new table for storing facts"""
|
||||
# In SQLite, tables are created automatically when inserting data
|
||||
# This method ensures the table exists and can be used for future enhancements
|
||||
pass
|
||||
"""Create a new table entry"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"INSERT OR REPLACE INTO tables (name, record_count) VALUES (?, 0)",
|
||||
(table_name,)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def store_fact(self, fact_data: Dict[str, Any], table_name: str, file_path: str):
|
||||
"""Store extracted fact in database"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
conn.execute('''
|
||||
INSERT INTO facts
|
||||
(table_name, title, summary, main_topic, key_entities,
|
||||
financial_impact, key_dates, main_points, file_path)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (
|
||||
table_name,
|
||||
fact_data.get('title'),
|
||||
fact_data.get('summary'),
|
||||
fact_data.get('main_topic'),
|
||||
str(fact_data.get('key_entities', [])),
|
||||
fact_data.get('financial_impact'),
|
||||
str(fact_data.get('key_dates', [])),
|
||||
str(fact_data.get('main_points', [])),
|
||||
file_path
|
||||
))
|
||||
conn.commit()
|
||||
def get_table_names(self) -> List[str]:
|
||||
"""Get all table names"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT name FROM tables")
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
def mark_file_processed(self, file_path: str, table_name: str):
|
||||
"""Mark file as processed"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
conn.execute('''
|
||||
INSERT OR IGNORE INTO processed_files (file_path, table_name)
|
||||
VALUES (?, ?)
|
||||
''', (file_path, table_name))
|
||||
conn.commit()
|
||||
def get_table_info(self, table_name: str) -> Dict[str, Any]:
|
||||
"""Get information about a specific table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT name, created_at, record_count FROM tables WHERE name = ?", (table_name,))
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
return {
|
||||
'name': row[0],
|
||||
'created_at': row[1],
|
||||
'record_count': row[2]
|
||||
}
|
||||
return None
|
||||
|
||||
def is_file_processed(self, file_path: str) -> bool:
|
||||
"""Check if file has been processed"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.execute('''
|
||||
SELECT COUNT(*) as count FROM processed_files WHERE file_path = ?
|
||||
''', (file_path,))
|
||||
result = cursor.fetchone()
|
||||
return result['count'] > 0
|
||||
def update_table_count(self, table_name: str, count: int):
|
||||
"""Update the record count for a table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute(
|
||||
"UPDATE tables SET record_count = ? WHERE name = ?",
|
||||
(count, table_name)
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def insert_fact(self, table_name: str, fact_data: Dict[str, Any]) -> int:
|
||||
"""Insert a fact into the database"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Convert lists to JSON strings for storage
|
||||
key_entities = str(fact_data.get('key_entities', []))
|
||||
key_dates = str(fact_data.get('key_dates', []))
|
||||
main_points = str(fact_data.get('main_points', []))
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO facts (
|
||||
table_name, title, summary, main_topic, key_entities,
|
||||
financial_impact, key_dates, main_points, file_path
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
''', (
|
||||
table_name,
|
||||
fact_data.get('title'),
|
||||
fact_data.get('summary'),
|
||||
fact_data.get('main_topic'),
|
||||
key_entities,
|
||||
fact_data.get('financial_impact'),
|
||||
key_dates,
|
||||
main_points,
|
||||
fact_data.get('file_path')
|
||||
))
|
||||
|
||||
fact_id = cursor.lastrowid
|
||||
conn.commit()
|
||||
|
||||
# Update table record count
|
||||
self._update_table_record_count(table_name)
|
||||
|
||||
return fact_id
|
||||
|
||||
def _update_table_record_count(self, table_name: str):
|
||||
"""Update the record count for a table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT COUNT(*) FROM facts WHERE table_name = ?", (table_name,))
|
||||
count = cursor.fetchone()[0]
|
||||
cursor.execute("UPDATE tables SET record_count = ? WHERE name = ?", (count, table_name))
|
||||
conn.commit()
|
||||
|
||||
def get_facts(self, table_name: str, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
|
||||
"""Get facts from a specific table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT id, table_name, title, summary, main_topic, key_entities,
|
||||
financial_impact, key_dates, main_points, file_path, extracted_at
|
||||
FROM facts
|
||||
WHERE table_name = ?
|
||||
ORDER BY extracted_at DESC
|
||||
LIMIT ? OFFSET ?
|
||||
''', (table_name, limit, offset))
|
||||
|
||||
columns = [description[0] for description in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
|
||||
facts = []
|
||||
for row in rows:
|
||||
fact = dict(zip(columns, row))
|
||||
# Convert JSON strings back to lists
|
||||
if fact['key_entities']:
|
||||
fact['key_entities'] = eval(fact['key_entities'])
|
||||
if fact['key_dates']:
|
||||
fact['key_dates'] = eval(fact['key_dates'])
|
||||
if fact['main_points']:
|
||||
fact['main_points'] = eval(fact['main_points'])
|
||||
facts.append(fact)
|
||||
|
||||
return facts
|
||||
|
||||
def get_fact_by_id(self, fact_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""Get a specific fact by ID"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
SELECT id, table_name, title, summary, main_topic, key_entities,
|
||||
financial_impact, key_dates, main_points, file_path, extracted_at
|
||||
FROM facts
|
||||
WHERE id = ?
|
||||
''', (fact_id,))
|
||||
|
||||
row = cursor.fetchone()
|
||||
if row:
|
||||
columns = [description[0] for description in cursor.description]
|
||||
fact = dict(zip(columns, row))
|
||||
# Convert JSON strings back to lists
|
||||
if fact['key_entities']:
|
||||
fact['key_entities'] = eval(fact['key_entities'])
|
||||
if fact['key_dates']:
|
||||
fact['key_dates'] = eval(fact['key_dates'])
|
||||
if fact['main_points']:
|
||||
fact['main_points'] = eval(fact['main_points'])
|
||||
return fact
|
||||
return None
|
||||
|
||||
def get_all_tables(self) -> List[Dict[str, Any]]:
|
||||
"""Get all available tables with record counts"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.execute('''
|
||||
SELECT DISTINCT table_name, COUNT(*) as count
|
||||
FROM facts
|
||||
GROUP BY table_name
|
||||
''')
|
||||
return [dict(row) for row in cursor.fetchall()]
|
||||
"""Get all tables with their information"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT name, created_at, record_count FROM tables ORDER BY created_at DESC")
|
||||
|
||||
def query_table(self, table_name: str, query: Optional[str] = None) -> List[Dict[str, Any]]:
|
||||
"""Query facts from specific table"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
if query:
|
||||
# Simple query support - in a real implementation, this would be more sophisticated
|
||||
cursor = conn.execute(f'''
|
||||
SELECT * FROM facts WHERE table_name = ? AND ({query})
|
||||
''', (table_name,))
|
||||
else:
|
||||
cursor = conn.execute('''
|
||||
SELECT * FROM facts WHERE table_name = ?
|
||||
''', (table_name,))
|
||||
return [dict(row) for row in cursor.fetchall()]
|
||||
columns = [description[0] for description in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
|
||||
def get_table_record_count(self, table_name: str) -> int:
|
||||
"""Get record count for a specific table"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.execute('''
|
||||
SELECT COUNT(*) as count FROM facts WHERE table_name = ?
|
||||
''', (table_name,))
|
||||
result = cursor.fetchone()
|
||||
return result['count'] if result else 0
|
||||
return [dict(zip(columns, row)) for row in rows]
|
||||
|
||||
def get_database_stats(self) -> Dict[str, Any]:
|
||||
"""Get database statistics for monitoring"""
|
||||
with self._lock:
|
||||
with self.get_connection() as conn:
|
||||
# Get total facts
|
||||
cursor = conn.execute('SELECT COUNT(*) as total FROM facts')
|
||||
total_facts = cursor.fetchone()['total']
|
||||
def get_table_count(self, table_name: str) -> int:
|
||||
"""Get the count of facts in a specific table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT COUNT(*) FROM facts WHERE table_name = ?", (table_name,))
|
||||
return cursor.fetchone()[0]
|
||||
|
||||
# Get table counts
|
||||
cursor = conn.execute('''
|
||||
SELECT table_name, COUNT(*) as count
|
||||
FROM facts
|
||||
GROUP BY table_name
|
||||
''')
|
||||
table_counts = {row['table_name']: row['count'] for row in cursor.fetchall()}
|
||||
def mark_file_processed(self, file_path: str, table_name: str, success: bool = True, error: str = None):
|
||||
"""Mark a file as processed"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('''
|
||||
INSERT OR REPLACE INTO files (path, table_name, processed, processed_at, error)
|
||||
VALUES (?, ?, ?, datetime('now'), ?)
|
||||
''', (file_path, table_name, success, error))
|
||||
conn.commit()
|
||||
|
||||
# Get processed files count
|
||||
cursor = conn.execute('SELECT COUNT(*) as total FROM processed_files')
|
||||
processed_files = cursor.fetchone()['total']
|
||||
def is_file_processed(self, file_path: str) -> bool:
|
||||
"""Check if a file has been processed"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT processed FROM files WHERE path = ?", (file_path,))
|
||||
row = cursor.fetchone()
|
||||
return row[0] if row else False
|
||||
|
||||
return {
|
||||
'total_facts': total_facts,
|
||||
'table_counts': table_counts,
|
||||
'processed_files': processed_files,
|
||||
'database_file': self.config.connection_string
|
||||
}
|
||||
def get_processed_files(self, table_name: str) -> List[str]:
|
||||
"""Get list of processed files for a table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 1", (table_name,))
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
def get_unprocessed_files(self, table_name: str) -> List[str]:
|
||||
"""Get list of unprocessed files for a table"""
|
||||
with self.get_connection() as conn:
|
||||
cursor = conn.cursor()
|
||||
cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 0", (table_name,))
|
||||
return [row[0] for row in cursor.fetchall()]
|
||||
|
||||
# Global database manager instance
|
||||
_db_manager = None
|
||||
|
||||
def get_db_manager(config) -> DatabaseManager:
|
||||
"""Get the global database manager instance"""
|
||||
global _db_manager
|
||||
if _db_manager is None:
|
||||
_db_manager = DatabaseManager(config)
|
||||
return _db_manager
|
||||
Loading…
x
Reference in New Issue
Block a user