dudong-v2/utils/db_schema.py

"""
Database Schema and Initialization

Defines SQLite database schema for storing analysis data.
Handles database creation and migrations.
"""

import sqlite3
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


def get_schema_version() -> int:
    """Get the current schema version."""
    return 1


def get_database_schema() -> str:
    """
    Get SQL schema for database initialization.
    
    Returns:
        str: SQL schema definition
    """
    return """
-- Analyses table: stores basic analysis information
CREATE TABLE IF NOT EXISTS analyses (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    report_id TEXT UNIQUE NOT NULL,
    timestamp DATETIME NOT NULL,
    device_id TEXT NOT NULL,
    overall_grade TEXT,
    grade_description TEXT,
    processing_time REAL,
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);

-- Analysis inputs table: stores input files information
CREATE TABLE IF NOT EXISTS analysis_inputs (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    analysis_id INTEGER NOT NULL,
    input_type TEXT NOT NULL,
    original_path TEXT NOT NULL,
    saved_path TEXT NOT NULL,
    file_size INTEGER,
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (analysis_id) REFERENCES analyses(id) ON DELETE CASCADE
);

-- Analysis results table: stores model predictions and results
CREATE TABLE IF NOT EXISTS analysis_results (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    analysis_id INTEGER NOT NULL,
    model_type TEXT NOT NULL,
    predicted_class TEXT,
    confidence REAL,
    probabilities TEXT,
    processing_time REAL,
    metadata TEXT,
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (analysis_id) REFERENCES analyses(id) ON DELETE CASCADE
);

-- Analysis visualizations table: stores visualization image metadata
CREATE TABLE IF NOT EXISTS analysis_visualizations (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    analysis_id INTEGER NOT NULL,
    visualization_type TEXT NOT NULL,
    file_path TEXT NOT NULL,
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (analysis_id) REFERENCES analyses(id) ON DELETE CASCADE
);

-- Create indexes for common queries
CREATE INDEX IF NOT EXISTS idx_analyses_report_id ON analyses(report_id);
CREATE INDEX IF NOT EXISTS idx_analyses_created_at ON analyses(created_at);
CREATE INDEX IF NOT EXISTS idx_analysis_inputs_analysis_id ON analysis_inputs(analysis_id);
CREATE INDEX IF NOT EXISTS idx_analysis_results_analysis_id ON analysis_results(analysis_id);
CREATE INDEX IF NOT EXISTS idx_analysis_visualizations_analysis_id ON analysis_visualizations(analysis_id);
"""


def init_database(db_path: Path) -> bool:
    """
    Initialize the database with schema.
    
    Args:
        db_path: Path to the database file
        
    Returns:
        bool: True if initialization successful, False otherwise
    """
    try:
        # Ensure parent directory exists
        db_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Connect to database
        conn = sqlite3.connect(str(db_path))
        cursor = conn.cursor()
        
        # Execute schema
        cursor.executescript(get_database_schema())
        
        # Create metadata table to track schema version
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS schema_version (
                version INTEGER PRIMARY KEY,
                updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # Check if version exists
        cursor.execute("SELECT version FROM schema_version LIMIT 1")
        result = cursor.fetchone()
        
        if not result:
            # Insert initial version
            cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (get_schema_version(),))
        
        conn.commit()
        conn.close()
        
        logger.info(f"Database initialized successfully: {db_path}")
        return True
        
    except Exception as e:
        logger.error(f"Error initializing database: {e}")
        return False


def get_database_connection(db_path: Path) -> Optional[sqlite3.Connection]:
    """
    Get a connection to the database.
    
    Args:
        db_path: Path to the database file
        
    Returns:
        sqlite3.Connection or None if connection failed
    """
    try:
        # Initialize database if doesn't exist
        if not db_path.exists():
            init_database(db_path)
        
        # Enable foreign keys
        conn = sqlite3.connect(str(db_path))
        conn.execute("PRAGMA foreign_keys = ON")
        
        # Use Row factory for easier dict-like access
        conn.row_factory = sqlite3.Row
        
        return conn
        
    except Exception as e:
        logger.error(f"Error connecting to database: {e}")
        return None


def close_database_connection(conn: sqlite3.Connection) -> None:
    """
    Close a database connection.
    
    Args:
        conn: Connection to close
    """
    try:
        if conn:
            conn.close()
    except Exception as e:
        logger.error(f"Error closing database connection: {e}")