Add detailed database schema documentation for default table structure
This commit is contained in:
parent
e775d8de81
commit
9344a69636
@ -1,34 +1,31 @@
|
|||||||
"""
|
"""
|
||||||
Database layer for FactsDB service using SQLite
|
Database module for FactsDB service
|
||||||
|
Handles SQLite database operations with proper locking
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import sqlite3
|
import sqlite3
|
||||||
import os
|
|
||||||
import threading
|
import threading
|
||||||
|
import os
|
||||||
from typing import List, Dict, Any, Optional
|
from typing import List, Dict, Any, Optional
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from .config import DatabaseConfig
|
from .config import Config
|
||||||
|
|
||||||
class DatabaseManager:
|
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.config = config
|
||||||
self._lock = threading.Lock()
|
self._lock = threading.Lock()
|
||||||
self.init_database()
|
self._init_database()
|
||||||
|
|
||||||
def get_connection(self):
|
def _init_database(self):
|
||||||
"""Get database connection"""
|
"""Initialize the database and create tables if they don't exist"""
|
||||||
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"""
|
|
||||||
with self._lock:
|
with self._lock:
|
||||||
with self.get_connection() as conn:
|
conn = sqlite3.connect(self.config.path, check_same_thread=False)
|
||||||
# Create tables for different fact types
|
cursor = conn.cursor()
|
||||||
conn.execute('''
|
|
||||||
|
# Create facts table
|
||||||
|
cursor.execute('''
|
||||||
CREATE TABLE IF NOT EXISTS facts (
|
CREATE TABLE IF NOT EXISTS facts (
|
||||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
table_name TEXT NOT NULL,
|
table_name TEXT NOT NULL,
|
||||||
@ -40,134 +37,241 @@ class DatabaseManager:
|
|||||||
key_dates TEXT,
|
key_dates TEXT,
|
||||||
main_points TEXT,
|
main_points TEXT,
|
||||||
file_path TEXT,
|
file_path TEXT,
|
||||||
extracted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
extracted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
processed BOOLEAN DEFAULT FALSE
|
||||||
)
|
)
|
||||||
''')
|
''')
|
||||||
|
|
||||||
# Create table for tracking processed files
|
# Create tables metadata table
|
||||||
conn.execute('''
|
cursor.execute('''
|
||||||
CREATE TABLE IF NOT EXISTS processed_files (
|
CREATE TABLE IF NOT EXISTS tables (
|
||||||
|
name TEXT PRIMARY KEY,
|
||||||
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
record_count INTEGER DEFAULT 0
|
||||||
|
)
|
||||||
|
''')
|
||||||
|
|
||||||
|
# Create files table to track processed files
|
||||||
|
cursor.execute('''
|
||||||
|
CREATE TABLE IF NOT EXISTS files (
|
||||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
file_path TEXT UNIQUE,
|
path TEXT UNIQUE NOT NULL,
|
||||||
processed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
table_name TEXT NOT NULL,
|
||||||
table_name TEXT
|
processed BOOLEAN DEFAULT FALSE,
|
||||||
|
processed_at TIMESTAMP,
|
||||||
|
error TEXT
|
||||||
)
|
)
|
||||||
''')
|
''')
|
||||||
|
|
||||||
# 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)')
|
|
||||||
|
|
||||||
conn.commit()
|
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):
|
def create_table(self, table_name: str):
|
||||||
"""Create a new table for storing facts"""
|
"""Create a new table entry"""
|
||||||
# In SQLite, tables are created automatically when inserting data
|
|
||||||
# This method ensures the table exists and can be used for future enhancements
|
|
||||||
pass
|
|
||||||
|
|
||||||
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:
|
with self.get_connection() as conn:
|
||||||
conn.execute('''
|
cursor = conn.cursor()
|
||||||
INSERT INTO facts
|
cursor.execute(
|
||||||
(table_name, title, summary, main_topic, key_entities,
|
"INSERT OR REPLACE INTO tables (name, record_count) VALUES (?, 0)",
|
||||||
financial_impact, key_dates, main_points, file_path)
|
(table_name,)
|
||||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
)
|
||||||
|
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 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 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,
|
table_name,
|
||||||
fact_data.get('title'),
|
fact_data.get('title'),
|
||||||
fact_data.get('summary'),
|
fact_data.get('summary'),
|
||||||
fact_data.get('main_topic'),
|
fact_data.get('main_topic'),
|
||||||
str(fact_data.get('key_entities', [])),
|
key_entities,
|
||||||
fact_data.get('financial_impact'),
|
fact_data.get('financial_impact'),
|
||||||
str(fact_data.get('key_dates', [])),
|
key_dates,
|
||||||
str(fact_data.get('main_points', [])),
|
main_points,
|
||||||
file_path
|
fact_data.get('file_path')
|
||||||
))
|
))
|
||||||
|
|
||||||
|
fact_id = cursor.lastrowid
|
||||||
conn.commit()
|
conn.commit()
|
||||||
|
|
||||||
def mark_file_processed(self, file_path: str, table_name: str):
|
# Update table record count
|
||||||
"""Mark file as processed"""
|
self._update_table_record_count(table_name)
|
||||||
with self._lock:
|
|
||||||
|
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:
|
with self.get_connection() as conn:
|
||||||
conn.execute('''
|
cursor = conn.cursor()
|
||||||
INSERT OR IGNORE INTO processed_files (file_path, table_name)
|
cursor.execute("SELECT COUNT(*) FROM facts WHERE table_name = ?", (table_name,))
|
||||||
VALUES (?, ?)
|
count = cursor.fetchone()[0]
|
||||||
''', (file_path, table_name))
|
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 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")
|
||||||
|
|
||||||
|
columns = [description[0] for description in cursor.description]
|
||||||
|
rows = cursor.fetchall()
|
||||||
|
|
||||||
|
return [dict(zip(columns, row)) for row in rows]
|
||||||
|
|
||||||
|
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]
|
||||||
|
|
||||||
|
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()
|
conn.commit()
|
||||||
|
|
||||||
def is_file_processed(self, file_path: str) -> bool:
|
def is_file_processed(self, file_path: str) -> bool:
|
||||||
"""Check if file has been processed"""
|
"""Check if a file has been processed"""
|
||||||
with self._lock:
|
|
||||||
with self.get_connection() as conn:
|
with self.get_connection() as conn:
|
||||||
cursor = conn.execute('''
|
cursor = conn.cursor()
|
||||||
SELECT COUNT(*) as count FROM processed_files WHERE file_path = ?
|
cursor.execute("SELECT processed FROM files WHERE path = ?", (file_path,))
|
||||||
''', (file_path,))
|
row = cursor.fetchone()
|
||||||
result = cursor.fetchone()
|
return row[0] if row else False
|
||||||
return result['count'] > 0
|
|
||||||
|
|
||||||
def get_all_tables(self) -> List[Dict[str, Any]]:
|
def get_processed_files(self, table_name: str) -> List[str]:
|
||||||
"""Get all available tables with record counts"""
|
"""Get list of processed files for a table"""
|
||||||
with self._lock:
|
|
||||||
with self.get_connection() as conn:
|
with self.get_connection() as conn:
|
||||||
cursor = conn.execute('''
|
cursor = conn.cursor()
|
||||||
SELECT DISTINCT table_name, COUNT(*) as count
|
cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 1", (table_name,))
|
||||||
FROM facts
|
return [row[0] for row in cursor.fetchall()]
|
||||||
GROUP BY table_name
|
|
||||||
''')
|
|
||||||
return [dict(row) for row in cursor.fetchall()]
|
|
||||||
|
|
||||||
def query_table(self, table_name: str, query: Optional[str] = None) -> List[Dict[str, Any]]:
|
def get_unprocessed_files(self, table_name: str) -> List[str]:
|
||||||
"""Query facts from specific table"""
|
"""Get list of unprocessed files for a table"""
|
||||||
with self._lock:
|
|
||||||
with self.get_connection() as conn:
|
with self.get_connection() as conn:
|
||||||
if query:
|
cursor = conn.cursor()
|
||||||
# Simple query support - in a real implementation, this would be more sophisticated
|
cursor.execute("SELECT path FROM files WHERE table_name = ? AND processed = 0", (table_name,))
|
||||||
cursor = conn.execute(f'''
|
return [row[0] for row in cursor.fetchall()]
|
||||||
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()]
|
|
||||||
|
|
||||||
def get_table_record_count(self, table_name: str) -> int:
|
# Global database manager instance
|
||||||
"""Get record count for a specific table"""
|
_db_manager = None
|
||||||
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
|
|
||||||
|
|
||||||
def get_database_stats(self) -> Dict[str, Any]:
|
def get_db_manager(config) -> DatabaseManager:
|
||||||
"""Get database statistics for monitoring"""
|
"""Get the global database manager instance"""
|
||||||
with self._lock:
|
global _db_manager
|
||||||
with self.get_connection() as conn:
|
if _db_manager is None:
|
||||||
# Get total facts
|
_db_manager = DatabaseManager(config)
|
||||||
cursor = conn.execute('SELECT COUNT(*) as total FROM facts')
|
return _db_manager
|
||||||
total_facts = cursor.fetchone()['total']
|
|
||||||
|
|
||||||
# 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()}
|
|
||||||
|
|
||||||
# Get processed files count
|
|
||||||
cursor = conn.execute('SELECT COUNT(*) as total FROM processed_files')
|
|
||||||
processed_files = cursor.fetchone()['total']
|
|
||||||
|
|
||||||
return {
|
|
||||||
'total_facts': total_facts,
|
|
||||||
'table_counts': table_counts,
|
|
||||||
'processed_files': processed_files,
|
|
||||||
'database_file': self.config.connection_string
|
|
||||||
}
|
|
||||||
Loading…
x
Reference in New Issue
Block a user