netbackup-insights/database.py

import os
import sys
import sqlite3
import pandas as pd

def get_db_path():
    """
    Returns the absolute path to the SQLite database.
    If the application is running inside a Docker container, it uses /app/data/nbu_insights.db.
    If running locally, it defaults to a local './data' directory.
    """
    # Check if we are running in the container workspace or local workspace
    container_data_dir = "/app/data"
    if os.path.exists(container_data_dir) or os.environ.get("NBU_INSIGHTS_RUNNING") == "1":
        db_dir = container_data_dir
    else:
        # Local development fallback
        base_dir = os.path.dirname(os.path.abspath(__file__))
        db_dir = os.path.join(base_dir, "data")
        
    os.makedirs(db_dir, exist_ok=True)
    return os.path.join(db_dir, "nbu_insights.db")

def get_connection():
    """
    Establishes and returns a connection to the SQLite database.
    Rows are configured to be accessible by column names like a dictionary.
    """
    db_path = get_db_path()
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    return conn

def init_db():
    """
    Initializes the SQLite database tables (processed_files, backup_jobs, and job_actions)
    according to the v2.0 spec.
    """
    conn = get_connection()
    cursor = conn.cursor()
    
    # Track processed files
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS processed_files (
        file_hash TEXT PRIMARY KEY,
        file_name TEXT,
        upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    );
    """)
    
    # Store backup jobs (with UPSERT mapping)
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS backup_jobs (
        job_id INTEGER PRIMARY KEY,
        client TEXT,
        policy TEXT,
        type TEXT,
        exit_code INTEGER,
        start_time TIMESTAMP,
        finish_time TIMESTAMP,
        duration_secs INTEGER,
        mbytes REAL,
        files_count INTEGER,
        primary_server TEXT,
        media_server TEXT,
        is_rerun_success INTEGER DEFAULT 0
    );
    """)
    
    # Store technician actions
    cursor.execute("""
    CREATE TABLE IF NOT EXISTS job_actions (
        job_id INTEGER PRIMARY KEY,
        action_taken TEXT,
        status TEXT DEFAULT 'Pendente',
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        FOREIGN KEY(job_id) REFERENCES backup_jobs(job_id)
    );
    """)
    
    conn.commit()
    conn.close()

def is_file_processed(file_hash):
    """
    Checks if a CSV file (based on its unique MD5 hash) has already been processed.
    """
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("SELECT 1 FROM processed_files WHERE file_hash = ?", (file_hash,))
    row = cursor.fetchone()
    conn.close()
    return row is not None

def mark_file_processed(file_hash, file_name):
    """
    Logs a file as processed.
    """
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute(
        "INSERT OR IGNORE INTO processed_files (file_hash, file_name) VALUES (?, ?)",
        (file_hash, file_name)
    )
    conn.commit()
    conn.close()

def save_jobs(jobs_df):
    """
    Saves or updates jobs into the backup_jobs table using SQL UPSERT.
    Preserves existing job actions by updating job details but not touching actions.
    """
    conn = get_connection()
    cursor = conn.cursor()
    
    for _, row in jobs_df.iterrows():
        # Handle nan values for start/end times
        start_time = str(row['Start Time']) if not pd.isna(row['Start Time']) else None
        finish_time = str(row['Finish Time']) if not pd.isna(row['Finish Time']) else None
        
        cursor.execute("""
            INSERT INTO backup_jobs (
                job_id, client, policy, type, exit_code, start_time, finish_time, 
                duration_secs, mbytes, files_count, primary_server, media_server, is_rerun_success
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(job_id) DO UPDATE SET
                client = excluded.client,
                policy = excluded.policy,
                type = excluded.type,
                exit_code = excluded.exit_code,
                start_time = excluded.start_time,
                finish_time = excluded.finish_time,
                duration_secs = excluded.duration_secs,
                mbytes = excluded.mbytes,
                files_count = excluded.files_count,
                primary_server = excluded.primary_server,
                media_server = excluded.media_server,
                is_rerun_success = excluded.is_rerun_success;
        """, (
            int(row['Job ID']),
            row['Client'],
            row['Policy'],
            row['Type'],
            int(row['Exit Code']),
            start_time,
            finish_time,
            int(row['Duration_Sec']),
            float(row['MBytes']),
            int(row['# of Files']),
            row['Primary Server'],
            row['Media Server'],
            int(row['is_rerun_success'])
        ))
        
    conn.commit()
    conn.close()

def save_action(job_id, action_taken, status):
    """
    Saves or updates a mitigation action for a failed backup job.
    """
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO job_actions (job_id, action_taken, status, updated_at)
        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
        ON CONFLICT(job_id) DO UPDATE SET
            action_taken = excluded.action_taken,
            status = excluded.status,
            updated_at = CURRENT_TIMESTAMP;
    """, (job_id, action_taken, status))
    conn.commit()
    conn.close()

def get_historical_jobs():
    """
    Retrieves all records from the backup_jobs table joined with job_actions.
    """
    conn = get_connection()
    cursor = conn.cursor()
    cursor.execute("""
        SELECT j.*, a.action_taken, COALESCE(a.status, 'Pendente') as status, a.updated_at
        FROM backup_jobs j
        LEFT JOIN job_actions a ON j.job_id = a.job_id
    """)
    rows = cursor.fetchall()
    conn.close()
    return [dict(row) for row in rows]