From 41d6283b83bea02f9dfe2fc837671e67f75e1cb4 Mon Sep 17 00:00:00 2001 From: ysolanky Date: Wed, 22 May 2024 16:45:05 -0400 Subject: [PATCH] update-sql-tool-phi-824 --- phi/tools/sql.py | 36 ++++++++++++++---------------------- 1 file changed, 14 insertions(+), 22 deletions(-) diff --git a/phi/tools/sql.py b/phi/tools/sql.py index b7f5a3311..ccdc3271c 100644 --- a/phi/tools/sql.py +++ b/phi/tools/sql.py @@ -1,15 +1,11 @@ +import json from typing import List, Optional, Dict, Any from phi.tools import Toolkit from phi.utils.log import logger try: - import simplejson as json -except ImportError: - raise ImportError("`simplejson` not installed") - -try: - from sqlalchemy import create_engine, Engine, Row + from sqlalchemy import create_engine, Engine from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.inspection import inspect from sqlalchemy.sql.expression import text @@ -128,20 +124,16 @@ def run_sql(self, sql: str, limit: Optional[int] = None) -> List[dict]: """ logger.debug(f"Running sql |\n{sql}") - result = None with self.Session() as sess, sess.begin(): - if limit: - result = sess.execute(text(sql)).fetchmany(limit) - else: - result = sess.execute(text(sql)).fetchall() - - logger.debug(f"SQL result: {result}") - if result is None: - return [] - elif isinstance(result, list): - return [row._asdict() for row in result] - elif isinstance(result, Row): - return [result._asdict()] - else: - logger.debug(f"SQL result type: {type(result)}") - return [] + result = sess.execute(text(sql)) + + # Check if the operation has returned rows. + try: + if limit: + rows = result.fetchmany(limit) + else: + rows = result.fetchall() + return [row._asdict() for row in rows] + except Exception as e: + logger.error(f"Error while executing SQL: {e}") + return []