database/sqlite_manager.py

import sqlite3
import pandas as pd
from typing import List, Tuple

class SqliteManager:
    """Manages the lifecycle and execution of an internal SQLite database."""
    def __init__(self, db_path: str = ":memory:"):
        self.db_path = db_path
        self.conn = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        """Establishes the connection and sets up initial dummy schema + data."""
        self.conn = sqlite3.Connection(self.db_path)
        self._initialize_schema()

    def close(self):
        if self.conn:
            self.conn.close()

    def _initialize_schema(self):
        cursor = self.conn.cursor()
        
        cursor.execute('''CREATE TABLE DEPARTMENTS (ID INTEGER PRIMARY KEY, NAME TEXT)''')
        cursor.execute('''CREATE TABLE EMPLOYEES (ID INTEGER PRIMARY KEY, NAME TEXT, DEPARTMENT_ID INTEGER, SALARY REAL)''')
        cursor.execute('''CREATE TABLE ROLES (ID INTEGER PRIMARY KEY, TITLE TEXT, MIN_SALARY REAL)''')
        
        cursor.executemany("INSERT INTO DEPARTMENTS (ID, NAME) VALUES (?, ?)", [
            (1, "HR"), (2, "Engineering"), (3, "Marketing"), (4, "Sales")
        ])
        
        cursor.executemany("INSERT INTO EMPLOYEES (ID, NAME, DEPARTMENT_ID, SALARY) VALUES (?, ?, ?, ?)", [
            (1, "Alice", 2, 120000), (2, "Bob", 2, 110000), (3, "Charlie", 1, 80000),
            (4, "David", 3, 90000), (5, "Eve", 4, 95000), (6, "Frank", 2, 105000),
            (7, "Grace", 4, 100000), (8, "Hank", 2, 85000)
        ])
        
        cursor.executemany("INSERT INTO ROLES (ID, TITLE, MIN_SALARY) VALUES (?, ?, ?)", [
            (1, "Associate", 60000), (2, "Senior", 100000), (3, "Staff", 130000)
        ])
        
        self.conn.commit()

    def get_schema_summary(self) -> str:
        return """
Database Schema:
1. DEPARTMENTS (ID INTEGER PRIMARY KEY, NAME TEXT)
2. EMPLOYEES (ID INTEGER PRIMARY KEY, NAME TEXT, DEPARTMENT_ID INTEGER REFERENCES DEPARTMENTS(ID), SALARY REAL)
3. ROLES (ID INTEGER PRIMARY KEY, TITLE TEXT, MIN_SALARY REAL)
"""

    def is_safe_query(self, query: str) -> bool:
        """Validates that a query isn't destructive."""
        forbidden = ["DROP", "DELETE", "UPDATE", "INSERT", "ALTER"]
        upper_query = query.upper()
        return not any(f in upper_query for f in forbidden)

    def execute_dataframe(self, query: str) -> pd.DataFrame:
        """Executes a SQL query and returns its result as a Pandas DataFrame."""
        if not self.conn:
            raise RuntimeError("Database connection not initialized.")
        if not self.is_safe_query(query):
            raise ValueError("Destructive operations are not allowed.")
        return pd.read_sql_query(query, self.conn)