""" Database layer for FactsDB service using SQLite """ import sqlite3 import os import threading from typing import List, Dict, Any, Optional from contextlib import contextmanager from .config import DatabaseConfig class DatabaseManager: """Manages database operations for FactsDB""" def __init__(self, config: DatabaseConfig): self.config = config self._lock = threading.Lock() 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""" 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 ) ''') # 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 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() 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 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 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 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 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()] 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()] 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 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'] # 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 }