58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
import pandas as pd
|
|
from sqlalchemy import inspect
|
|
from typing import Any, List, Union, Optional
|
|
|
|
|
|
def format_query_results(query_results: Any) -> str:
|
|
"""
|
|
Format query results as a readable string using pandas
|
|
|
|
Args:
|
|
query_results: SQLAlchemy query, query results, or model instance(s)
|
|
|
|
Returns:
|
|
Formatted string representation of the query results
|
|
"""
|
|
try:
|
|
# If it's a query object, execute it
|
|
if hasattr(query_results, 'all'):
|
|
results = query_results.all()
|
|
elif not isinstance(query_results, list):
|
|
results = [query_results]
|
|
else:
|
|
results = query_results
|
|
|
|
# Handle different types of results
|
|
if results and hasattr(results[0], '__table__'):
|
|
# SQLAlchemy ORM objects
|
|
data = []
|
|
for item in results:
|
|
row = {}
|
|
for column in inspect(item).mapper.column_attrs:
|
|
row[column.key] = getattr(item, column.key)
|
|
data.append(row)
|
|
df = pd.DataFrame(data)
|
|
elif results and isinstance(results[0], tuple):
|
|
# Join query results (tuples)
|
|
if hasattr(results[0], '_fields'): # Named tuples
|
|
df = pd.DataFrame(results)
|
|
else:
|
|
# Regular tuples - try to get column names from query
|
|
if hasattr(query_results, 'statement'):
|
|
columns = query_results.statement.columns.keys()
|
|
df = pd.DataFrame(results, columns=columns)
|
|
else:
|
|
df = pd.DataFrame(results)
|
|
else:
|
|
# Fallback for other types
|
|
df = pd.DataFrame(results)
|
|
|
|
# Format the output with pandas
|
|
with pd.option_context('display.max_rows', 20, 'display.max_columns', None,
|
|
'display.width', 1000):
|
|
formatted_output = f"Query returned {len(df)} results:\n{df}"
|
|
|
|
return formatted_output
|
|
except Exception as e:
|
|
return f"Error formatting query results: {str(e)}"
|