Spaces:
Sleeping
Sleeping
File size: 2,823 Bytes
9b9c562 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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)
|