backend/tools/database_tools.py

# backend/tasks/tools/database_tools.py
import sqlite3
import asyncio

try:
    import psycopg2
    from psycopg2.extras import RealDictCursor
except ImportError:
    psycopg2 = None

async def query_sqlite(db_path: str, sql: str) -> list[dict]:
    def _query():
        conn = sqlite3.connect(db_path)
        conn.execute('PRAGMA journal_mode=WAL')
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()
        cursor.execute(sql)
        
        # Check if it's a mutation
        if sql.strip().upper().startswith(("INSERT", "UPDATE", "DELETE")):
            conn.commit()
            rows = [{"affected_rows": cursor.rowcount}]
        else:
            rows = [dict(row) for row in cursor.fetchall()]
            
        conn.close()
        return rows
        
    return await asyncio.to_thread(_query)

async def query_postgres(connection_string: str, sql: str) -> list[dict]:
    if psycopg2 is None:
        return [{"error": "psycopg2 is not installed"}]
        
    def _query():
        try:
            conn = psycopg2.connect(connection_string)
            cursor = conn.cursor(cursor_factory=RealDictCursor)
            cursor.execute(sql)
            
            if sql.strip().upper().startswith(("INSERT", "UPDATE", "DELETE")):
                conn.commit()
                rows = [{"affected_rows": cursor.rowcount}]
            else:
                rows = [dict(row) for row in cursor.fetchall()]
                
            conn.close()
            return rows
        except Exception as e:
            return [{"error": str(e)}]
            
    return await asyncio.to_thread(_query)