Skip to content

Commit

Permalink
Merge pull request #979 from phidatahq/update-sql-tool-phi-824
Browse files Browse the repository at this point in the history
update-sql-tool-phi-824
  • Loading branch information
ashpreetbedi authored May 29, 2024
2 parents e3dafa3 + 41d6283 commit 529bf5a
Showing 1 changed file with 14 additions and 22 deletions.
36 changes: 14 additions & 22 deletions phi/tools/sql.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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 []

0 comments on commit 529bf5a

Please sign in to comment.