Spaces:
Running
Running
File size: 1,686 Bytes
a31f556 | 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 | # 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)
|