FactsDB/factsdb/database.py

"""
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
                }