Skip to content

Commit 9d6fddc

Browse files
tiberiu92Tiberiu Scarlatysolanky
authored
Fix schema handling in SQLTools for table inspection (#3060)
## Summary Fixes SQLAlchemy schema handling in the SQLTools class. Currently, the schema parameter is accepted in the constructor but not properly used when listing tables or describing table structure. This causes errors when working with databases where tables exist in non-default schemas (outside of "public" in PostgreSQL). ## Type of change - [x] Bug fix - [ ] New feature - [ ] Breaking change - [ ] Improvement - [ ] Model update - [ ] Other: --- ## Checklist - [x] Code complies with style guidelines - [x] Ran format/validation scripts (`./scripts/format.sh` and `./scripts/validate.sh`) - [x] Self-review completed - [x] Documentation updated (comments, docstrings) - [ ] Examples and guides: Relevant cookbook examples have been included or updated (if applicable) - [x] Tested in clean environment - [ ] Tests added/updated (if applicable) --- ## Additional Notes The core issue is that SQLAlchemy defaults to the "public" schema in PostgreSQL when no schema is specified in the inspector methods. Even though a schema parameter is provided to the SQLTools constructor, it wasn't being passed to the underlying inspector calls. This fix ensures that: 1. When listing tables with `list_tables()`, we properly pass the schema parameter to `get_table_names()` 2. When describing a table with `describe_table()`, we correctly use the schema parameter with `get_columns()` 3. Error logging has been improved throughout the class This resolves the "table not found" errors that occur when working with databases where tables are located in non-default schemas, such as when using a schema other than "public" in PostgreSQL. Tested with PostgreSQL database with tables in custom schemas to verify the fix works properly. --------- Co-authored-by: Tiberiu Scarlat <tiberiu.scarlat@arnia.ro> Co-authored-by: Yash Pratap Solanky <101447028+ysolanky@users.noreply.github.com>
1 parent 75d0d50 commit 9d6fddc

File tree

1 file changed

+15
-4
lines changed

1 file changed

+15
-4
lines changed

libs/agno/agno/tools/sql.py

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,8 @@ def __init__(
4949
self.db_engine: Engine = _engine
5050
self.Session: sessionmaker[Session] = sessionmaker(bind=self.db_engine)
5151

52+
self.schema = schema
53+
5254
# Tables this toolkit can access
5355
self.tables: Optional[Dict[str, Any]] = tables
5456

@@ -71,7 +73,11 @@ def list_tables(self) -> str:
7173

7274
try:
7375
log_debug("listing tables in the database")
74-
table_names = inspect(self.db_engine).get_table_names()
76+
inspector = inspect(self.db_engine)
77+
if self.schema:
78+
table_names = inspector.get_table_names(schema=self.schema)
79+
else:
80+
table_names = inspector.get_table_names()
7581
log_debug(f"table_names: {table_names}")
7682
return json.dumps(table_names)
7783
except Exception as e:
@@ -90,9 +96,14 @@ def describe_table(self, table_name: str) -> str:
9096

9197
try:
9298
log_debug(f"Describing table: {table_name}")
93-
table_names = inspect(self.db_engine)
94-
table_schema = table_names.get_columns(table_name)
95-
return json.dumps([str(column) for column in table_schema])
99+
inspector = inspect(self.db_engine)
100+
table_schema = inspector.get_columns(table_name, schema=self.schema)
101+
return json.dumps(
102+
[
103+
{"name": column["name"], "type": str(column["type"]), "nullable": column["nullable"]}
104+
for column in table_schema
105+
]
106+
)
96107
except Exception as e:
97108
logger.error(f"Error getting table schema: {e}")
98109
return f"Error getting table schema: {e}"

0 commit comments

Comments
 (0)